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
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 ** (String, (s ** relevance s = r))))
  -> 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
doApp : {domain : _} ->
  NormalForm (function domain codomain) ctx ->
  NormalForm domain ctx ->
  LogNormalForm ann codomain ctx
doApp (Ntrl t) u = pure (Ntrl $ App _ t u)
doApp (Cnstr (Lambda s var t)) u = inScope "doApp" $ do
  trace $ pretty {ann} "substituting" <++> pretty u <+> softline <+> pretty "for \{var} in" <++> pretty t
  let Yes Refl = decEq domain (relevance s)
    | No _ => fatal "internal relevance mismatch"
  subst' t (add (LogNormalForm' ann) (Left $ pure u) Right)
doApp (Cnstr t) u = inScope "wrong constructor for apply" $ fatal t
doApp Irrel u = pure Irrel

export
doFst : (fst, snd : _) -> NormalForm (pair fst snd) ctx -> LogNormalForm ann fst ctx
doFst Irrelevant snd t = pure Irrel
doFst Relevant snd (Ntrl t) = pure (Ntrl $ Fst snd t)
doFst Relevant snd (Cnstr (Pair (Set _) s' prf t u)) = pure t
doFst Relevant snd (Cnstr t) = inScope "wrong constructor for fst" $ fatal t

export
doSnd : (fst, snd : _) -> NormalForm (pair fst snd) ctx -> LogNormalForm ann snd ctx
doSnd fst Irrelevant t = pure Irrel
doSnd fst Relevant t =
  let t' : NormalForm Relevant ctx
      t' = rewrite sym $ pairRelevantRight fst in t
  in case t' of
  Ntrl t => pure (Ntrl $ Snd fst t)
  Cnstr (Pair _ (Set _) prf t u) => pure u
  Cnstr t => inScope "wrong constructor for snd" $ fatal t

export
doIf : {r : _} ->
  NormalForm Relevant ctx ->
  NormalForm r ctx ->
  NormalForm r ctx ->
  LogNormalForm ann r ctx
doIf {r = Irrelevant} t u v = pure Irrel
doIf {r = Relevant} (Ntrl t) u v = pure (Ntrl $ If t u v)
doIf {r = Relevant} (Cnstr True) u v = pure u
doIf {r = Relevant} (Cnstr False) u v = pure v
doIf {r = Relevant} (Cnstr t) u v = inScope "wrong constructor for case" $ fatal t

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

export
doCast : (r : _) -> (a, b : TypeNormalForm ctx) -> NormalForm r ctx -> LogNormalForm ann r ctx

doCastR : (a : Constructor ctx) ->
  (b : TypeNormalForm ctx) ->
  NormalForm Relevant ctx ->
  LogNormalForm ann Relevant ctx
doCastR a (Ntrl b) t = pure (Ntrl $ CastR a b t)
doCastR (Universe _) (Cnstr (Universe _)) t = pure t
doCastR ty@(Pi s s'@(Set _) var a b) (Cnstr ty'@(Pi l l' _ a' b')) t =
  let y : NormalForm (relevance s) (relevance s :: ctx)
      y = point (var, (s ** Refl)) Here
  in do
  let Yes Refl = decEq (s, s') (l, l')
    | No _ => pure (Ntrl $ CastStuck ty ty' t)
  x <- assert_total (doCast (relevance s) (weaken [relevance s] a') (weaken [relevance s] a) y)
  b <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure x) (Right . There)))
  b' <- assert_total (subst' b' (add (LogNormalForm' ann) (Left $ pure y) (Right . There)))
  t <- assert_total (doApp (Sorted.weaken [relevance s] t) x)
  t <- assert_total (doCast Relevant b b' t)
  pure (Cnstr $ Lambda s var t)
doCastR ty@(Sigma s@(Set k) s' var a b) (Cnstr ty'@(Sigma l l' _ a' b')) t = do
  let Yes Refl = decEq (s, s') (l, l')
    | No _ => pure (Ntrl $ CastStuck ty ty' t)
  t1 <- doFst Relevant (relevance s') t
  u1 <- assert_total (doCast Relevant a a' t)
  b <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure t1) Right))
  b' <- assert_total (subst' b' (add (LogNormalForm' ann) (Left $ pure u1) Right))
  t2 <- doSnd Relevant (relevance s') t
  u2 <- assert_total (doCast (relevance s') b b' t2)
  pure (Cnstr $ Pair (Set k) s' Set u1 u2)
doCastR ty@(Sigma Prop s'@(Set k) var a b) (Cnstr ty'@(Sigma Prop l' _ a' b')) t = do
  let Yes Refl = decEq s' l'
    | No _ => pure (Ntrl $ CastStuck ty ty' t)
  b <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure Irrel) Right))
  b' <- assert_total (subst' b' (add (LogNormalForm' ann) (Left $ pure Irrel) Right))
  t2 <- doSnd Irrelevant Relevant t
  u2 <- assert_total (doCast Relevant b b' t2)
  pure (Cnstr $ Pair Prop (Set k) Set Irrel u2)
doCastR Bool (Cnstr Bool) t = pure t
doCastR a (Cnstr b) t = pure (Ntrl $ CastStuck a b t)

doCast Irrelevant a b t = pure Irrel
doCast Relevant (Ntrl a) b t = pure (Ntrl $ CastL a b t)
doCast Relevant (Cnstr a) b t = doCastR a b t

export
doEqual : (r : _) ->
  (a : TypeNormalForm ctx) ->
  NormalForm r ctx ->
  NormalForm r ctx ->
  LogNormalForm ann Relevant ctx

-- Relies heavily on typing
doEqualR : (a : Constructor ctx) -> (b : TypeNormalForm ctx) -> LogNormalForm ann Relevant ctx
doEqualR a (Ntrl b) = pure (Ntrl $ EqualR a b)
doEqualR (Universe _) (Cnstr (Universe s)) = pure (Cnstr Top)
doEqualR ty@(Pi s s' var a b) (Cnstr ty'@(Pi l l' _ a' b')) =
  let u : NormalForm (relevance s) (relevance s :: ctx)
      u = point (var, (s ** Refl)) Here
  in do
  let Yes Refl = decEq (s, s') (l, l')
    | No _ => pure (Ntrl $ EqualStuck ty ty')
  eq1 <- assert_total (doEqual Relevant (cast s) a' a)
  t <- doCast (relevance s) (weaken [relevance s] a') (weaken [relevance s] a) u
  b <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure t) (Right . There)))
  b' <- assert_total (subst' b' (add (LogNormalForm' ann) (Left $ pure u) (Right . There)))
  eq2 <- assert_total (doEqual Relevant (cast s') b b')
  pure (Cnstr $ Sigma Prop Prop "_" eq1 (Cnstr $ Unsorted.weaken [Irrelevant] $ Pi s Prop var a eq2))
doEqualR ty@(Sigma s s' var a b) (Cnstr ty'@(Sigma l l' _ a' b')) =
  let t : NormalForm (relevance s) (relevance s :: ctx)
      t = point (var, (s ** Refl)) Here
  in do
  let Yes Refl = decEq (s, s') (l, l')
    | No _ => pure (Ntrl $ EqualStuck ty ty')
  eq1 <- assert_total (doEqual Relevant (cast s) a a')
  u <- doCast (relevance s) (weaken [relevance s] a) (weaken [relevance s] a') t
  b <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure t) (Right . There)))
  b' <- assert_total (subst' b' (add (LogNormalForm' ann) (Left $ pure u) (Right . There)))
  eq2 <- assert_total (doEqual Relevant (cast s') b b')
  pure (Cnstr $ Sigma Prop Prop "_" eq1 (Cnstr $ Unsorted.weaken [Irrelevant] $ Pi s Prop var a eq2))
doEqualR Bool (Cnstr Bool) = pure (Cnstr Top)
doEqualR Top (Cnstr Top) = pure (Cnstr Top)
doEqualR Bottom (Cnstr Bottom) = pure (Cnstr Top)
doEqualR a (Cnstr b) = pure (Ntrl $ EqualStuck a b)

export
doEqualSet : (a, b : TypeNormalForm ctx) -> LogNormalForm ann Relevant ctx
doEqualSet (Ntrl a) b = pure (Ntrl $ EqualL a b)
doEqualSet (Cnstr a) b = doEqualR a b

doEqual Irrelevant a t u = pure (Cnstr Top)
doEqual Relevant (Ntrl a) t u = pure (Ntrl $ Equal a t u)
doEqual Relevant (Cnstr (Universe Prop)) t u = do
  pure (Cnstr $ Sigma Prop Prop ""
    (Cnstr $ Pi Prop Prop "" t (Sorted.weaken [Irrelevant] u))
    (Cnstr $ Unsorted.weaken [Irrelevant] $ Pi Prop Prop "" u (Sorted.weaken [Irrelevant] t)))
doEqual Relevant (Cnstr (Universe (Set _))) t u = doEqualSet t u
doEqual Relevant (Cnstr (Pi s (Set k) var a b)) t u =
  let x : NormalForm (relevance s) (relevance s :: ctx)
      x = point (var, (s ** Refl)) Here
  in do
  t <- assert_total (doApp (weaken [relevance s] t) x)
  u <- assert_total (doApp (weaken [relevance s] u) x)
  eq <- doEqual Relevant b t u
  pure (Cnstr $ Pi s Prop var a eq)
doEqual Relevant (Cnstr (Sigma s@(Set _) s' var a b)) t u = do
  t1 <- doFst Relevant (relevance s') t
  u1 <- doFst Relevant (relevance s') u
  t2 <- doSnd Relevant (relevance s') t
  u2 <- doSnd Relevant (relevance s') u
  eq1 <- doEqual Relevant a t1 u1
  bt1 <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure t1) Right))
  bu1 <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure u1) Right))
  t2' <- doCast (relevance s') bt1 bu1 t2
  eq2 <- doEqual (relevance s') (assert_smaller b bu1) t2' u2
  pure (Cnstr $ Sigma Prop Prop "_" eq1 (Sorted.weaken [Irrelevant] eq2))
doEqual Relevant (Cnstr (Sigma Prop (Set k) var a b)) t u = do
  t2 <- doSnd Irrelevant Relevant t
  u2 <- doSnd Irrelevant Relevant u
  bt1 <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure $ Irrel) Right))
  bu1 <- assert_total (subst' b (add (LogNormalForm' ann) (Left $ pure $ Irrel) Right))
  t2' <- doCast Relevant bt1 bu1 t2
  eq2 <- doEqual Relevant (assert_smaller b bu1) t2' u2
  pure (Cnstr $ Sigma Prop Prop "_" (Cnstr Top) (Sorted.weaken [Irrelevant] eq2))
