module Obs.NormalForm.Normalise

import Data.Bool
import Data.Nat

import Decidable.Equality

import Obs.Logging
import Obs.NormalForm
import Obs.Substitution
import Obs.Universe

import Text.PrettyPrint.Prettyprinter

%default total

-- Aliases ---------------------------------------------------------------------

public export 0
LogConstructor : Type -> Unsorted.Family Relevance
LogConstructor ann ctx = Logging ann (Constructor ctx)

public export 0
LogNormalForm : Type -> Sorted.Family Relevance
LogNormalForm ann b ctx = Logging ann (NormalForm b ctx)

0
LogDeclForm : Type -> Unsorted.Family Relevance
LogDeclForm ann ctx = Logging ann (DeclForm ctx)

0
LogNormalForm' : Type -> Sorted.Family Relevance
LogNormalForm' ann b ctx = Either (Logging ann (NormalForm b ctx)) (Elem b ctx)

-- Copied and specialised from Obs.Substitution
lift : (ctx : List (r ** Maybe String))
  -> Map ctx' (LogNormalForm' ann) ctx''
  -> Map (map DPair.fst ctx ++ ctx') (LogNormalForm' ann) (map DPair.fst ctx ++ ctx'')
lift [] f = f
lift ((s ** y) :: ctx) f = add (LogNormalForm' ann)
  (Left $ pure $ point y Here)
  (\i => bimap (\t => pure (rename !t There)) There (lift ctx f i))

-- Normalisation ---------------------------------------------------------------

subst' : NormalForm ~|> Hom (LogNormalForm' ann) (LogNormalForm ann)

export
subst1 : {s' : _} -> NormalForm s ctx -> NormalForm s' (s :: ctx) -> LogNormalForm ann s' ctx
subst1 t u = subst' u (add (LogNormalForm' ann) (Left $ pure t) Right)

export
map1 : {s' : _} ->
  NormalForm s (s :: ctx) ->
  NormalForm s' (s :: ctx) ->
  LogNormalForm ann s' (s :: ctx)
map1 t u = subst' u (add (LogNormalForm' ann) (Left $ pure t) (Right . There))

export
doApp : {domainRel : _} ->
  (fun : NormalForm (function domainRel codomainRel) ctx) ->
  (arg : NormalForm domainRel ctx) ->
  LogNormalForm ann codomainRel ctx

