module Soat.FirstOrder.Term

import Data.Setoid.Indexed

import Soat.Data.Product
import Soat.FirstOrder.Algebra
import Soat.FirstOrder.Signature

import Syntax.PreorderReasoning.Setoid

%default total

public export
data Term : (0 sig : Signature) -> (0 _ : sig.T -> Type) -> sig.T -> Type where
  Done : (i : v t) -> Term sig v t
  Call : (op : Op sig t) -> (ts : Term sig v ^ op.arity) -> Term sig v t

public export
bindTerm : {auto a : RawAlgebra sig} -> ((t : sig.T) -> v t -> a.U t)
  -> {t : _} -> Term sig v t -> a.U t

public export
bindTerms : {auto a : RawAlgebra sig} -> ((t : sig.T) -> v t -> a.U t)
  -> {ts : _} -> Term sig v ^ ts -> a.U ^ ts

bindTerm         env (Done i)     = env _ i
bindTerm {a = a} env (Call op ts) = a.sem op (bindTerms env ts)

bindTerms env []        = []
bindTerms env (t :: ts) = bindTerm env t :: bindTerms env ts

public export
bindTermsIsMap : {auto a : RawAlgebra sig}
  -> (env : (t : sig.T) -> v t -> a.U t) -> {ts : _} -> (tms : Term sig v ^ ts)
  -> bindTerms (\tm => env tm) tms = map (\t => bindTerm (\tm => env tm)) tms
bindTermsIsMap env []        = Refl
bindTermsIsMap env (t :: ts) = cong ((::) (bindTerm env t)) (bindTermsIsMap env ts)

