module Soat.SecondOrder.Algebra.Lift

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

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

import Syntax.PreorderReasoning.Setoid

%default total
%ambiguity_depth 4

public export
project : SecondOrder.Algebra.RawAlgebra (lift sig) -> (ctx : List sig.T)
  -> FirstOrder.Algebra.RawAlgebra sig
project a ctx = MkRawAlgebra
  (\t => a.U t ctx)
  (\op => a.sem ctx (MkOp (Op op.op)) . wrap (MkPair []))

public export
projectAlgebra : (0 sig : _) -> Algebra (lift sig) -> (ctx : List sig.T) -> Algebra sig
projectAlgebra sig a ctx = MkAlgebra
  { raw = project a.raw ctx
  , algebra = MkIsAlgebra
    { equivalence = (reindex (flip MkPair ctx) a.setoid).equivalence
    , semCong     = \op => a.algebra.semCong ctx (MkOp (Op op.op)) . wrapIntro
    }
  }

public export
projectHomo : {a, b : SecondOrder.Algebra.Algebra (lift sig)} -> a ~> b
  -> (ctx : _) -> projectAlgebra sig a ctx ~> projectAlgebra sig b ctx
projectHomo f ctx = MkHomomorphism
  { func    = MkIndexedSetoidHomomorphism
    { H           = \t => f.func.H (t, ctx)
    , homomorphic = \t => f.func.homomorphic (t, ctx)
    }
  , semHomo = \op, tms => CalcWith (index b.setoid _) $
    |~ f.func.H (_, ctx) (a.raw.sem ctx (MkOp (Op op.op)) (wrap (MkPair []) tms))
    ~~ b.raw.sem ctx (MkOp (Op op.op)) (map (\ty => f.func.H (snd ty, fst ty ++ ctx)) (wrap (MkPair []) tms))
      ...(f.semHomo ctx (MkOp (Op op.op)) (wrap (MkPair []) tms))
    ~~ b.raw.sem ctx (MkOp (Op op.op)) (wrap (MkPair []) (map (\t => f.func.H (t, ctx)) tms))
      .=.(cong (b.raw.sem ctx (MkOp (Op op.op))) $
          mapWrap (MkPair []) {f = \ty => f.func.H (snd ty, fst ty ++ ctx)} tms)
  }

