module Obs.Typing.Conversion

import Control.Monad.Maybe
import Control.Monad.Trans

import Data.Bool
import Data.Nat

import Decidable.Equality

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

%default total

-- Conversion ------------------------------------------------------------------

export
convert : (s : Universe)
  -> (a : TypeNormalForm ctx)
  -> (t, u : NormalForm (relevance s) ctx)
  -> MaybeT (Logging ann) (NormalForm (relevance s) ctx)

untypedConvert : (t, u : NormalForm Relevant ctx) -> MaybeT (Logging ann) (NormalForm Relevant ctx)
-- Diagonals
untypedConvert (Cnstr (Universe s)) (Cnstr (Universe s')) =
  if s == s' then pure (Cnstr (Universe s)) else nothing
untypedConvert (Cnstr (Pi s s' var a b)) (Cnstr (Pi l l' _ a' b')) = do
  let Yes Refl = decEq s l
    | No _ => nothing
  lift $ trace "going back to type-directed conversion"
  a <- assert_total (convert (succ s) (cast s) a a')
  b <- assert_total (convert (succ s') (cast s') b b')
  pure (Cnstr $ Pi s s' var a b)
untypedConvert (Cnstr (Lambda s var t)) (Cnstr (Lambda s' _ u)) = do
  let Yes Refl = decEq s s'
    | No _ => nothing
  lift $ trace "converting under lambda"
  t <- untypedConvert t u
  pure (Cnstr $ Lambda s var t)
untypedConvert (Cnstr (Sigma s s' var a b)) (Cnstr (Sigma l l' _ a' b')) = do
  let Yes Refl = decEq s l
    | No _ => nothing
  lift $ trace "going back to type-directed conversion"
  a <- assert_total (convert (succ s) (cast s) a a')
  b <- assert_total (convert (succ s') (cast s') b b')
  pure (Cnstr $ Sigma s s' var a b)
untypedConvert (Cnstr (Pair s@(Set _) s'@(Set _) prf t u)) (Cnstr (Pair l l' prf' t' u')) = do
  let Yes Refl = decEq (s, s') (l, l')
    | No _ => nothing
  lift $ trace "converting pair pointwise"
  t <- untypedConvert t t'
  u <- untypedConvert u u'
  pure (Cnstr $ Pair l l' Set t u)
untypedConvert (Cnstr (Pair s@(Set _) Prop prf t u)) (Cnstr (Pair l Prop prf' t' u')) = do
  let Yes Refl = decEq s l
    | No _ => nothing
  lift $ trace "converting pair pointwise"
  t <- untypedConvert t t'
  pure (Cnstr $ Pair l Prop Set t Irrel)
untypedConvert (Cnstr (Pair Prop s'@(Set _) prf t u)) (Cnstr (Pair Prop l' prf' t' u')) = do
  let Yes Refl = decEq s' l'
    | No _ => nothing
  lift $ trace "converting pair pointwise"
  u <- untypedConvert u u'
  pure (Cnstr $ Pair Prop l' Set Irrel u)
untypedConvert (Cnstr Bool) (Cnstr Bool) = pure (Cnstr Bool)
untypedConvert (Cnstr True) (Cnstr True) = pure (Cnstr True)
untypedConvert (Cnstr False) (Cnstr False) = pure (Cnstr False)
untypedConvert (Cnstr Top) (Cnstr Top) = pure (Cnstr Top)
untypedConvert (Cnstr Bottom) (Cnstr Bottom) = pure (Cnstr Bottom)
untypedConvert t@(Ntrl (Var _ _ _ i)) (Ntrl (Var _ _ _ j)) =
  if elemToNat i == elemToNat j then pure t else nothing
untypedConvert lhs@(Ntrl (App r t u)) rhs@(Ntrl (App r' t' u')) = do
  let Yes Refl = decEq r r'
    | No _ => nothing
  lift $ trace "checking parts of a spine"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  let Relevant = r'
    | Irrelevant => lift $ doApp t Irrel
  u <- untypedConvert u u'
  lift $ doApp t u
untypedConvert lhs@(Ntrl (Fst r t)) rhs@(Ntrl (Fst r' t')) = do
  lift $ trace "checking full pair for fst"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  lift $ doFst Relevant r t
untypedConvert lhs@(Ntrl (Snd r t)) rhs@(Ntrl (Snd r' t')) = do
  lift $ trace "checking full pair for snd"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  lift $ doSnd r Relevant (rewrite pairRelevantRight r in t)
untypedConvert lhs@(Ntrl (If t f g)) rhs@(Ntrl (If t' f' g')) = do
  lift $ trace "checking all branches of if"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  f <- untypedConvert f f'
  g <- untypedConvert g g'
  lift $ doIf t f g
untypedConvert (Ntrl Absurd) (Ntrl Absurd) = pure (Ntrl Absurd)
untypedConvert lhs@(Ntrl (Equal a t u)) rhs@(Ntrl (Equal a' t' u')) = do
  lift $ trace "checking equal (stuck on type)"
  a <- untypedConvert (assert_smaller lhs (Ntrl a)) (assert_smaller rhs (Ntrl a'))
  t <- untypedConvert t t'
  u <- untypedConvert u u'
  lift $ doEqual _ a t u
untypedConvert lhs@(Ntrl (EqualL t u)) rhs@(Ntrl (EqualL t' u')) = do
  lift $ trace "checking equal (stuck on left)"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  u <- untypedConvert u u'
  lift $ doEqualSet t u
untypedConvert lhs@(Ntrl (EqualR t u)) rhs@(Ntrl (EqualR t' u')) = do
  lift $ trace "checking equal (stuck on right)"
  t <- untypedConvert (assert_smaller lhs (Cnstr t)) (assert_smaller rhs (Cnstr t'))
  u <- untypedConvert (assert_smaller lhs (Ntrl u)) (assert_smaller rhs (Ntrl u'))
  lift $ doEqualSet t u
untypedConvert lhs@(Ntrl (EqualStuck t u)) rhs@(Ntrl (EqualStuck t' u')) = do
  lift $ trace "checking equal (head mismatch)"
  t <- untypedConvert (assert_smaller lhs (Cnstr t)) (assert_smaller rhs (Cnstr t'))
  u <- untypedConvert (assert_smaller lhs (Cnstr u)) (assert_smaller rhs (Cnstr u'))
  lift $ doEqualSet t u
untypedConvert lhs@(Ntrl (CastL a b t)) rhs@(Ntrl (CastL a' b' t')) = do
  lift $ trace "checking cast (stuck on left)"
  a <- untypedConvert (assert_smaller lhs (Ntrl a)) (assert_smaller rhs (Ntrl a'))
  b <- untypedConvert b b'
  t <- untypedConvert t t'
  lift $ doCast _ a b t
untypedConvert lhs@(Ntrl (CastR a b t)) rhs@(Ntrl (CastR a' b' t')) = do
  lift $ trace "checking cast (stuck on right)"
  a <- untypedConvert (assert_smaller lhs (Cnstr a)) (assert_smaller rhs (Cnstr a'))
  b <- untypedConvert (assert_smaller lhs (Ntrl b)) (assert_smaller rhs (Ntrl b'))
  t <- untypedConvert t t'
  lift $ doCast _ a b t
untypedConvert lhs@(Ntrl (CastStuck a b t)) rhs@(Ntrl (CastStuck a' b' t')) = do
  lift $ trace "checking cast (stuck on left)"
  a <- untypedConvert (assert_smaller lhs (Cnstr a)) (assert_smaller rhs (Cnstr a'))
  b <- untypedConvert (assert_smaller lhs (Cnstr b)) (assert_smaller rhs (Cnstr b'))
  t <- untypedConvert t t'
  lift $ doCast _ a b t
-- Pi types
untypedConvert t@(Cnstr (Lambda s var _)) (Ntrl u) = do
  u <- lift $ doApp (Ntrl $ weaken [relevance s] u) (point (var, (s ** Refl)) Here)
  assert_total (untypedConvert t (Cnstr (Lambda s var u)))
untypedConvert (Ntrl t) u@(Cnstr (Lambda s var _)) = do
  t <- lift $ doApp (Ntrl $ weaken [relevance s] t) (point (var, (s ** Refl)) Here)
  assert_total (untypedConvert (Cnstr (Lambda s var t)) u)
-- Sigma types
untypedConvert t@(Cnstr (Pair s s' prf _ _)) t'@(Ntrl _) =
  let t'' : NormalForm (pair (relevance s) (relevance s')) ctx
      t'' = rewrite trans (sym $ pairRelevance s s') (forgetIsRelevant prf) in t'
  in do
  t' <- lift $ doFst (relevance s) (relevance s') t''
  u' <- lift $ doSnd (relevance s) (relevance s') t''
  assert_total (untypedConvert t (Cnstr (Pair s s' prf t' u')))
untypedConvert t@(Ntrl _) t'@(Cnstr (Pair s s' prf _ _)) =
  let t'' : NormalForm (pair (relevance s) (relevance s')) ctx
      t'' = rewrite trans (sym $ pairRelevance s s') (forgetIsRelevant prf) in t
  in do
  t <- lift $ doFst (relevance s) (relevance s') t''
  u <- lift $ doSnd (relevance s) (relevance s') t''
  assert_total (untypedConvert (Cnstr (Pair s s' prf t u)) t')
-- Bools
-- no cases because diagonals complete
-- Fallback
untypedConvert _ _ = nothing

convert Prop a Irrel Irrel = pure Irrel
convert (Set k) (Ntrl a) t u = do
  lift $ trace $ pretty {ann} "cannot do type-directed conversion on" <++> pretty a
  untypedConvert t u
convert (Set k) (Cnstr (Universe s)) t u = do
  lift $ trace $ pretty {ann} "converting in type" <++> pretty s
  untypedConvert t u
convert (Set k) (Cnstr (Pi s s'@(Set _) var a b)) t u = do
  -- s' must be Set, otherwise ty is in Prop
  lift $ trace "converting pi type"
  t' <- lift $ doApp (Sorted.weaken [relevance s] t) (point (var, (s ** Refl)) Here)
  u' <- lift $ doApp (Sorted.weaken [relevance s] u) (point (var, (s ** Refl)) Here)
  t <- convert s' b t' u'
  pure (Cnstr $ Lambda s var t)
convert (Set k) ty@(Cnstr (Sigma s@(Set j) s' var a b)) t u = do
  lift $ trace "converting sigma type"
  t1 <- lift $ doFst Relevant (relevance s') t
  u1 <- lift $ doFst Relevant (relevance s') u
  t1 <- convert s a t1 u1
  t2 <- lift $ doSnd Relevant (relevance s') t
  u2 <- lift $ doSnd Relevant (relevance s') u
  b <- lift $ subst1 t1 b
  t2 <- convert s' (assert_smaller ty b) t2 u2
  pure (Cnstr $ Pair s s' Set t1 t2)
convert (Set k) ty@(Cnstr (Sigma Prop s'@(Set _) var a b)) t u = do
  lift $ trace "converting sigma type (snd only)"
  t2 <- lift $ doSnd Irrelevant Relevant t
  u2 <- lift $ doSnd Irrelevant Relevant u
  b <- lift $ subst1 Irrel b
  t2 <- convert s' (assert_smaller ty b) t2 u2
  pure (Cnstr $ Pair Prop s' Set Irrel t2)
convert (Set 0) (Cnstr Bool) (Cnstr True) (Cnstr True) = do
  lift $ trace "converting bool (true)"
  pure (Cnstr True)
convert (Set 0) (Cnstr Bool) (Cnstr False) (Cnstr False) = do
  lift $ trace "converting bool (false)"
  pure (Cnstr False)
convert (Set k) (Cnstr Bool) t@(Ntrl _) u@(Ntrl _) = do
  lift $ trace "converting bool (neutral)"
  untypedConvert t u
convert (Set k) (Cnstr a) t u = lift $ inScope "bad type constructor" $ fatal a