module Soat.SecondOrder.Algebra.Lift

import Data.List.Elem
import Data.Morphism.Indexed
import Data.Setoid.Indexed

import Soat.Data.Product
import Soat.Data.Sublist
import Soat.FirstOrder.Algebra
import Soat.FirstOrder.Term
import Soat.SecondOrder.Algebra
import Soat.SecondOrder.Signature.Lift

import Syntax.PreorderReasoning.Setoid

%default total

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

public export
projectIsAlgebra : IsAlgebra (lift sig) a -> (ctx : List sig.T) -> IsAlgebra sig (project a ctx)
projectIsAlgebra a ctx = MkIsAlgebra
  (\t => a.relation (t, ctx))
  (\t => a.equivalence (t, ctx))
  (\op => a.semCong ctx _ . wrapIntro)

public export
projectAlgebra : (0 sig : _) -> Algebra (lift sig) -> (ctx : List sig.T) -> Algebra sig
projectAlgebra sig a ctx = MkAlgebra _ (projectIsAlgebra a.algebra ctx)

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    = \t => f.func t ctx
  , cong    = \t => f.cong t ctx
  , semHomo = \op, tms => CalcWith (b.setoid.index _) $
    |~ f.func _ ctx (a.raw.sem ctx (MkOp (Op op.op)) (wrap (MkPair []) tms))
    ~~ b.raw.sem ctx (MkOp (Op op.op)) (map (extendFunc f.func 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 t ctx) tms))
      .=.(cong (b.raw.sem ctx (MkOp (Op op.op))) $ mapWrap (MkPair []) {f = extendFunc f.func 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    = \t => a.raw.rename t f
  , cong    = \t => a.algebra.renameCong t f
  , semHomo = \op, tms => CalcWith (a.setoid.index _) $
    |~ 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 (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)
            .=.(cong (wrap (MkPair [])) $
                pwEqImpliesEqual $
                mapIntro'' (\t, tm, _, Refl => cong (\f => a.raw.rename t f tm) $ uncurryCurry f) $
                equalImpliesPwEq Refl))
  }

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    = \t, tm => a.raw.subst t ctx tm tms
  , cong    = \t, eq =>
    a.algebra.substCong t ctx eq $
    pwRefl (\_ => (a.algebra.equivalence _).refl)
  , semHomo = \op, tms' => CalcWith (a.setoid.index _) $
    |~ 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 (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 (pwSetoid (a.setoidAt _) _) $
                  |~ map (\t => a.raw.rename t ([] {ys = []} ++ reflexive)) tms
                  ~~ map (\t => a.raw.rename t reflexive) tms
                    .=.(pwEqImpliesEqual $
                        mapIntro' (\t => cong2 (a.raw.rename t) $ uncurryCurry reflexive) $
                        equalImpliesPwEq Refl)
                  ~~ map (\t => id) tms
                    ...(mapIntro' (\t, Refl => a.algebra.renameId t ctx _) $
                        equalImpliesPwEq Refl)
                  ~~ tms
                    .=.(mapId tms)) $
                pwRefl (\t => (a.algebra.equivalence _).refl)))
  }