public export
(.renameHomo) : (a : SecondOrder.Algebra.Algebra (lift sig)) -> {ctx, ctx' : _}
  -> (f : ctx `Sublist` ctx')
  -> projectAlgebra sig a ctx ~> projectAlgebra sig a ctx'
(.renameHomo) a f = MkHomomorphism
  { func    = MkIndexedSetoidHomomorphism
    { H           = \t => a.raw.rename t f
    , homomorphic = \t, _, _ => a.algebra.renameCong t f
    }
  , semHomo = \op, tms => CalcWith (index a.setoid _) $
    |~ a.raw.rename _ f (a.raw.sem _ (MkOp (Op op.op)) (wrap (MkPair []) tms))
    ~~ a.raw.sem _ (MkOp (Op op.op)) (map (\ty => a.raw.rename (snd ty) (reflexive {x = fst ty} ++ f)) (wrap (MkPair []) tms))
      ...(a.algebra.semNat f (MkOp (Op op.op)) (wrap (MkPair []) tms))
    ~~ a.raw.sem _ (MkOp (Op op.op)) (wrap (MkPair []) (map (\t => a.raw.rename t f) tms))
      ...(a.algebra.semCong _ (MkOp (Op op.op)) $
          CalcWith (index (pwSetoid (reindex (\ty => (snd ty, fst ty ++ _)) a.setoid)) _) $
          |~ map (\ty => a.raw.rename (snd ty) (reflexive {x = fst ty} ++ f)) (wrap (MkPair []) tms)
          ~~ wrap (MkPair []) (map (\t => a.raw.rename t (reflexive {x = []} ++ f)) tms)
            .=.(mapWrap (MkPair []) tms)
          ~~ wrap (MkPair []) (map (\t => a.raw.rename t f) tms)
            ...(wrapIntro $
                mapIntro'' (\t, tm, tm', eq =>
                  CalcWith (index a.setoid (t, _)) $
                  |~ a.raw.rename t (reflexive {x = []} ++ f) tm
                  ~~ a.raw.rename t f tm
                    .=.(cong (\f => a.raw.rename t f tm) $ uncurryCurry f)
                  ~~ a.raw.rename t f tm'
                    ...(a.algebra.renameCong t f eq)) $
                (pwSetoid (a.setoidAt ctx)).equivalence.reflexive _ tms))
  }

public export
(.substHomo1) : (a : SecondOrder.Algebra.Algebra (lift sig)) -> (ctx : List sig.T)
  -> {ctx' : List sig.T} -> (tms : (\t => a.raw.U t ctx) ^ ctx')
  -> projectAlgebra sig a ctx' ~> projectAlgebra sig a ctx
(.substHomo1) a ctx tms = MkHomomorphism
  { func    = MkIndexedSetoidHomomorphism
    { H           = \t, tm => a.raw.subst t ctx tm tms
    , homomorphic = \t, _, _, eq =>
      a.algebra.substCong t ctx eq $
      (pwSetoid (a.setoidAt _)).equivalence.reflexive _ tms
    }
  , semHomo = \op, tms' => CalcWith (index a.setoid _) $
    |~ a.raw.subst _ ctx (a.raw.sem ctx' (MkOp (Op op.op)) (wrap (MkPair []) tms')) tms
    ~~ a.raw.sem ctx (MkOp (Op op.op)) (map (a.raw.extendSubst ctx tms) (wrap (MkPair []) tms'))
      ...(a.algebra.substCompat ctx (MkOp (Op op.op)) (wrap (MkPair []) tms') tms)
    ~~ a.raw.sem ctx (MkOp (Op op.op)) (wrap (MkPair []) (map (\t, tm => a.raw.subst t ctx tm tms) tms'))
      ...(a.algebra.semCong ctx (MkOp (Op op.op)) $
          CalcWith (index (pwSetoid (reindex (\ty => (snd ty, fst ty ++ ctx)) a.setoid)) _) $
          |~ map (a.raw.extendSubst ctx tms) (wrap (MkPair []) tms')
          ~~ wrap (MkPair []) (map (\t, tm => a.raw.subst t ctx tm (map (\t => a.raw.rename t ([] {ys = []} ++ reflexive)) tms)) tms')
            .=.(mapWrap (MkPair []) tms')
          ~~ wrap (MkPair []) (map (\t, tm => a.raw.subst t ctx tm tms) tms')
            ...(wrapIntro $
                mapIntro' (\t, eq =>
                  a.algebra.substCong t ctx eq $
                  CalcWith (index (pwSetoid (a.setoidAt _)) _) $
                  |~ map (\t => a.raw.rename t ([] {ys = []} ++ reflexive)) tms
                  ~~ map (\t => id) tms
                    ...(mapIntro'' (\t, tm, tm', eq =>
                          CalcWith (index a.setoid (t, ctx)) $
                          |~ a.raw.rename t ([] {ys = []} ++ reflexive) tm
                          ~~ a.raw.rename t reflexive tm
                            .=.(cong (\f => a.raw.rename t f tm) $ uncurryCurry reflexive)
                          ~~ tm
                            ...(a.algebra.renameId t ctx tm)
                          ~~ tm'
                            ...(eq)) $
                        (pwSetoid (a.setoidAt _)).equivalence.reflexive _ _)
                  ~~ tms
                    .=.(mapId tms)) $
                (pwSetoid (a.setoidAt _)).equivalence.reflexive _ _))
  }

-- renameBodyFunc : (f : ctx `Sublist` ctx')
--   -> irrelevantCast (flip Elem ctx) ~>
--      (FreeAlgebra {sig = sig} (irrelevantCast (flip Elem ctx'))).setoid
-- renameBodyFunc f = mate (\_ => Done . curry f)

-- indexFunc : {ctx : List sig.T} -> (tms : Term sig (flip Elem ctx) ^ ts)
--   -> irrelevantCast (flip Elem ts) ~>
--      (FreeAlgebra {sig = sig} (irrelevantCast (flip Elem ctx))).setoid
-- indexFunc tms = mate (\_ => index tms)

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

public export
Initial : (0 sig : _) -> SecondOrder.Algebra.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 (pwSetoid (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) _)) $
          (pwSetoid (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 (pwSetoid (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' _
                      }))) $
          (pwSetoid (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 (pwSetoid ((InitialAlgebra sig).setoidAt ctx)) _) $
    sym $
    trans (unwrapWrap (extend (Initial sig).U ctx) _) (mapId tms)
  }

public export
fromInitial : (a : SecondOrder.Algebra.RawAlgebra (lift sig)) -> (t : sig.T) -> (ctx : List sig.T)
  -> (Initial sig).U t ctx -> a.U t ctx
fromInitial a t ctx = bindTerm {a = project a ctx} (\_ => a.var)

public export
fromInitialHomo : (a : Algebra (lift sig)) -> InitialAlgebra sig ~> a
fromInitialHomo a = MkHomomorphism
  { func       = MkIndexedSetoidHomomorphism
    { H           = \(t, ctx) => fromInitial a.raw t ctx
    , homomorphic = \(t, ctx), _, _ => bindTermCong {a = projectAlgebra sig a ctx} (a.varFunc ctx)
    }
  , renameHomo = \t, f => bindUnique'
    {v = irrelevantCast (flip Elem _)}
    {a = projectAlgebra sig a _}
    (bindHomo (a.varFunc _) . bindHomo (mate (\_ => Done . curry f)))
    (a.renameHomo f . bindHomo (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 (pwSetoid (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.raw t ctx) (unwrap (MkPair []) tms))
      .=.(cong (wrap _) $ bindTermsIsMap {a = project a.raw ctx} _ _)
    ~~ wrap (MkPair []) (unwrap (MkPair []) (map (\ty => fromInitial a.raw (snd ty) (fst ty ++ ctx)) tms))
      .=.(cong (wrap _) $ mapUnwrap (MkPair []) tms)
    ~~ map (\ty => fromInitial a.raw (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 (a.varFunc _) . bindHomo (mate (\_ => index tms)))
    (a.substHomo1 ctx _ . bindHomo (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.raw t ctx tm)
fromInitialUnique {sig = sig} {a = a} f t ctx = bindUnique
  {v = irrelevantCast (flip Elem _)}
  {a = projectAlgebra sig a ctx}
  (a.varFunc ctx)
  (projectHomo f ctx . freeToInitialHomo sig ctx)
  f.varHomo