doEqual Relevant (Cnstr Bool) t u = do
  relevant <- doIf u (Cnstr Top) (Cnstr Bottom)
  irrelevant <- doIf u (Cnstr Bottom) (Cnstr Top)
  doIf t relevant irrelevant
doEqual Relevant (Cnstr a) t u = inScope "wrong constructor for equal" $ fatal a

substCnstr : Constructor ~|> Hom (LogNormalForm' ann) (LogConstructor ann)
substCnstr (Universe s) f = pure (Universe s)
substCnstr (Pi s s' var a b) f = do
  a <- subst' a f
  b <- subst' b (lift [(_ ** (var, (s ** Refl)))] f)
  pure (Pi s s' var a b)
substCnstr (Lambda s var t) f = do
  t <- subst' t (lift [(_ ** (var, (s ** Refl)))] f)
  pure (Lambda s var t)
substCnstr (Sigma s s' var a b) f = do
  a <- subst' a f
  b <- subst' b (lift [(_ ** (var, (s ** Refl)))] f)
  pure (Sigma s s' var a b)
substCnstr (Pair s s' prf t u) f = do
  t <- subst' t f
  u <- subst' u f
  pure (Pair s s' prf t u)
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 sort prf i) f = case f i of
  Left t => t
  Right j => pure (Ntrl $ Var var sort prf j)
substNtrl (App r t u) f = do
  t <- substNtrl t f
  u <- subst' u f
  assert_total (doApp t u)
substNtrl (Fst r t) f = do
  t <- substNtrl t f
  doFst Relevant r t
substNtrl (Snd r t) f = do
  t <- substNtrl t f
  doSnd r Relevant $ rewrite pairRelevantRight r in t
substNtrl (If t u v) f = do
  t <- substNtrl t f
  u <- subst' u f
  v <- subst' v f
  doIf t u v
substNtrl Absurd f = pure (doAbsurd Relevant)
substNtrl (Equal a t u) f = do
  a <- substNtrl a f
  t <- subst' t f
  u <- subst' u f
  doEqual _ a t u
substNtrl (EqualL a b) f = do
  a <- substNtrl a f
  b <- subst' b f
  doEqualSet a b
substNtrl (EqualR a b) f = do
  a <- substCnstr a f
  b <- substNtrl b f
  doEqualR a b
substNtrl (EqualStuck a b) f = do
  a <- substCnstr a f
  b <- substCnstr b f
  pure (Ntrl $ EqualStuck a b)
substNtrl (CastL a b t) f = do
  a <- substNtrl a f
  b <- subst' b f
  t <- subst' t f
  doCast _ a b t
substNtrl (CastR a b t) f = do
  a <- substCnstr a f
  b <- substNtrl b f
  t <- subst' t f
  doCastR a b t
substNtrl (CastStuck a b t) f = do
  a <- substCnstr a f
  b <- substCnstr b f
  t <- subst' t f
  pure (Ntrl $ CastStuck a b t)

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)

-- Utilities -------------------------------------------------------------------

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)