module Data.Term

import public Data.DPair
import public Data.Vect

import Data.Fin.Extra
import Data.Vect.Properties.Map

import Syntax.PreorderReasoning

%prefix_record_projections off

-- Definition ------------------------------------------------------------------

public export
record Signature where
  constructor MkSignature
  Operator : Nat -> Type

%name Signature sig

public export
(.Op) : Signature -> Type
sig .Op = Exists sig.Operator

public export
interface DecOp (0 sig : Signature) where
  decOp : (op, op' : sig.Op) -> Dec (op = op')

public export
data Term : Signature -> Nat -> Type where
  Var : Fin n -> Term sig n
  Op : sig.Operator k -> Vect k (Term sig n) -> Term sig n

%name Term t, u, v

export
Uninhabited (Op op ts = Var x) where uninhabited prf impossible

-- Substitution ----------------------------------------------------------------

export
pure : (Fin k -> Fin n) -> Fin k -> Term sig n
pure r = Var . r

public export
(<$>) : (Fin k -> Term sig n) -> Term sig k -> Term sig n
f <$> Var i = f i
f <$> Op op ts = Op op (map (\t => f <$> assert_smaller ts t) ts)

-- Extensional Equality --------------------------------------------------------

public export
0 (.=.) : Rel (Fin k -> Term sig n)
f .=. g = (i : Fin k) -> f i = g i

export
subExtensional : {f, g : Fin k -> Term sig n} -> f .=. g -> (t : Term sig k) -> f <$> t = g <$> t
subExtensional prf (Var i) = prf i
subExtensional prf (Op op ts) =
  cong (Op op) $
  mapExtensional (f <$>) (g <$>) (\t => subExtensional prf (assert_smaller ts t)) ts

-- Substitution is Monadic -----------------------------------------------------

public export
(.) : (Fin k -> Term sig n) -> (Fin j -> Term sig k) -> Fin j -> Term sig n
f . g = (f <$>) . g

-- Properties

export
subUnit : (t : Term sig n) -> Var <$> t = t
subUnit (Var i) = Refl
subUnit (Op op ts) = cong (Op op) $ Calc $
  |~ map (Var <$>) ts
  ~~ map id ts        ...(mapExtensional (Var <$>) id (\t => subUnit (assert_smaller ts t)) ts)
  ~~ ts               ...(mapId ts)

export
subAssoc :
  (f : Fin k -> Term sig n) ->
  (g : Fin j -> Term sig k) ->
  (t : Term sig j) ->
  (f . g) <$> t = f <$> g <$> t
subAssoc f g (Var i) = Refl
subAssoc f g (Op op ts) = cong (Op op) $ Calc $
  |~ map ((f . g) <$>) ts
  ~~ map ((f <$>) . (g <$>)) ts   ...(mapExtensional ((f . g) <$>) ((f <$>) . (g <$>)) (\t => subAssoc f g (assert_smaller ts t)) ts)
  ~~ map (f <$>) (map (g <$>) ts) ..<(mapFusion (f <$>) (g <$>) ts)

export
subComp :
  (0 f : Fin k -> Term sig n) ->
  (0 r : Fin j -> Fin k) ->
  f . pure r .=. f . r
subComp f r i = Refl

export
subAssoc' :
  (f : Fin k -> Term sig n) ->
  (g : Fin j -> Fin k) ->
  (t : Term sig j) ->
  (f . g) <$> t = f <$> pure g <$> t
subAssoc' f g = subAssoc f (Var . g)

export
compCong2 :
  {0 f1, f2 : Fin k -> Term sig n} ->
  {0 g1, g2 : Fin j -> Term sig k} ->
  f1 .=. f2 ->
  g1 .=. g2 ->
  f1 . g1 .=. f2 . g2
compCong2 prf1 prf2 i = Calc $
  |~ f1 <$> g1 i
  ~~ f1 <$> g2 i ...(cong (f1 <$>) $ prf2 i)
  ~~ f2 <$> g2 i ...(subExtensional prf1 (g2 i))

compCongL :
  {0 f1, f2 : Fin k -> Term sig n} ->
  f1 .=. f2 ->
  (0 g : Fin j -> Term sig k) ->
  f1 . g .=. f2 . g
compCongL prf g = compCong2 prf (\_ => Refl)

compCongR :
  (0 f : Fin k -> Term sig n) ->
  {0 g1, g2 : Fin j -> Term sig k} ->
  g1 .=. g2 ->
  f . g1 .=. f . g2
compCongR f prf = compCong2 (\_ => Refl) prf

-- Substitution is a Semigroupoid ----------------------------------------------

export
(++) :
  {j : Nat} ->
  (Fin j -> Term sig n) ->
  (Fin k -> Term sig n) ->
  Fin (j + k) -> Term sig n
f ++ g = either f g . splitSum

export
appendWeakenN :
  {j : Nat} ->
  (0 f : Fin j -> Term sig n) ->
  (0 g : Fin k -> Term sig n) ->
  (f ++ g) . weakenN k .=. f
appendWeakenN f g i = cong (either f g) $ splitSumOfWeakenN i

export
appendShift :
  {j : Nat} ->
  (0 f : Fin j -> Term sig n) ->
  (0 g : Fin k -> Term sig n) ->
  (f ++ g) . shift j .=. g
appendShift f g i = cong (either f g) $ splitSumOfShift i

export
appendCong2 :
  {j : Nat} ->
  {0 f : Fin j -> Term sig n} ->
  {0 g : Fin k -> Term sig n} ->
  {0 h : Fin (j + k) -> Term sig n} ->
  h . weakenN k .=. f ->
  h . shift j .=. g ->
  h .=. f ++ g
appendCong2 prf1 prf2 x with (splitSum x) proof prf
  _ | Left i = Calc $
    |~ h x
    ~~ h (indexSum (splitSum x)) ..<(cong h (indexOfSplitSumInverse x))
    ~~ h (weakenN k i)           ...(cong (h . indexSum) prf)
    ~~ f i                       ...(prf1 i)
  _ | Right i = Calc $
    |~ h x
    ~~ h (indexSum (splitSum x)) ..<(cong h (indexOfSplitSumInverse x))
    ~~ h (shift j i)             ...(cong (h . indexSum) prf)
    ~~ g i                       ...(prf2 i)

-- Injectivity -----------------------------------------------------------------

export
opInjectiveLen : (prf : Op {k} op ts = Op {k = n} op' us) -> k = n
opInjectiveLen Refl = Refl

export
opInjectiveOp :
  (prf : Op {sig, k} op ts = Op {sig, k = n} op' us) ->
  the sig.Op (Evidence k op) = Evidence n op'
opInjectiveOp Refl = Refl

export
opInjectiveTs' : {0 ts, us : Vect k (Term sig n)} -> (prf : Op op ts = Op op us) -> ts = us
opInjectiveTs' Refl = Refl

export
opInjectiveTs :
  {0 ts : Vect k (Term sig n)} ->
  {0 us : Vect j (Term sig n)} ->
  (prf : Op op ts = Op op' us) -> ts = (rewrite opInjectiveLen prf in us)
opInjectiveTs Refl = Refl