module Soat.FirstOrder.Algebra.Coproduct

import Data.Morphism.Indexed
import Data.Setoid.Indexed
import Data.Setoid.Either

import Soat.FirstOrder.Algebra
import Soat.FirstOrder.Signature
import Soat.FirstOrder.Term

import Syntax.PreorderReasoning.Setoid

%default total

public export
CoproductSet : (0 sig : Signature) -> (0 x, y : sig.T -> Type) -> sig.T -> Type
CoproductSet sig x y = Term sig (\i => Either (x i) (y i))

public export
data (~=~) : (0 x, y : RawSetoidAlgebra sig) -> IRel (CoproductSet sig x.raw.U y.raw.U) where
  DoneL    : {0 x, y : RawSetoidAlgebra sig}
    -> {t : sig.T} -> {i, j : x.raw.U t} -> x.relation t i j
    -> (~=~) x y t (Done $ Left i) (Done $ Left j)
  DoneR    : {0 x, y : RawSetoidAlgebra sig}
    -> {t : sig.T} -> {i, j : y.raw.U t} -> y.relation t i j
    -> (~=~) x y t (Done $ Right i) (Done $ Right j)
  Call     : {t : sig.T} -> (op : Op sig t) -> {xs, ys : _} -> Pointwise ((~=~) x y) xs ys
    -> (~=~) x y t (Call op xs) (Call op ys)
  InjectL  : {t : sig.T} -> (op : Op sig t) -> (xs : _)
    -> (~=~) x y t (Done $ Left $ x.raw.sem op xs) (Call op (map (\_ => Done . Left) xs))
  InjectR  : {t : sig.T} -> (op : Op sig t) -> (xs : _)
    -> (~=~) x y t (Done $ Right $ y.raw.sem op xs) (Call op (map (\_ => Done . Right) xs))
  Sym      : {0 x, y : RawSetoidAlgebra sig} -> {t : sig.T} -> {tm, tm' : _}
    -> (~=~) x y t tm tm' -> (~=~) x y t tm' tm
  Trans    : {0 x, y : RawSetoidAlgebra sig} -> {t : sig.T} -> {tm, tm', tm'' : _}
    -> (~=~) x y t tm tm' -> (~=~) x y t tm' tm'' -> (~=~) x y t tm tm''

parameters {0 x, y : RawSetoidAlgebra sig}
  coprodRelRefl : IReflexive x.raw.U x.relation -> IReflexive y.raw.U y.relation
    -> {t : _} -> (tm : _) -> (~=~) x y t tm tm
  coprodsRelRefl : IReflexive x.raw.U x.relation -> IReflexive y.raw.U y.relation
    -> {ts : _} -> (tms : _ ^ ts) -> Pointwise ((~=~) x y) tms tms

  coprodRelRefl x y (Done (Left i))  = DoneL $ reflexive @{x _}
  coprodRelRefl x y (Done (Right i)) = DoneR $ reflexive @{y _}
  coprodRelRefl x y (Call op ts)     = Call op $ coprodsRelRefl x y ts

  coprodsRelRefl x y []        = []
  coprodsRelRefl x y (t :: ts) = coprodRelRefl x y t :: coprodsRelRefl x y ts

  coprodsRelReflexive : IReflexive x.raw.U x.relation -> IReflexive y.raw.U y.relation
    -> {ts : _} -> {tms, tms' : _ ^ ts} -> tms = tms' -> Pointwise ((~=~) x y) tms tms'
  coprodsRelReflexive x y Refl = coprodsRelRefl x y _

  tmRelImpliesCoprodRel : Term.Rel.(~=~) (\i => Either (x.relation i) (y.relation i)) t tm tm'
    -> (~=~) x y t tm tm'
  tmRelsImpliesCoprodRels : {0 tms, tms' : _ ^ ts}
    -> Pointwise (Term.Rel.(~=~) (\i => Either (x.relation i) (y.relation i))) tms tms'
    -> Pointwise ((~=~) x y) tms tms'

  tmRelImpliesCoprodRel (Done' (Left eq))  = DoneL eq
  tmRelImpliesCoprodRel (Done' (Right eq)) = DoneR eq
  tmRelImpliesCoprodRel (Call' op eqs)     = Call op $ tmRelsImpliesCoprodRels eqs

  tmRelsImpliesCoprodRels []          = []
  tmRelsImpliesCoprodRels (eq :: eqs) = tmRelImpliesCoprodRel eq :: tmRelsImpliesCoprodRels eqs

public export
Coproduct : (x, y : RawAlgebra sig) -> RawAlgebra sig
Coproduct x y = MkRawAlgebra
  { U   = CoproductSet sig x.U y.U
  , sem = Call
  }

public export
CoproductIsAlgebra : IsAlgebra sig x -> IsAlgebra sig y
  -> IsAlgebra sig (Coproduct x y)
CoproductIsAlgebra xIsAlgebra yIsAlgebra = MkIsAlgebra
  { relation =
    (~=~)
      (MkRawSetoidAlgebra x xIsAlgebra.relation)
      (MkRawSetoidAlgebra y yIsAlgebra.relation)
  , equivalence = \t => MkEquivalence
    { refl  = MkReflexive $ coprodRelRefl
        (\t => (xIsAlgebra.equivalence t).refl)
        (\t => (yIsAlgebra.equivalence t).refl)
        _
    , sym   = MkSymmetric $ Sym
    , trans = MkTransitive $ Trans
    }
  , semCong = Call
  }

public export
CoproductAlgebra : (x, y : Algebra sig) -> Algebra sig
CoproductAlgebra x y = MkAlgebra (Coproduct x.raw y.raw) $ CoproductIsAlgebra x.algebra y.algebra

public export
injectLHomo : x ~> CoproductAlgebra x y
injectLHomo = MkHomomorphism
  { func    = \_ => Done . Left
  , cong    = \_ => DoneL
  , semHomo = InjectL
  }

public export
injectRHomo : y ~> CoproductAlgebra x y
injectRHomo = MkHomomorphism
  { func    = \_ => Done . Right
  , cong    = \_ => DoneR
  , semHomo = InjectR
  }

public export
coproduct : {z : RawAlgebra sig} -> IFunc x z.U -> IFunc y z.U -> IFunc (CoproductSet sig x y) z.U
coproduct f g _ = bindTerm (\i => either (f i) (g i))

public export
coproducts : {z : RawAlgebra sig} -> IFunc x z.U -> IFunc y z.U -> (ts : _)
  -> CoproductSet sig x y ^ ts -> z.U ^ ts
coproducts f g _ = bindTerms (\i => either (f i) (g i))

public export
coproductCong : {x, y, z : Algebra sig} -> (f : x ~> z) -> (g : y ~> z)
  -> (t : sig.T) -> {tm, tm' : _} -> (~=~) x.rawSetoid y.rawSetoid t tm tm'
  -> (z.relation t `on` coproduct {z = z.raw} f.func g.func t) tm tm'

public export
coproductsCong : {x, y, z : Algebra sig} -> (f : x ~> z) -> (g : y ~> z)
  -> (ts : List sig.T) -> {tms, tms' : _ ^ ts}
  -> Pointwise ((~=~) x.rawSetoid y.rawSetoid) tms tms'
  -> (Pointwise z.relation `on` (coproducts {z = z.raw} f.func g.func ts)) tms tms'

coproductCong f g _ (DoneL eq)       = f.cong _ eq
coproductCong f g _ (DoneR eq)       = g.cong _ eq
coproductCong f g _ (Call op eqs)    =
  z.algebra.semCong op $
  coproductsCong f g _ eqs
coproductCong f g _ (InjectL op ts)  = CalcWith (z.setoid.index _) $
  |~ f.func _ (x.raw.sem op ts)
  ~~ z.raw.sem op (map f.func ts)                                              ...(f.semHomo op ts)
  ~~ z.raw.sem op (map (coproduct f.func g.func) (map (\_ => Done . Left) ts)) .=.(cong (z.raw.sem op) $ mapComp ts)
  ~~ z.raw.sem op (coproducts f.func g.func _ (map (\_ => Done . Left) ts))    .=<(cong (z.raw.sem op) $ bindTermsIsMap {a = z.raw} _ _)
coproductCong f g _ (InjectR op ts)  = CalcWith (z.setoid.index _) $
  |~ g.func _ (y.raw.sem op ts)
  ~~ z.raw.sem op (map g.func ts)                                               ...(g.semHomo op ts)
  ~~ z.raw.sem op (map (coproduct f.func g.func) (map (\_ => Done . Right) ts)) .=.(cong (z.raw.sem op) $ mapComp ts)
  ~~ z.raw.sem op (coproducts f.func g.func _ (map (\_ => Done . Right) ts))    .=<(cong (z.raw.sem op) $ bindTermsIsMap {a = z.raw} _ _)
coproductCong f g _ (Sym eq)         = (z.algebra.equivalence _).symmetric $ coproductCong f g _ eq
coproductCong f g _ (Trans eq eq')   = (z.algebra.equivalence _).transitive
  (coproductCong f g _ eq)
  (coproductCong f g _ eq')

coproductsCong f g _ []          = []
coproductsCong f g _ (eq :: eqs) =
  coproductCong f g _ eq :: coproductsCong f g _ eqs

public export
coproductHomo : {x, y, z : Algebra sig} -> x ~> z -> y ~> z
  -> CoproductAlgebra x y ~> z
coproductHomo f g = MkHomomorphism
  { func = coproduct {z = z.raw} f.func g.func
  , cong = coproductCong f g
  , semHomo = \op, tms =>
    (z.algebra.equivalence _).equalImpliesEq $
    cong (z.raw.sem op) $
    bindTermsIsMap {a = z.raw} _ tms
  }

public export
termToCoproduct : (x, y : Algebra sig)
  -> FreeAlgebra (fromIndexed (\i => EitherSetoid (x.setoid.index i) (y.setoid.index i))) ~>
     CoproductAlgebra x y
termToCoproduct x y = MkHomomorphism
  { func    = \_ => id
  , cong    = \t => tmRelImpliesCoprodRel
  , semHomo = \op, tms =>
    Call op $
    coprodsRelReflexive
      (\i => (x.algebra.equivalence i).refl)
      (\i => (y.algebra.equivalence i).refl) $
    sym $
    mapId tms
  }

public export
coproductUnique' : {x, y, z : Algebra sig}
  -> (f : CoproductAlgebra x y ~> z)
  -> (g : CoproductAlgebra x y ~> z)
  -> (congL : {t : sig.T} -> (i : x.raw.U t)
        -> z.relation t (f.func t (Done (Left i))) (g.func t (Done (Left i))))
  -> (congR : {t : sig.T} -> (i : y.raw.U t)
        -> z.relation t (f.func t (Done (Right i))) (g.func t (Done (Right i))))
  -> {t : sig.T} -> (tm : _)
  -> z.relation t (f.func t tm) (g.func t tm)
coproductUnique' f g congL congR tm = bindUnique'
  { v = fromIndexed (\i => EitherSetoid (x.setoid.index i) (y.setoid.index i))
  , f = compHomo f $ termToCoproduct x y
  , g = compHomo g $ termToCoproduct x y
  , cong = \x => case x of
      (Left x) => congL x
      (Right x) => congR x
  , tm = tm
  }

public export
coproductUnique : {x, y, z : Algebra sig} -> (f : x ~> z) -> (g : y ~> z)
  -> (h : CoproductAlgebra x y ~> z)
  -> (congL : {t : sig.T} -> (i : x.raw.U t)
        -> z.relation t (h.func t (Done (Left i))) (f.func t i))
  -> (congR : {t : sig.T} -> (i : y.raw.U t)
        -> z.relation t (h.func t (Done (Right i))) (g.func t i))
  -> {t : sig.T} -> (tm : _)
  -> z.relation t (h.func t tm) (coproduct {z = z.raw} f.func g.func t tm)
coproductUnique f g h = coproductUnique' h (coproductHomo f g)