module Soat.SecondOrder.Algebra

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.SecondOrder.Signature

infix 5 ~>

public export
extend : (U : a -> List a -> Type) -> (ctx : List a) -> Pair (List a) a -> Type
extend x ctx ty = x (snd ty) (fst ty ++ ctx)

public export
extendRel : {U : a -> List a -> Type} -> (rel : (ty : _) -> Rel (uncurry U ty))
  -> (ctx : List a) -> (ty : _) -> Rel (extend U ctx ty)
extendRel rel ctx ty = rel (snd ty, fst ty ++ ctx)

public export 0
algebraOver : (sig : Signature) -> (U : sig.T -> List sig.T -> Type) -> Type
algebraOver sig x = (ctx : List sig.T) -> {t : sig.T} -> (op : Op sig t)
  -> extend x ctx ^ op.arity -> x t ctx

public export
record RawAlgebra (0 sig : Signature) where
  constructor MkRawAlgebra
  0 U    : (t : sig.T) -> (ctx : List sig.T) -> Type
  rename : (t : sig.T) -> forall ctx, ctx' . (f : ctx `Sublist` ctx') -> U t ctx -> U t ctx'
  sem    : sig `algebraOver` U
  var    : forall t, ctx . (i : Elem t ctx) -> U t ctx
  subst  : (t : sig.T) -> (ctx : List sig.T)
    -> forall ctx' . U t ctx' -> (\t => U t ctx) ^ ctx' -> U t ctx

