module Soat.SecondOrder.Algebra.Lift.Initial

import Data.List.Elem
import Data.List.Sublist
import Data.Product
import Data.Setoid.Indexed
import Data.Setoid.Pair
import Data.Setoid.Product

import Soat.FirstOrder.Algebra
import Soat.FirstOrder.Term
import Soat.SecondOrder.Algebra
import Soat.SecondOrder.Algebra.Lift
import Soat.SecondOrder.Signature.Lift

import Syntax.PreorderReasoning.Setoid

%default total
%ambiguity_depth 4

public export
freeAlg : List sig.T -> FirstOrder.Algebra.Algebra sig
freeAlg ctx = FreeAlgebra (irrelevantCast (flip Elem ctx))

public export
Initial : (0 sig : _) -> RawAlgebra (lift sig)
Initial sig = MkRawAlgebra
  (\t, ctx => (freeAlg ctx).raw.U t)
  (\t, f => bindTerm {a = Free _} (\_ => Done . curry f))
  (\_, (MkOp (Op op)) => Call (MkOp op) . unwrap (MkPair []))
  Done
  (\_, _, t, ts => bindTerm {a = Free _} (\_ => index ts) t)

public export
InitialIsAlgebra : (0 sig : _) -> SecondOrder.Algebra.IsAlgebra (lift sig) (Initial sig)
InitialIsAlgebra sig = MkIsAlgebra
  { equivalence = MkIndexedEquivalence
    { relation   = \(t, ctx) => (freeAlg ctx).relation t
    , reflexive  = \(t, ctx) => (freeAlg ctx).algebra.equivalence.reflexive t
    , symmetric  = \(t, ctx) => (freeAlg ctx).algebra.equivalence.symmetric t
    , transitive = \(t, ctx) => (freeAlg ctx).algebra.equivalence.transitive t
    }
  , renameCong  = \t, f => bindTermCong
    { a   = freeAlg _
    , env = mate (\_ => Done . curry f)
    }
  , semCong     = \_ , (MkOp (Op op)) => Call (MkOp op) . unwrapIntro
  , substCong   = \_, _, eq, eqs => bindTermCong'
    { a    = freeAlg _
    , cong = \_, Refl => index eqs _
    , eq
    }
  , renameId    = \t, ctx, tm =>
    (freeAlg _).setoid.equivalence.symmetric t _ _ $
    bindUnique (mate (\_ => Done . curry reflexive)) id (\i => Done $ sym $ curryUncurry id i) tm
  , renameComp  = \t, f, g, tm =>
    (freeAlg _).setoid.equivalence.symmetric t _ _ $
    bindUnique
      { a    = freeAlg _
      , env  = mate (\_ => Done . curry (transitive g f))
      , f    = bindHomo (mate (\_ => Done . curry f)) . bindHomo (mate (\_ => Done . curry g))
      , cong = \i => Done $ sym $ curryUncurry (curry f . curry g) i
      , tm
      }
  , semNat      = \f, (MkOp (Op op)), tms =>
    Call (MkOp op) $
    CalcWith (index (Product (freeAlg _).setoid) _) $
    |~ bindTerms {a = Free _} (\_ => Done . curry f) (unwrap (MkPair []) tms)
    ~~ map (\_ => bindTerm {a = Free _} (\_ => Done . curry f)) (unwrap (MkPair []) tms)
      .=.(bindTermsIsMap {a = Free _} _ _)
    ~~ map (\_ => bindTerm {a = Free _} (\_ => Done . curry (reflexive {x = []} ++ f))) (unwrap (MkPair []) tms)
      ..<(mapIntro' (\t =>
            bindTermCong'
              {rel = \_ => Equal}
              {a = freeAlg _}
              (\_, Refl => Done $ curryUncurry (curry f) _)) $
          (Product (freeAlg _).setoid).equivalence.reflexive _ (unwrap (MkPair []) tms))
    ~~ unwrap (MkPair []) (map (\ty => bindTerm {a = Free _} (\_ => Done . curry (reflexive {x = fst ty} ++ f))) tms)
      .=.(mapUnwrap (MkPair []) tms)
  , varNat      = \_, _ => Done Refl
  , substNat    = \t, f, tm, tms => bindUnique
    { a = freeAlg _
    , env = mate (\_ => index $ map (\_ => bindTerm {a = Free _} (\_ => Done . curry f)) tms)
    , f   = bindHomo (mate (\_ => Done . curry f)) . bindHomo (mate (\_ => index tms))
    , cong = \i =>
      reflect (index (freeAlg _).setoid _) $
      sym $
      indexMap tms i
    , tm
    }
  , substExnat  = \t, ctx, f, tm, tms => bindUnique
    { a    = freeAlg _
    , env  = mate (\_ => index $ shuffle f tms)
    , f    = bindHomo (mate (\_ => index tms)) . bindHomo (mate (\_ => Done . curry f))
    , cong = \i =>
      reflect (index (freeAlg _).setoid _) $
      sym $
      indexShuffle f i
    , tm
    }
  , substComm   = \t, ctx, tm, tms, tms' => bindUnique
    { a    = freeAlg _
    , env  = mate (\_ => index $ map (\_ => bindTerm {a = Free _} (\_ => index tms')) tms)
    , f    = bindHomo (mate (\_ => index tms')) . bindHomo (mate (\_ => index tms))
    , cong = \i =>
      reflect (index (freeAlg _).setoid _) $
      sym $
      indexMap tms i
    , tm
    }
  , substVarL   = \_, _, _ => (freeAlg _).setoid.equivalence.reflexive _ _
  , substVarR   = \t, ctx, tm =>
    (freeAlg _).setoid.equivalence.symmetric t _ _ $
    bindUnique
      { v    = irrelevantCast (flip Elem ctx)
      , a    = freeAlg ctx
      , env  = mate (\_ => index $ tabulate (Done {sig = sig, v = flip Elem ctx}))
      , f    = id
      , cong = \i =>
        reflect (index (freeAlg ctx).setoid _) $
        sym $
        indexTabulate Done i
      , tm
      }
  , substCompat = \ctx, (MkOp (Op op)), tms, tms' =>
    Call (MkOp op) $
    CalcWith (index (Product (freeAlg _).setoid) _) $
    |~ bindTerms {a = Free _} (\_ => index tms') (unwrap (MkPair []) tms)
    ~~ map (\_ => bindTerm {a = Free _} (\_ => index tms')) (unwrap (MkPair []) tms)
      .=.(bindTermsIsMap {a = Free _} _ _)
    ~~ map (\t => (Initial sig).extendSubst ctx tms' ([], t)) (unwrap (MkPair []) tms)
      ..<(mapIntro' (\t => bindTermCong'
            {rel = \_ => Equal}
            {a = freeAlg _}
            (\t, Refl => CalcWith (index (freeAlg _).setoid _) $
              |~ index (map (\_ => bindTerm {a = Free _} (\_ => Done . curry ([] {ys = []} ++ reflexive))) tms') _
              ~~ bindTerm {a = Free _} (\_ => Done . curry ([] {ys = []} ++ reflexive)) (index tms' _)
                .=.(indexMap tms' _)
              ~~ index tms' _
                ..<(bindUnique
                      { env  = mate (\_ => Done . curry ([] {ys = []} ++ reflexive))
                      , f    = id
                      , cong = \i => Done $ sym $ trans (curryUncurry _ i) (curryUncurry id i)
                      , tm   = index tms' _
                      }))) $
          (Product (freeAlg _).setoid).equivalence.reflexive _ (unwrap (MkPair []) tms))
    ~~ unwrap (MkPair []) (map ((Initial sig).extendSubst ctx tms') tms)
      .=.(mapUnwrap (MkPair []) tms)
  }

public export
InitialAlgebra : (0 sig : _) -> SecondOrder.Algebra.Algebra (lift sig)
InitialAlgebra sig = MkAlgebra (Initial sig) (InitialIsAlgebra sig)

public export
freeToInitialHomo : (0 sig : _) -> (ctx : List sig.T)
  -> FreeAlgebra (irrelevantCast (flip Elem ctx)) ~>
     projectAlgebra sig (InitialAlgebra sig) ctx
freeToInitialHomo sig ctx = MkHomomorphism
  { func    = MkIndexedSetoidHomomorphism
    { H           = \_ => id
    , homomorphic = \_, _, _ => id
    }
  , semHomo = \(MkOp op), tms =>
    Call (MkOp op) $
    reflect (index (Product ((InitialAlgebra sig).setoidAt ctx)) _) $
    sym $
    trans (unwrapWrap (extend (Initial sig).U ctx) _) (mapId tms)
  }

public export
fromInitial : (a : Algebra (lift sig)) -> (InitialAlgebra sig).setoid ~> a.setoid
fromInitial a = bundle (\(t, ctx) =>
  index (bindFunc {a = projectAlgebra sig a ctx} (reindex (flip MkPair ctx) a.varFunc)) t)

public export
fromInitialHomo : (a : Algebra (lift sig)) -> InitialAlgebra sig ~> a
fromInitialHomo a = MkHomomorphism
  { func       = fromInitial a
  , renameHomo = \t, f => bindUnique'
    {v = irrelevantCast (flip Elem _)}
    {a = projectAlgebra sig a _}
    (bindHomo (reindex (flip MkPair _) a.varFunc) . bindHomo (mate (\_ => Done . curry f)))
    (a.renameHomo f . bindHomo (reindex (flip MkPair _) a.varFunc))
    (\i => a.algebra.equivalence.symmetric _ _ _ $ a.algebra.varNat f i)
  , semHomo    = \ctx, (MkOp (Op op)), tms =>
    a.algebra.semCong ctx (MkOp (Op op)) $
    CalcWith (index (Product (reindex (\ty => (snd ty, fst ty ++ ctx)) a.setoid)) _) $
    |~ wrap (MkPair []) (bindTerms {a = project a.raw ctx} (\_ => a.raw.var) (unwrap (MkPair []) tms))
    ~~ wrap (MkPair []) (map (\t => (fromInitial a).H (t, ctx)) (unwrap (MkPair []) tms))
      .=.(cong (wrap _) $ bindTermsIsMap {a = project a.raw ctx} _ _)
    ~~ wrap (MkPair []) (unwrap (MkPair []) (map (\ty => (fromInitial a).H (snd ty, fst ty ++ ctx)) tms))
      .=.(cong (wrap _) $ mapUnwrap (MkPair []) tms)
    ~~ map (\ty => (fromInitial a).H (snd ty, fst ty ++ ctx)) tms
      .=.(wrapUnwrap _)
  , varHomo    = \i => a.algebra.equivalence.reflexive _ $ a.raw.var i
  , substHomo  = \t, ctx, tm, tms => bindUnique'
    {v = irrelevantCast (flip Elem _)}
    {a = projectAlgebra sig a _}
    (bindHomo (reindex (flip MkPair _) a.varFunc) . bindHomo (mate (\_ => index tms)))
    (a.substHomo1 ctx _ . bindHomo (reindex (flip MkPair _) a.varFunc))
    (\i => CalcWith (index a.setoid _) $
        |~ bindTerm {a = project a.raw _} (\_ => a.raw.var) (index tms i)
        ~~ index (map (\_ => bindTerm {a = project a.raw _} (\_ => a.raw.var)) tms) i
          .=<(indexMap {f = \_ => bindTerm {a = project a.raw _} (\_ => a.raw.var)} tms i)
        ~~ a.raw.subst _ ctx (a.raw.var i) (map (\_ => bindTerm {a = project a.raw _} (\_ => a.raw.var)) tms)
          ..<(a.algebra.substVarL ctx i _))
    tm
  }

public export
fromInitialUnique : {a : SecondOrder.Algebra.Algebra (lift sig)}
  -> (f : InitialAlgebra sig ~> a)
  -> (t : sig.T) -> (ctx : List sig.T) -> (tm : Term sig (flip Elem ctx) t)
  -> a.relation (t, ctx) (f.func.H (t, ctx) tm) ((fromInitial a).H (t, ctx) tm)
fromInitialUnique {sig = sig} {a = a} f t ctx = bindUnique
  {v = irrelevantCast (flip Elem _)}
  {a = projectAlgebra sig a ctx}
  (reindex (flip MkPair ctx) a.varFunc)
  (projectHomo f ctx . freeToInitialHomo sig ctx)
  f.varHomo