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} -> (f : (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)

-- FIXME: these names shouldn't be public. Indication of bad API design
namespace Rel
  public export
  data (~=~) : forall v . (0 r : (t : _) -> Rel (v t)) -> (t : _) -> Rel (Term sig v t)
    where
    Done' : {0 v : sig.T -> Type} -> {0 r : (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 : _) -> 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')

  public export
  tmRelEqualIsEqual : (~=~) (\_ => Equal) t tm tm' -> tm = tm'

  public export
  tmsRelEqualIsEqual : Pointwise ((~=~) (\_ => Equal)) tms tms' -> tms = tms'

  tmRelEqualIsEqual (Done' eq)     = cong Done eq
  tmRelEqualIsEqual (Call' op eqs) = cong (Call op) $ tmsRelEqualIsEqual eqs

  tmsRelEqualIsEqual []          = Refl
  tmsRelEqualIsEqual (eq :: eqs) = cong2 (::) (tmRelEqualIsEqual eq) (tmsRelEqualIsEqual eqs)

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

  public export
  tmsRelRefl : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((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

  public export
  tmRelReflexive : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((t : sig.T) -> (x : V t) -> rel t x x)
    -> {t : sig.T} -> {tm, tm' : Term sig V t} -> tm = tm' -> (~=~) rel t tm tm'
  tmRelReflexive refl Refl = tmRelRefl refl _

  public export
  tmsRelReflexive : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((t : sig.T) -> (x : V t) -> rel t x x)
    -> {ts : List sig.T} -> {tms, tms' : Term sig V ^ ts}
    -> tms = tms' -> Pointwise ((~=~) rel) tms tms'
  tmsRelReflexive refl Refl = tmsRelRefl refl _

  public export
  tmRelSym : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((t : sig.T) -> Symmetric (V t) (rel t))
    -> {t : sig.T} -> {tm, tm' : Term sig V t} -> (~=~) rel t tm tm' -> (~=~) rel t tm' tm

  public export
  tmsRelSym : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((t : sig.T) -> Symmetric (V t) (rel t))
    -> {ts : List sig.T} -> {tms, tms' : Term sig V ^ ts}
    -> Pointwise ((~=~) rel) tms tms' -> Pointwise ((~=~) rel) tms' tms

  tmRelSym sym (Done' i)     = Done' (symmetric i)
  tmRelSym sym (Call' op ts) = Call' op (tmsRelSym sym ts)

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

  public export
  tmRelTrans : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((t : sig.T) -> Transitive (V t) (rel t))
    -> {t : sig.T} -> {tm, tm', tm'' : Term sig V t}
    -> (~=~) rel t tm tm' -> (~=~) rel t tm' tm'' -> (~=~) rel t tm tm''

  public export
  tmsRelTrans : {0 V : sig.T -> Type} -> {0 rel : (t : _) -> Rel (V t)}
    -> ((t : sig.T) -> Transitive (V t) (rel t))
    -> {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' (transitive 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   = (~=~) {sig, v} eq.relation
    , reflexive  = \t => tmRelRefl eq.reflexive
    , symmetric  = \t, _, _ => tmRelSym (\t => MkSymmetric $ eq.symmetric t _ _)
    , transitive = \t, _, _, _ => tmRelTrans (\t => MkTransitive $ eq.transitive t _ _ _)
    }

public export
Free : (0 V : sig.T -> Type) -> RawAlgebra sig
Free v = MakeRawAlgebra (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 : _) -> 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} -> (~=~) 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 : _) -> 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} -> 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} -> (~=~) 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}
  -> 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
bindIsHomo : (a : Algebra sig) -> (env : v ~> a.setoid)
  -> IsHomomorphism (FreeAlgebra v) a (\t => bindTerm {a = a.raw} env.H)
bindIsHomo a env = MkIsHomomorphism
  (\t, eq => bindTermCong env eq)
  (\op, tms =>
      a.algebra.semCong op $
      map (\t => reflect $ index a.setoid t) $
      equalImpliesPwEq $
      bindTermsIsMap {a = a.raw} env.H tms)

public export
bindHomo : {a : Algebra sig} -> (env : v ~> a.setoid) -> Homomorphism (FreeAlgebra v) a
bindHomo env = MkHomomorphism _ (bindIsHomo _ env)

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

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

bindUnique' f g cong (Done i)     = cong i
bindUnique' f g cong (Call op ts) = CalcWith (index a.setoid _) $
  |~ f.func _ (Call op ts)
  ~~ a.raw.sem op (map f.func ts) ...( f.homo.semHomo op ts )
  ~~ a.raw.sem op (map g.func ts) ...( a.algebra.semCong op $ bindsUnique' f g cong ts )
  ~~ g.func _ (Call op ts)        ..<( g.homo.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 : Homomorphism (FreeAlgebra v) a)
  -> (cong : {t : sig.T} -> (i : v.U t) -> a.relation t (f.func t (Done i)) (env.H t i))
  -> {t : sig.T} -> (tm : Term sig v.U t)
  -> a.relation t (f.func 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 : Homomorphism (FreeAlgebra v) a)
  -> (cong : {t : sig.T} -> (i : v.U t) -> a.relation t (f.func t (Done i)) (env.H t i))
  -> {ts : List sig.T} -> (tms : Term sig v.U ^ ts)
  -> Pointwise a.relation (map f.func tms) (bindTerms {a = a.raw} env.H tms)
bindsUnique env f cong ts = CalcWith (pwSetoid a.setoid _) $
  |~ map f.func 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 )