renameBodyFunc : (f : ctx `Sublist` ctx')
  -> IFunction
       (isetoid (flip Elem ctx))
       (FreeAlgebra {sig = sig} (isetoid (flip Elem ctx'))).setoid
renameBodyFunc f = MkIFunction (\_ => Done . curry f) (\_ => Done' . cong (curry f))

indexFunc : {ctx : List sig.T} -> (tms : Term sig (flip Elem ctx) ^ ts)
  -> IFunction
       (isetoid (flip Elem ts))
       (FreeAlgebra {sig = sig} (isetoid (flip Elem ctx))).setoid
indexFunc tms = MkIFunction
  (\_ => index tms)
  (\_ => ((FreeIsAlgebra (isetoid (flip Elem _))).equivalence _).equalImpliesEq . cong (index tms))

public export
Initial : (0 sig : _) -> SecondOrder.Algebra.RawAlgebra (lift sig)
Initial sig = MkRawAlgebra
  (\t, ctx => Term sig (flip Elem ctx) t)
  (\t, f => bindTerm {a = Free _} (renameBodyFunc f).func)
  (\_, (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
  { relation    = \(t, ctx) => (~=~) {sig = sig} {v = flip Elem ctx} (\_ => Equal) t
  , equivalence = \(t, ctx) => tmRelEq (\_ => equiv) t
  , renameCong  = \t, f => bindTermCong {a = FreeAlgebra (isetoid (flip Elem _))} (renameBodyFunc f)
  , semCong     = \_ , (MkOp (Op op)) => Call' (MkOp op) . unwrapIntro
  , substCong   = \_, _, eq, eqs => bindTermCong'
    {a = FreeAlgebra (isetoid (flip Elem _))}
    (\t, Refl => index eqs _)
    eq
  , renameId    = \t, ctx, tm =>
    tmRelSym (\_ => MkSymmetric symmetric) $
    bindUnique (renameBodyFunc reflexive) id (\i => Done' $ sym $ curryUncurry id i) $
    tm
  , renameComp  = \t, f, g, tm =>
    tmRelSym (\_ => MkSymmetric symmetric) $
    bindUnique
      {a = FreeAlgebra (isetoid (flip Elem _))}
      (renameBodyFunc (transitive g f))
      (bindHomo (renameBodyFunc f) . bindHomo (renameBodyFunc g))
      (\i => Done' $ sym $ curryUncurry (curry f . curry g) i) $
    tm
  , semNat      = \f, (MkOp (Op op)), tms =>
    Call' (MkOp op) $
    CalcWith (pwSetoid (FreeAlgebra (isetoid (flip Elem _))).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 = FreeAlgebra (isetoid (flip Elem _))}
              (\_, Refl => Done' $ curryUncurry (curry f) _)) $
          tmsRelRefl (\_ => MkReflexive 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 = FreeAlgebra (isetoid (flip Elem _))}
      (indexFunc _)
      (bindHomo (renameBodyFunc f) . bindHomo (indexFunc tms))
      (\i =>
        tmRelReflexive (\_ => MkReflexive reflexive) $
        sym $
        indexMap tms i)
      tm
  , substExnat  = \t, ctx, f, tm, tms =>
    bindUnique
      {a = FreeAlgebra (isetoid (flip Elem _))}
      (indexFunc _)
      (bindHomo (indexFunc tms) . bindHomo (renameBodyFunc f))
      (\i =>
        tmRelReflexive (\_ => MkReflexive reflexive) $
        sym $
        indexShuffle f i)
      tm
  , substComm   = \t, ctx, tm, tms, tms' =>
    bindUnique
      {a = FreeAlgebra (isetoid (flip Elem _))}
      (indexFunc _)
      (bindHomo (indexFunc tms') . bindHomo (indexFunc tms))
      (\i =>
        tmRelReflexive (\_ => MkReflexive reflexive) $
        sym $
        indexMap tms i)
      tm
  , substVarL   = \_, _, _ => tmRelRefl (\_ => MkReflexive reflexive) _
  , substVarR   = \t, ctx, tm =>
    tmRelSym (\_ => MkSymmetric symmetric) $
    bindUnique
      {a = FreeAlgebra (isetoid (flip Elem _))}
      (indexFunc _)
      id
      (\i =>
        tmRelReflexive (\_ => MkReflexive reflexive) $
        sym $
        indexTabulate Done i)
      tm
  , substCompat = \ctx, (MkOp (Op op)), tms, tms' =>
    Call' (MkOp op) $
    CalcWith (pwSetoid (FreeAlgebra (isetoid (flip Elem _))).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 = FreeAlgebra (isetoid (flip Elem _))}
            (\t, Refl => CalcWith ((FreeAlgebra (isetoid (flip Elem _))).setoid.index _) $
              |~ index (map (\_ => bindTerm {a = Free _} (\_ => Done . curry ([] {ys = []} ++ reflexive))) tms') _
              ~~ bindTerm {a = Free _} (\_ => Done . curry ([] {ys = []} ++ reflexive)) (index tms' _)
                .=.(indexMap tms' _)
              ~~ index tms' _
                ..<(bindUnique
                      (renameBodyFunc ([] {ys = []} ++ reflexive))
                      id
                      (\i => Done' $ sym $ trans (curryUncurry _ i) (curryUncurry id i))
                      (index tms' _)))) $
          tmsRelRefl (\_ => MkReflexive 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 (isetoid (flip Elem ctx)) ~> projectAlgebra sig (InitialAlgebra sig) ctx
freeToInitialHomo sig ctx = MkHomomorphism
  { func    = \_ => id
  , cong    = \_ => id
  , semHomo = \(MkOp op), tms =>
    Call' (MkOp op) $
    tmsRelSym (\_ => MkSymmetric symmetric) $
    tmsRelReflexive (\_ => MkReflexive reflexive) $
    transitive (unwrapWrap _ _) (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       = fromInitial a.raw
  , cong       = \t , ctx => bindTermCong {a = projectAlgebra sig a ctx} (a.varFunc ctx)
  , renameHomo = \t, f => bindUnique'
    {v = isetoid (flip Elem _)}
    {a = projectAlgebra sig a _}
    (bindHomo (a.varFunc _) . bindHomo (renameBodyFunc 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 (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 (extendFunc (fromInitial a.raw) ctx) tms))
      .=.(cong (wrap _) $ mapUnwrap (MkPair []) tms)
    ~~ map (extendFunc (fromInitial a.raw) ctx) tms
      .=.(wrapUnwrap _)
  , varHomo    = \_ => (a.algebra.equivalence _).reflexive
  , substHomo  = \t, ctx, tm, tms => bindUnique'
    {v = isetoid (flip Elem _)}
    {a = projectAlgebra sig a _}
    (bindHomo (a.varFunc _) . bindHomo (indexFunc tms))
    (a.substHomo1 ctx _ . bindHomo (a.varFunc _))
    (\i => CalcWith (a.setoid.index _) $
        |~ 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 t ctx tm) (fromInitial a.raw t ctx tm)
fromInitialUnique {sig = sig} {a = a} f t ctx = bindUnique
  {v = isetoid (flip Elem _)}
  {a = projectAlgebra sig a ctx}
  (a.varFunc ctx)
  (projectHomo f ctx . freeToInitialHomo sig ctx)
  f.varHomo