namespace Rel
  public export
  data (~=~) : forall v . (0 r : (t : sig.T) -> Rel (v t)) -> (t : sig.T) -> Rel (Term sig v t)
    where
    Done : {0 v : sig.T -> Type} -> {0 r : (t : sig.T) -> Rel (v t)}
      -> {t : sig.T} -> {i, j : v t}
      -> r t i j
      -> (~=~) {sig} {v} r t (Done i) (Done j)
    Call : {0 v : sig.T -> Type} -> {0 r : (t : sig.T) -> Rel (v t)}
      -> {t : sig.T} -> (op : Op sig t) -> {tms, tms' : Term sig v ^ op.arity}
      -> Pointwise ((~=~) {sig} {v} r) tms tms'
      -> (~=~) {sig} {v} r t (Call op tms) (Call op tms')

  tmRelRefl : {0 V : sig.T -> Type} -> {0 rel : (t : sig.T) -> Rel (V t)}
    -> (refl : (t : sig.T) -> (x : V t) -> rel t x x)
    -> {t : sig.T} -> (tm : Term sig V t) -> (~=~) rel t tm tm

  tmsRelRefl : {0 V : sig.T -> Type} -> {0 rel : (t : sig.T) -> Rel (V t)}
    -> (refl : (t : sig.T) -> (x : V t) -> rel t x x)
    -> {ts : List sig.T} -> (tms : Term sig V ^ ts) -> Pointwise ((~=~) rel) tms tms

  tmRelRefl refl (Done i)     = Done (refl _ i)
  tmRelRefl refl (Call op ts) = Call op (tmsRelRefl refl ts)

  tmsRelRefl refl []        = []
  tmsRelRefl refl (t :: ts) = tmRelRefl refl t :: tmsRelRefl refl ts

  tmRelSym : {0 V : sig.T -> Type} -> {0 rel : (t : sig.T) -> Rel (V t)}
    -> (sym : (t : sig.T) -> (x, y : V t) -> rel t x y -> rel t y x)
    -> {t : sig.T} -> {tm, tm' : Term sig V t} -> (~=~) rel t tm tm' -> (~=~) rel t tm' tm

  tmsRelSym : {0 V : sig.T -> Type} -> {0 rel : (t : sig.T) -> Rel (V t)}
    -> (sym : (t : sig.T) -> (x, y : V t) -> rel t x y -> rel t y x)
    -> {ts : List sig.T} -> {tms, tms' : Term sig V ^ ts}
    -> Pointwise ((~=~) rel) tms tms' -> Pointwise ((~=~) rel) tms' tms

  tmRelSym sym (Done i)     = Done (sym _ _ _ i)
  tmRelSym sym (Call op ts) = Call op (tmsRelSym sym ts)

  tmsRelSym sym []        = []
  tmsRelSym sym (t :: ts) = tmRelSym sym t :: tmsRelSym sym ts

  tmRelTrans : {0 V : sig.T -> Type} -> {0 rel : (t : sig.T) -> Rel (V t)}
    -> (trans : (t : sig.T) -> (x, y, z : V t) -> rel t x y -> rel t y z -> rel t x z)
    -> {t : sig.T} -> {tm, tm', tm'' : Term sig V t}
    -> (~=~) rel t tm tm' -> (~=~) rel t tm' tm'' -> (~=~) rel t tm tm''

  tmsRelTrans : {0 V : sig.T -> Type} -> {0 rel : (t : sig.T) -> Rel (V t)}
    -> (trans : (t : sig.T) -> (x, y, z : V t) -> rel t x y -> rel t y z -> rel t x z)
    -> {ts : List sig.T} -> {tms, tms', tms'' : Term sig V ^ ts}
    -> Pointwise ((~=~) rel) tms tms' -> Pointwise ((~=~) rel) tms' tms''
    -> Pointwise ((~=~) rel) tms tms''

  tmRelTrans trans (Done i)     (Done j)     = Done (trans _ _ _ _ i j)
  tmRelTrans trans (Call op ts) (Call _ ts') = Call op (tmsRelTrans trans ts ts')

  tmsRelTrans trans []        []          = []
  tmsRelTrans trans (t :: ts) (t' :: ts') = tmRelTrans trans t t' :: tmsRelTrans trans ts ts'

  public export
  tmRelEq : IndexedEquivalence sig.T v -> IndexedEquivalence sig.T (Term sig v)
  tmRelEq eq = MkIndexedEquivalence
    { relation   = (~=~) eq.relation
    , reflexive  = \t => tmRelRefl eq.reflexive
    , symmetric  = \t, x, y => tmRelSym eq.symmetric
    , transitive = \t, x, y, z => tmRelTrans eq.transitive
    }

public export
Free : (0 V : sig.T -> Type) -> RawAlgebra sig
Free v = MkRawAlgebra (Term sig v) Call

public export
FreeIsAlgebra : (v : IndexedSetoid sig.T) -> IsAlgebra sig (Free v.U)
FreeIsAlgebra v = MkIsAlgebra (tmRelEq v.equivalence) Call

public export
FreeAlgebra : IndexedSetoid sig.T -> Algebra sig
FreeAlgebra v = MkAlgebra _ (FreeIsAlgebra v)

public export
bindTermCong' : {a : Algebra sig} -> {env, env' : (t : sig.T) -> v t -> a.raw.U t}
  -> {0 rel : (t : sig.T) -> Rel (v t)}
  -> (cong : (t : sig.T) -> {x, y : v t} -> rel t x y -> a.relation t (env t x) (env' t y))
  -> {t : sig.T} -> {tm, tm' : Term sig v t} -> (eq : (~=~) rel t tm tm')
  -> a.relation t (bindTerm {a = a.raw} env tm) (bindTerm {a = a.raw} env' tm')

public export
bindTermsCong' : {a : Algebra sig} -> {env, env' : (t : sig.T) -> v t -> a.raw.U t}
  -> {0 rel : (t : sig.T) -> Rel (v t)}
  -> (cong : (t : sig.T) -> {x, y : v t} -> rel t x y -> a.relation t (env t x) (env' t y))
  -> {ts : List sig.T} -> {tms, tms' : Term sig v ^ ts} -> (eqs : Pointwise ((~=~) rel) tms tms')
  -> Pointwise a.relation (bindTerms {a = a.raw} env tms) (bindTerms {a = a.raw} env' tms')

bindTermCong'         cong (Done i)     = cong _ i
bindTermCong' {a = a} cong (Call op ts) = a.algebra.semCong op (bindTermsCong' cong ts)

bindTermsCong' cong []        = []
bindTermsCong' cong (t :: ts) = bindTermCong' cong t :: bindTermsCong' cong ts

public export
bindTermCong : {a : Algebra sig} -> (env : v ~> a.setoid)
  -> {t : sig.T} -> {tm, tm' : Term sig v.U t} -> (eq : (~=~) v.equivalence.relation t tm tm')
  -> a.relation t (bindTerm {a = a.raw} env.H tm) (bindTerm {a = a.raw} env.H tm')
bindTermCong env = bindTermCong' env.homomorphic

public export
bindTermsCong : {a : Algebra sig} -> (env : v ~> a.setoid)
  -> {ts : List sig.T} -> {tms, tms' : Term sig v.U ^ ts}
  -> (eqs : Pointwise ((~=~) v.equivalence.relation) tms tms')
  -> Pointwise a.relation (bindTerms {a = a.raw} env.H tms) (bindTerms {a = a.raw} env.H tms')
bindTermsCong env = bindTermsCong' env.homomorphic

public export
bindHomo : {a : Algebra sig} -> (env : v ~> a.setoid) -> FreeAlgebra v ~> a
bindHomo env = MkHomomorphism
  { func    = MkIndexedSetoidHomomorphism
    { H           = \_ => bindTerm {a = a.raw} env.H
    , homomorphic = \_, _, _ => bindTermCong env
    }
  , semHomo = \op, tms =>
    a.algebra.semCong op $
    reflect (index (pwSetoid a.setoid) _) $
    bindTermsIsMap {a = a.raw} env.H tms
  }

public export
bindUnique' : {v : IndexedSetoid sig.T} -> {a : Algebra sig}
  -> (f, g : FreeAlgebra v ~> a)
  -> (cong : {t : sig.T} -> (i : v.U t)
        -> a.relation t (f.func.H t (Done i)) (g.func.H t (Done i)))
  -> {t : sig.T} -> (tm : Term sig v.U t)
  -> a.relation t (f.func.H t tm) (g.func.H t tm)

public export
bindsUnique' : {v : IndexedSetoid sig.T} -> {a : Algebra sig}
  -> (f, g : FreeAlgebra v ~> a)
  -> (cong : {t : sig.T} -> (i : v.U t)
        -> a.relation t (f.func.H t (Done i)) (g.func.H t (Done i)))
  -> {ts : List sig.T} -> (tms : Term sig v.U ^ ts)
  -> Pointwise a.relation (map f.func.H tms) (map g.func.H tms)

bindUnique' f g cong (Done i)     = cong i
bindUnique' f g cong (Call op ts) = CalcWith (index a.setoid _) $
  |~ f.func.H _ (Call op ts)
  ~~ a.raw.sem op (map f.func.H ts) ...( f.semHomo op ts )
  ~~ a.raw.sem op (map g.func.H ts) ...( a.algebra.semCong op $ bindsUnique' f g cong ts )
  ~~ g.func.H _ (Call op ts)        ..<( g.semHomo op ts )

bindsUnique' f g cong []        = []
bindsUnique' f g cong (t :: ts) = bindUnique' f g cong t :: bindsUnique' f g cong ts

public export
bindUnique : {v : IndexedSetoid sig.T} -> {a : Algebra sig} -> (env : v ~> a.setoid)
  -> (f : FreeAlgebra v ~> a)
  -> (cong : {t : sig.T} -> (i : v.U t) -> a.relation t (f.func.H t (Done i)) (env.H t i))
  -> {t : sig.T} -> (tm : Term sig v.U t)
  -> a.relation t (f.func.H t tm) (bindTerm {a = a.raw} env.H tm)
bindUnique env f = bindUnique' f (bindHomo env)

public export
bindsUnique : {v : IndexedSetoid sig.T} -> {a : Algebra sig} -> (env : v ~> a.setoid)
  -> (f : FreeAlgebra v ~> a)
  -> (cong : {t : sig.T} -> (i : v.U t) -> a.relation t (f.func.H t (Done i)) (env.H t i))
  -> {ts : List sig.T} -> (tms : Term sig v.U ^ ts)
  -> Pointwise a.relation (map f.func.H tms) (bindTerms {a = a.raw} env.H tms)
bindsUnique env f cong ts = CalcWith (index (pwSetoid a.setoid) _) $
  |~ map f.func.H ts
  ~~ map (\_ => bindTerm {a = a.raw} env.H) ts ...( bindsUnique' f (bindHomo env) cong ts )
  ~~ bindTerms {a = a.raw} env.H ts            .=<( bindTermsIsMap {a = a.raw} env.H ts )