doApp (Ntrl fun) arg = pure $ Ntrl $ App {argRel = _, fun, arg}
doApp (Cnstr (Lambda {domainRel = domainRel', var, body})) arg = do
  let Yes Refl = decEq domainRel domainRel'
    | No _ => fatal "internal relevance mismatch"

  subst1 arg body
doApp (Cnstr t) arg = inScope "wrong constructor for apply" $ fatal t
doApp Irrel arg = pure Irrel

export
doFst : (firstRel, secondRel : _) ->
  (arg : NormalForm (pair firstRel secondRel) ctx) ->
  LogNormalForm ann firstRel ctx

doFst Irrelevant secondRel arg = pure Irrel
doFst Relevant secondRel (Ntrl arg) = pure $ Ntrl $ First {secondRel, arg}
doFst Relevant secondRel (Cnstr (Pair {indexRel = Relevant, elementRel, prf, first, second})) =
  pure first
doFst Relevant secondRel (Cnstr t) = inScope "wrong constructor for fst" $ fatal t

export
doSnd : (firstRel, secondRel : _) ->
  (arg : NormalForm (pair firstRel secondRel) ctx) ->
  LogNormalForm ann secondRel ctx

doSnd firstRel Irrelevant arg = pure Irrel
doSnd firstRel Relevant arg =
  let arg' : NormalForm Relevant ctx
      arg' = rewrite sym $ pairRelevantRight firstRel in arg
  in case arg' of
  Ntrl arg => pure $ Ntrl $ Second {firstRel, arg}
  Cnstr (Pair {indexRel, elementRel = Relevant, prf, first, second}) => pure second
  Cnstr t => inScope "wrong constructor for snd" $ fatal t

export
doIf : {rel : _} ->
  (discriminant : NormalForm Relevant ctx) ->
  (true, false : NormalForm rel ctx) ->
  LogNormalForm ann rel ctx
doIf {rel = Irrelevant} discriminant true false = pure Irrel
doIf {rel = Relevant} (Ntrl discriminant) true false = pure $ Ntrl $ If {discriminant, true, false}
doIf {rel = Relevant} (Cnstr True) true false = pure true
doIf {rel = Relevant} (Cnstr False) true false = pure false
doIf {rel = Relevant} (Cnstr t) true false = inScope "wrong constructor for if" $ fatal t

export
doAbsurd : (rel : _) -> NormalForm rel ctx
doAbsurd Irrelevant = Irrel
doAbsurd Relevant = Ntrl Absurd

export
doCast : (rel : _) ->
  (oldType, newType : TypeNormalForm ctx) ->
  (value : NormalForm rel ctx) ->
  LogNormalForm ann rel ctx

doCastHelper : (oldType, newType : Constructor ctx) ->
  (value : NormalForm Relevant ctx) ->
  LogNormalForm ann Relevant ctx
doCastHelper (Universe {s}) (Universe {s = s'}) value = pure value
doCastHelper
  oldType@(Pi {domainSort, codomainSort, domain, codomain})
  newType@(Pi
    { domainSort = domainSort'
    , codomainSort = codomainSort'
    , domain = domain'
    , codomain = codomain'
    })
  value =
  let y : NormalForm (relevance domainSort) (relevance domainSort :: ctx)
      y = point domain.var Here
  in do

  let Yes Refl = decEq (domainSort, codomainSort) (domainSort', codomainSort')
    | No _ => pure $ Ntrl $ CastStuck {oldType, newType, value}

  let domainType = weaken [relevance domainSort] domain.type
  let domainType' = weaken [relevance domainSort] domain'.type

  x <- doCast
    { rel = relevance domainSort
    , oldType = assert_smaller oldType domainType'
    , newType = assert_smaller newType domainType
    , value = y
    }

  codomainType <- map1 x codomain
  codomainType' <- map1 y codomain'

  call <- doApp (weaken [relevance domainSort] value) x
  body <- doCast
    { rel = Relevant
    , oldType = assert_smaller oldType codomainType
    , newType = assert_smaller newType codomainType'
    , value = call
    }

  pure $ Cnstr $ Lambda {domainRel = relevance domainSort, var = Nothing, body}
doCastHelper
  oldType@(Sigma {indexSort = indexSort@(Set k), elementSort, index, element})
  newType@(Sigma
    { indexSort = indexSort'
    , elementSort = elementSort'
    , index = index'
    , element = element'
    })
  value = do
  let Yes Refl = decEq (indexSort, elementSort) (indexSort', elementSort')
    | No _ => pure $ Ntrl $ CastStuck {oldType, newType, value}

  first <- doFst Relevant (relevance elementSort) value
  second <- doSnd Relevant (relevance elementSort) value

  first' <- doCast
    { rel = Relevant
    , oldType = assert_smaller oldType index.type
    , newType = assert_smaller newType index'.type
    , value = first
    }

  elementType <- subst1 first element
  elementType' <- subst1 first' element'

  second' <- doCast
    { rel = relevance elementSort
    , oldType = assert_smaller oldType elementType
    , newType = assert_smaller newType elementType'
    , value = second
    }

  pure $ Cnstr $ Pair
    { indexRel = Relevant
    , elementRel = relevance elementSort
    , prf = Relevant
    , first = first'
    , second = second'
    }
doCastHelper
  oldType@(Sigma {indexSort = Prop, elementSort = elementSort@(Set k), index, element})
  newType@(Sigma
    { indexSort = Prop
    , elementSort = elementSort'
    , index = index'
    , element = element'
    })
  value = do
  let Yes Refl = decEq elementSort elementSort'
    | No _ => pure $ Ntrl $ CastStuck {oldType, newType, value}

  first <- doFst Irrelevant Relevant value
  second <- doSnd Irrelevant Relevant value

  first' <- doCast
    { rel = Irrelevant
    , oldType = assert_smaller oldType index.type
    , newType = assert_smaller newType index'.type
    , value = first
    }

  elementType <- subst1 first element
  elementType' <- subst1 first' element'

  second' <- doCast
    { rel = Relevant
    , oldType = assert_smaller oldType elementType
    , newType = assert_smaller newType elementType'
    , value = second
    }

  pure $ Cnstr $ Pair
    { indexRel = Irrelevant
    , elementRel = Relevant
    , prf = Relevant
    , first = first'
    , second = second'
    }
doCastHelper Bool Bool value = pure value
doCastHelper oldType newType value = pure $ Ntrl $ CastStuck {oldType, newType, value}

doCast Irrelevant oldType newType value = pure Irrel
doCast Relevant (Ntrl oldType) newType value = pure $ Ntrl $ CastL {oldType, newType, value}
doCast Relevant (Cnstr oldType) (Ntrl newType) value =
  pure $ Ntrl $ CastR {oldType, newType, value}
doCast Relevant (Cnstr oldType) (Cnstr newType) value = doCastHelper oldType newType value

export
doEqual : (rel : _) ->
  (type : TypeNormalForm ctx) ->
  (left, right : NormalForm rel ctx) ->
  LogNormalForm ann Relevant ctx

equalHelper : (left, right : Constructor ctx) -> LogNormalForm ann Relevant ctx
equalHelper (Universe {s}) (Universe {s = s'}) = pure (Cnstr Top)
equalHelper
  left@(Pi {domainSort, codomainSort, domain, codomain})
  right@(Pi
    { domainSort = domainSort'
    , codomainSort = codomainSort'
    , domain = domain'
    , codomain = codomain'
    }) =
  let y : NormalForm (relevance domainSort) (relevance domainSort :: Irrelevant :: ctx)
      y = point domain.var Here
  in do

  let Yes Refl = decEq (domainSort, codomainSort) (domainSort', codomainSort')
    | No _ => pure $ Ntrl $ EqualStuck {left, right}

  domainEqual <- doEqual
    { rel = Relevant
    , type = cast domainSort
    , left = assert_smaller right domain'.type
    , right = assert_smaller left domain.type
    }

  let domainType = Sorted.weaken [relevance domainSort, Irrelevant] domain.type
  let domainType' = Sorted.weaken [relevance domainSort, Irrelevant] domain'.type

  x <- doCast
    { rel = relevance domainSort
    , oldType = domainType'
    , newType = domainType
    , value = y
    }

  codomainType <- map1 x (rename codomain (add Elem Here (There . There)))
  codomainType' <- map1 y (rename codomain' (add Elem Here (There . There)))

  codomainEqual <- doEqual
    { rel = Relevant
    , type = cast codomainSort
    , left = assert_smaller left codomainType
    , right = assert_smaller right codomainType'
    }

  let returnElement = Cnstr $ Pi
       { domainSort = domainSort
       , codomainSort = Prop
       , domain = MkDecl Nothing (Sorted.weaken [Irrelevant] domain.type)
       , codomain = codomainEqual
       }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing domainEqual
    , element = returnElement
    }
equalHelper
  left@(Sigma {indexSort, elementSort, index, element})
  right@(Sigma
    { indexSort = indexSort'
    , elementSort = elementSort'
    , index = index'
    , element = element'
    }) =
  let x : NormalForm (relevance indexSort) (relevance indexSort :: Irrelevant :: ctx)
      x = point index.var Here
  in do

  let Yes Refl = decEq (indexSort, elementSort) (indexSort', elementSort')
    | No _ => pure $ Ntrl $ EqualStuck {left, right}

  indexEqual <- doEqual
    { rel = Relevant
    , type = cast indexSort
    , left = assert_smaller left index.type
    , right = assert_smaller right index'.type
    }

  let indexType = Sorted.weaken [relevance indexSort, Irrelevant] index.type
  let indexType' = Sorted.weaken [relevance indexSort, Irrelevant] index'.type

  y <- doCast
    { rel = relevance indexSort
    , oldType = indexType
    , newType = indexType'
    , value = x
    }

  elementType <- map1 x (rename element (add Elem Here (There . There)))
  elementType' <- map1 y (rename element' (add Elem Here (There . There)))

  elementEqual <- doEqual
    { rel = Relevant
    , type = cast elementSort
    , left = assert_smaller left elementType
    , right = assert_smaller right elementType'
    }

  let returnElement = Cnstr $ Pi
       { domainSort = indexSort
       , codomainSort = Prop
       , domain = MkDecl Nothing (Sorted.weaken [Irrelevant] index.type)
       , codomain = elementEqual
       }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing indexEqual
    , element = returnElement
    }
equalHelper Bool Bool = pure (Cnstr Top)
equalHelper left right = pure $ Ntrl $ EqualStuck {left, right}

doEqualType : (left, right : TypeNormalForm ctx) -> LogNormalForm ann Relevant ctx
doEqualType (Ntrl left) right = pure $ Ntrl $ EqualL {left, right}
doEqualType (Cnstr left) (Ntrl right) = pure $ Ntrl $ EqualR {left, right}
doEqualType (Cnstr left) (Cnstr right) = equalHelper left right

doEqual Irrelevant type left right = pure $ Cnstr Top
doEqual Relevant (Ntrl type) left right = pure $ Ntrl $ Equal {type, left, right}
doEqual Relevant (Cnstr (Universe {s = Prop})) left right = do
  let leftToRight = Cnstr $ Pi
        { domainSort = Prop
        , codomainSort = Prop
        , domain = MkDecl Nothing left
        , codomain = weaken [Irrelevant] right
        }
  let rightToLeft = Cnstr $ Pi
        { domainSort = Prop
        , codomainSort = Prop
        , domain = MkDecl Nothing right
        , codomain = weaken [Irrelevant] left
        }
  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing leftToRight
    , element = Sorted.weaken [Irrelevant] rightToLeft
    }
doEqual Relevant (Cnstr (Universe {s = Set _})) left right = doEqualType left right
doEqual Relevant (Cnstr (Pi {domainSort, codomainSort, domain, codomain})) left right =
  let var : NormalForm (relevance domainSort) (relevance domainSort :: ctx)
      var = point domain.var Here
  in do

  leftApp <- doApp (weaken [relevance domainSort] left) var
  rightApp <- doApp (weaken [relevance domainSort] right) var

  equality <- doEqual
    { rel = Relevant
    , type = codomain
    , left = assert_smaller left leftApp
    , right = assert_smaller right rightApp
    }

  pure $ Cnstr $ Pi {domainSort, codomainSort = Prop, domain, codomain = equality}
doEqual
  Relevant
  (Cnstr (Sigma {indexSort = indexSort@(Set _), elementSort, index, element}))
  left
  right = do
  leftFirst <- doFst Relevant (relevance elementSort) left
  rightFirst <- doFst Relevant (relevance elementSort) right
  leftSecond <- doSnd Relevant (relevance elementSort) left
  rightSecond <- doSnd Relevant (relevance elementSort) right

  leftEquality <- doEqual
    { rel = Relevant
    , type = index.type
    , left = assert_smaller left leftFirst
    , right = assert_smaller right rightFirst
    }

  leftElementType <- subst1 leftFirst element
  rightElementType <- subst1 rightFirst element

  leftSecond <- doCast (relevance elementSort) leftElementType rightElementType leftSecond

  rightEquality <- doEqual
    { rel = relevance elementSort
    , type = rightElementType
    , left = assert_smaller left leftSecond
    , right = assert_smaller right rightSecond
    }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing leftEquality
    , element = Sorted.weaken [Irrelevant] rightEquality
    }
doEqual
  Relevant
  (Cnstr (Sigma {indexSort = Prop, elementSort = elementSort@(Set _), index, element}))
  left
  right = do
  leftFirst <- doFst Irrelevant Relevant left
  rightFirst <- doFst Irrelevant Relevant right
  leftSecond <- doSnd Irrelevant Relevant left
  rightSecond <- doSnd Irrelevant Relevant right

  leftEquality <- doEqual
    { rel = Irrelevant
    , type = index.type
    , left = assert_smaller left leftFirst
    , right = assert_smaller right rightFirst
    }

  leftElementType <- subst1 leftFirst element
  rightElementType <- subst1 rightFirst element

  leftSecond <- doCast Relevant leftElementType rightElementType leftSecond

  rightEquality <- doEqual
    { rel = Relevant
    , type = rightElementType
    , left = assert_smaller left leftSecond
    , right = assert_smaller right rightSecond
    }

  pure $ Cnstr $ Sigma
    { indexSort = Prop
    , elementSort = Prop
    , index = MkDecl Nothing leftEquality
    , element = Sorted.weaken [Irrelevant] rightEquality
    }
doEqual Relevant (Cnstr Bool) left right = do
  whenLeftTrue <- doIf right (Cnstr Top) (Cnstr Bottom)
  whenLeftFalse <- doIf right (Cnstr Bottom) (Cnstr Top)
  doIf left whenLeftTrue whenLeftFalse
doEqual Relevant (Cnstr t) left right = inScope "wrong constructor for equal" $ fatal t

substDecl : DeclForm ~|> Hom (LogNormalForm' ann) (LogDeclForm ann)
substDecl (MkDecl var type) f = pure (MkDecl var !(subst' type f))

substCnstr : Constructor ~|> Hom (LogNormalForm' ann) (LogConstructor ann)
substCnstr (Universe {s}) f = pure (Universe {s})
substCnstr (Pi {domainSort, codomainSort, domain, codomain}) f = do
  domain <- substDecl domain f
  codomain <- subst' codomain (lift [(_ ** domain.var)] f)
  pure (Pi {domainSort, codomainSort, domain, codomain})
substCnstr (Lambda {domainRel, var, body}) f = do
  body <- subst' body (lift [(_ ** var)] f)
  pure (Lambda {domainRel, var, body})
substCnstr (Sigma {indexSort, elementSort, index, element}) f = do
  index <- substDecl index f
  element <- subst' element (lift [(_ ** index.var)] f)
  pure (Sigma {indexSort, elementSort, index, element})
substCnstr (Pair {indexRel, elementRel, prf, first, second}) f = do
  first <- subst' first f
  second <- subst' second f
  pure (Pair {indexRel, elementRel, prf, first, second})
substCnstr Bool f = pure Bool
substCnstr True f = pure True
substCnstr False f = pure False
substCnstr Top f = pure Top
substCnstr Bottom f = pure Bottom

substNtrl : Neutral ~|> Hom (LogNormalForm' ann) (LogNormalForm ann Relevant)
substNtrl (Var {var, i}) f = case f i of
  Left t => t
  Right i => pure $ Ntrl $ Var {var, i}
substNtrl (App {argRel, fun, arg}) f = do
  fun <- substNtrl fun f
  arg <- subst' arg f
  assert_total (doApp fun arg)
substNtrl (First {secondRel, arg}) f = do
  arg <- substNtrl arg f
  doFst Relevant secondRel arg
substNtrl (Second {firstRel, arg}) f = do
  arg <- substNtrl arg f
  let arg = rewrite pairRelevantRight firstRel in arg
  doSnd firstRel Relevant arg
substNtrl (If {discriminant, true, false}) f = do
  discriminant <- substNtrl discriminant f
  true <- subst' true f
  false <- subst' false f
  doIf discriminant true false
substNtrl Absurd f = pure (doAbsurd Relevant)
substNtrl (Equal {type, left, right}) f = do
  type <- substNtrl type f
  left <- subst' left f
  right <- subst' right f
  assert_total (doEqual Relevant type left right)
substNtrl (EqualL {left, right}) f = do
  left <- substNtrl left f
  right <- subst' right f
  assert_total (doEqualType left right)
substNtrl (EqualR {left, right}) f = do
  left <- substCnstr left f
  right <- substNtrl right f
  assert_total (doEqualType (Cnstr left) right)
substNtrl (EqualStuck {left, right}) f = do
  left <- substCnstr left f
  right <- substCnstr right f
  assert_total (doEqualType (Cnstr left) (Cnstr right))
substNtrl (CastL {oldType, newType, value}) f = do
  oldType <- substNtrl oldType f
  newType <- subst' newType f
  value <- subst' value f
  assert_total (doCast Relevant oldType newType value)
substNtrl (CastR {oldType, newType, value}) f = do
  oldType <- substCnstr oldType f
  newType <- substNtrl newType f
  value <- subst' value f
  assert_total (doCast Relevant (Cnstr oldType) newType value)
substNtrl (CastStuck {oldType, newType, value}) f = do
  oldType <- substCnstr oldType f
  newType <- substCnstr newType f
  value <- subst' value f
  assert_total (doCast Relevant (Cnstr oldType) (Cnstr newType) value)

subst' (Ntrl t) f = substNtrl t f
subst' (Cnstr t) f = pure $ Cnstr !(substCnstr t f)
subst' Irrel f = pure Irrel

export
subst : NormalForm ~|> Hom (LogNormalForm ann) (LogNormalForm ann)
subst t f = subst' t (Left . f)