module Soat.FirstOrder.Algebra.Coproduct

import Data.Setoid.Indexed
import Data.Setoid.Either
import Data.Setoid.Product

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) -> (t : sig.T)
  -> Rel (CoproductSet sig x.raw.U y.raw.U t)
  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 : (xRefl : (t : sig.T) -> (a : x.raw.U t) -> x.relation t a a)
    -> (yRefl : (t : sig.T) -> (a : y.raw.U t) -> y.relation t a a)
    -> {t : _} -> (tm : _) -> (~=~) x y t tm tm
  coprodsRelRefl : (xRefl : (t : sig.T) -> (a : x.raw.U t) -> x.relation t a a)
    -> (yRefl : (t : sig.T) -> (a : y.raw.U t) -> y.relation t a a)
    -> {ts : _} -> (tms : _ ^ ts) -> Pointwise ((~=~) x y) tms tms

  coprodRelRefl x y (Done (Left i))  = DoneL $ x _ i
  coprodRelRefl x y (Done (Right i)) = DoneR $ y _ i
  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

  tmRelImpliesCoprodRel : Term.Rel.(~=~) (\i => x.relation i `Or` y.relation i) t tm tm'
    -> (~=~) x y t tm tm'
  tmRelsImpliesCoprodRels : {0 tms, tms' : _ ^ ts}
    -> Pointwise (Term.Rel.(~=~) (\i => x.relation i `Or` 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

namespace Rel
  public export
  CoproductSet : (0 sig : Signature) -> (x, y : Algebra sig) -> IndexedSetoid sig.T
  CoproductSet sig x y = MkIndexedSetoid
    { U           = CoproductSet sig x.raw.U y.raw.U
    , equivalence = MkIndexedEquivalence
      { relation   = (~=~) x.rawSetoid y.rawSetoid
      , reflexive  = \_ =>
        coprodRelRefl
          x.algebra.equivalence.reflexive
          y.algebra.equivalence.reflexive
      , symmetric  = \_, _, _ => Sym
      , transitive = \_, _, _, _ => Trans {x = x.rawSetoid, y = y.rawSetoid}
      }
    }

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

public export
CoproductAlgebra : (x, y : Algebra sig) -> Algebra sig
CoproductAlgebra x y = MkAlgebra
  { U   = CoproductSet sig x y
  , sem = \op => MkSetoidHomomorphism (Call op) (\_, _ => Call op)

  }

public export
fromTerm : (0 x, y : Algebra sig)
  -> Term sig (bundle $ \i => Either (index x.setoid i) (index y.setoid i))
  ~> CoproductSet sig x y
fromTerm x y = MkIndexedSetoidHomomorphism (\_ => id) (\_, _, _ => tmRelImpliesCoprodRel)

public export
fromFree : (x, y : Algebra sig)
  -> FreeAlgebra (bundle (\t => Either (index x.setoid t) (index y.setoid t)))
  ~> CoproductAlgebra x y
fromFree x y = MkHomomorphism
  { func    = fromTerm x y
  , semHomo = \op, tms =>
    Call op $
    reflect (index (Product (CoproductAlgebra x y).setoid) _) $
    sym $
    mapId tms
  }

public export
injectLHomo : {x, y : Algebra sig} -> x ~> CoproductAlgebra x y
injectLHomo = MkHomomorphism
  { func    = bundle (\t =>
    index
      {x = bundle (\t => Either (index x.setoid t) (index y.setoid t))}
      (fromTerm x y . injectFunc) t .
    Left)
  , semHomo = InjectL
  }

public export
injectRHomo : {x, y : Algebra sig} -> y ~> CoproductAlgebra x y
injectRHomo = MkHomomorphism
  { func    = bundle (\t =>
    index
      {x = bundle (\t => Either (index x.setoid t) (index y.setoid t))}
      (fromTerm x y . injectFunc) t .
    Right)
  , semHomo = InjectR
  }

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

public export
coproducts : {z : RawAlgebra sig}
  -> (f : (t : sig.T) -> x t -> z.U t) -> (g : (t : sig.T) -> y t -> z.U t)
  -> (ts : List sig.T) -> 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.H g.func.H 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.H g.func.H ts)) tms tms'

coproductCong f g (DoneL eq)       = f.func.homomorphic _ _ _ eq
coproductCong f g (DoneR eq)       = g.func.homomorphic _ _ _ eq
coproductCong f g (Call op eqs)    = z.algebra.semCong op $ coproductsCong f g eqs
coproductCong f g (InjectL op ts)  = CalcWith (index z.setoid _) $
  |~ f.func.H _ (x.raw.sem op ts)
  ~~ z.raw.sem op (map f.func.H ts)                                                ...(f.semHomo op ts)
  ~~ z.raw.sem op (map (coproduct f.func.H g.func.H) (map (\_ => Done . Left) ts)) .=.(cong (z.raw.sem op) $ mapComp ts)
  ~~ z.raw.sem op (coproducts f.func.H g.func.H _ (map (\_ => Done . Left) ts))    .=<(cong (z.raw.sem op) $ bindTermsIsMap {a = z.raw} _ _)
coproductCong f g (InjectR op ts)  = CalcWith (index z.setoid _) $
  |~ g.func.H _ (y.raw.sem op ts)
  ~~ z.raw.sem op (map g.func.H ts)                                                 ...(g.semHomo op ts)
  ~~ z.raw.sem op (map (coproduct f.func.H g.func.H) (map (\_ => Done . Right) ts)) .=.(cong (z.raw.sem op) $ mapComp ts)
  ~~ z.raw.sem op (coproducts f.func.H g.func.H _ (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
coproductFunc : {x, y, z : Algebra sig} -> x ~> z -> y ~> z
  -> CoproductSet sig x y ~> z.setoid
coproductFunc f g = MkIndexedSetoidHomomorphism
  { H           = coproduct {z = z.raw} f.func.H g.func.H
  , homomorphic = \_, _, _ => coproductCong f g
  }

public export
coproductHomo : {x, y, z : Algebra sig} -> x ~> z -> y ~> z
  -> CoproductAlgebra x y ~> z
coproductHomo f g = MkHomomorphism
  { func = coproductFunc f g
  , semHomo = \op, tms =>
    reflect (index z.setoid _) $
    cong (z.raw.sem op) $
    bindTermsIsMap {a = z.raw} _ 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.H t (Done (Left i))) (g.func.H t (Done (Left i))))
  -> (congR : {t : sig.T} -> (i : y.raw.U t)
        -> z.relation t (f.func.H t (Done (Right i))) (g.func.H t (Done (Right i))))
  -> {t : sig.T} -> (tm : _)
  -> z.relation t (f.func.H t tm) (g.func.H t tm)
coproductUnique' f g congL congR tm = bindUnique'
  { v = bundle (\t => Either (index x.setoid t) (index y.setoid t))
  , f = f . fromFree x y
  , g = g . fromFree 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.H t (Done (Left i))) (f.func.H t i))
  -> (congR : {t : sig.T} -> (i : y.raw.U t)
        -> z.relation t (h.func.H t (Done (Right i))) (g.func.H t i))
  -> {t : sig.T} -> (tm : _)
  -> z.relation t (h.func.H t tm) (coproduct {z = z.raw} f.func.H g.func.H t tm)
coproductUnique f g h = coproductUnique' h (coproductHomo f g)