public export
(.extendSubst) : (a : RawAlgebra sig) -> (ctx : List sig.T) -> {ctx' : List sig.T}
  -> (tms : (\t => a.U t ctx) ^ ctx') -> (ty : _) -> extend a.U ctx' ty -> extend a.U ctx ty
a .extendSubst ctx tms ty tm = a.subst
  (snd ty)
  (fst ty ++ ctx)
  tm
  (tabulate {is = fst ty} (a.var . Sublist.elemJoinL {ys = ctx}) ++
   map (\t => a.rename t ([] {ys = fst ty} ++ Relation.reflexive {x = ctx})) tms)

public export
record IsAlgebra (0 sig : Signature) (0 a : RawAlgebra sig) where
  constructor MkIsAlgebra
  equivalence : IndexedEquivalence _ (uncurry a.U)
  -- Congruences
  renameCong  : (t : sig.T) -> {ctx, ctx' : _} -> (f : ctx `Sublist` ctx')
    -> {tm, tm' : a.U t ctx} -> equivalence.relation (t, ctx) tm tm'
    -> equivalence.relation (t, ctx') (a.rename t f tm) (a.rename t f tm')
  semCong     : (ctx : List sig.T) -> {t : sig.T} -> (op : Op sig t)
    -> {tms, tms' : extend a.U ctx ^ op.arity}
    -> Pointwise (extendRel {U = a.U} equivalence.relation ctx) tms tms'
    -> equivalence.relation (t, ctx) (a.sem ctx op tms) (a.sem ctx op tms')
  substCong   : (t : sig.T) -> (ctx : List sig.T)
    -> {ctx' : _} -> {tm, tm' : a.U t ctx'} -> equivalence.relation (t, ctx') tm tm'
    -> {tms, tms' : (\t => a.U t ctx) ^ ctx'}
    -> Pointwise (\t => equivalence.relation (t, ctx)) tms tms'
    -> equivalence.relation (t, ctx) (a.subst t ctx tm tms) (a.subst t ctx tm' tms')
  -- rename is functorial
  renameId    : (t : sig.T) -> (ctx : List sig.T) -> (tm : a.U t ctx)
    -> equivalence.relation (t, ctx) (a.rename t {ctx = ctx} Relation.reflexive tm) tm
  renameComp  : (t : sig.T) -> {ctx, ctx', ctx'' : _}
    -> (f : ctx' `Sublist` ctx'') -> (g : ctx `Sublist` ctx')
    -> (tm : a.U t ctx)
    -> equivalence.relation (t, ctx'')
         (a.rename t (transitive g f) tm)
         (a.rename t f $ a.rename t g tm)
  -- sem are natural transformation
  semNat      : {ctx, ctx' : _} -> (f : ctx `Sublist` ctx') -> {t : sig.T} -> (op : Op sig t)
    -> (tms : extend a.U ctx ^ op.arity)
    -> equivalence.relation (t, ctx')
         (a.rename t f $ a.sem ctx op tms)
         (a.sem ctx' op $
          map (\ty => a.rename (snd ty) (Relation.reflexive {x = fst ty} ++ f)) $
          tms)
  -- var is natural transformation
  varNat      : {t : _} -> {ctx, ctx' : _} -> (f : ctx `Sublist` ctx') -> (i : Elem t ctx)
    -> equivalence.relation (t, ctx') (a.rename t f $ a.var i) (a.var $ curry f i)
  -- subst is natural transformation
  substNat    : (t : sig.T) -> {ctx, ctx' : _} -> (f : ctx `Sublist` ctx')
    -> {ctx'' : _} -> (tm : a.U t ctx'') -> (tms : (\t => a.U t ctx) ^ ctx'')
    -> equivalence.relation (t, ctx')
         (a.rename t f $ a.subst t ctx tm tms)
         (a.subst t ctx' tm $ map (\t => a.rename t f) tms)
  -- subst is extranatural transformation
  substExnat  : (t : sig.T) -> (ctx : List sig.T)
    -> {ctx', ctx'' : _} -> (f : ctx' `Sublist` ctx'')
    -> (tm : a.U t ctx') -> (tms : (\t => a.U t ctx) ^ ctx'')
    -> equivalence.relation (t, ctx)
         (a.subst t ctx (a.rename t f tm) tms)
         (a.subst t ctx tm (shuffle f tms))
  -- var, subst is a monoid
  substComm   : (t : sig.T) -> (ctx : List sig.T)
    -> {ctx', ctx'' : _} -> (tm : a.U t ctx'')
    -> (tms : (\t => a.U t ctx') ^ ctx'') -> (tms' : (\t => a.U t ctx) ^ ctx')
    -> equivalence.relation (t, ctx)
         (a.subst t ctx (a.subst t ctx' tm tms) tms')
         (a.subst t ctx tm $ map (\t, tm => a.subst t ctx tm tms') tms)
  substVarL   : {t : _} -> (ctx : List sig.T) -> {ctx' : _} -> (i : Elem t ctx')
    -> (tms : (\t => a.U t ctx) ^ ctx')
    -> equivalence.relation (t, ctx) (a.subst t ctx (a.var i) tms) (index tms i)
  substVarR   : (t : sig.T) -> (ctx : List sig.T) -> (tm : a.U t ctx)
    -> equivalence.relation (t, ctx) (a.subst t ctx {ctx' = ctx} tm $ tabulate a.var) tm
  -- sem, var and subst are compatible
  substCompat : (ctx : List sig.T) -> {t : sig.T} -> (op : Op sig t)
    -> {ctx' : _} -> (tms : extend a.U ctx' ^ op.arity) -> (tms' : (\t => a.U t ctx) ^ ctx')
    -> equivalence.relation (t, ctx)
         (a.subst t ctx (a.sem ctx' op tms) tms')
         (a.sem ctx op $
          map (\ty, tm =>
            a.subst (snd ty) (fst ty ++ ctx) tm
              (tabulate {is = fst ty} (a.var . Sublist.elemJoinL {ys = ctx}) ++
               map (\t => a.rename t ([] {ys = fst ty} ++ Relation.reflexive {x = ctx})) tms')) $
          tms)

public export
record Algebra (0 sig : Signature) where
  constructor MkAlgebra
  raw        : RawAlgebra sig
  algebra    : IsAlgebra sig raw

public export 0
(.relation) : (0 a : Algebra sig) -> (ty : _) -> Rel (uncurry a.raw.U ty)
(.relation) a = a.algebra.equivalence.relation

public export
(.setoid) : Algebra sig -> IndexedSetoid (Pair sig.T (List sig.T))
(.setoid) a = MkIndexedSetoid (uncurry a.raw.U) a.algebra.equivalence

public export
(.setoidAt) : Algebra sig -> (ctx : List sig.T) -> IndexedSetoid sig.T
(.setoidAt) a ctx = reindex (flip MkPair ctx) a.setoid

public export
(.renameFunc) : (a : Algebra sig) -> {ctx, ctx' : _} -> (f : ctx `Sublist` ctx')
  -> a.setoidAt ctx ~> a.setoidAt ctx'
(.renameFunc) a f = MkIndexedSetoidHomomorphism
  { H           = \t => a.raw.rename t f
  , homomorphic = \t, _, _ => a.algebra.renameCong t f
  }

public export
(.semFunc) : (a : Algebra sig) -> (ctx : _) -> {t : _} -> (op : Op sig t)
  -> index (Product (reindex (\ty => (snd ty, fst ty ++ ctx)) a.setoid)) op.arity
       ~> index a.setoid (t, ctx)
(.semFunc) a ctx op = MkSetoidHomomorphism
  { H           = a.raw.sem ctx op
  , homomorphic = \_, _ => a.algebra.semCong ctx op
  }

public export
(.varFunc) : (a : Algebra sig) -> irrelevantCast (uncurry Elem) ~> a.setoid
(.varFunc) a = mate (\(_, _) => a.raw.var)

public export
(.substFunc) : (a : Algebra sig) -> (t : _) -> (ctx, ctx' : _)
  -> Pair (index a.setoid (t, ctx')) (index (Product (a.setoidAt ctx)) ctx')
       ~> index a.setoid (t, ctx)
(.substFunc) a t ctx ctx' = MkSetoidHomomorphism
  { H           = \x => a.raw.subst t ctx (fst x) (snd x)
  , homomorphic = \_, _, eq => a.algebra.substCong t ctx (fst eq) (snd eq)
  }

public export
record (~>) {0 sig : Signature} (a, b : Algebra sig)
  where
  constructor MkHomomorphism
  func       : a.setoid ~> b.setoid
  renameHomo : (t : sig.T) -> {ctx, ctx' : _} -> (g : ctx `Sublist` ctx') -> (tm : a.raw.U t ctx)
    -> b.relation (t, ctx')
         (func.H (t, ctx') $ a.raw.rename t g tm)
         (b.raw.rename t g $ func.H (t, ctx) tm)
  semHomo    : (ctx : List sig.T) -> {t : sig.T} -> (op : Op sig t)
    -> (tms : extend a.raw.U ctx ^ op.arity)
    -> b.relation (t, ctx)
         (func.H (t, ctx) $ a.raw.sem ctx op tms)
         (b.raw.sem ctx op $ map (\ty => func.H (snd ty, fst ty ++ ctx)) tms)
  varHomo    : {t : _} -> {ctx : _} -> (i : Elem t ctx)
    -> b.relation (t, ctx) (func.H (t, ctx) $ a.raw.var i) (b.raw.var i)
  substHomo  : (t : sig.T) -> (ctx : List sig.T) -> {ctx' : _} -> (tm : a.raw.U t ctx')
    -> (tms : (\t => a.raw.U t ctx) ^ ctx')
    -> b.relation (t, ctx)
         (func.H (t, ctx) $ a.raw.subst t ctx tm tms)
         (b.raw.subst t ctx (func.H (t, ctx') tm) $ map (\t => func.H (t, ctx)) tms)