module Data.Term.Property

import Data.Term

import Data.Vect.Properties.Index
import Data.Vect.Quantifiers
import Data.Vect.Quantifiers.Extra

import Syntax.PreorderReasoning

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

public export
record Property (sig : Signature) (k : Nat) where
  constructor MkProp
  0 Prop : forall n. (Fin k -> Term sig n) -> Type
  cong : forall n. {f, g : Fin k -> Term sig n} -> f .=. g -> Prop f -> Prop g

%name Property p, q

-- Instances -------------------------------------------------------------------

public export
Unifies : Term sig k -> Term sig k -> Property sig k
Unifies t u = MkProp
  { Prop = \f => f <$> t = f <$> u
  , cong = \cong, prf => Calc $
    |~ _ <$> t
    ~~ _ <$> t ..<(subCong cong t)
    ~~ _ <$> u ...(prf)
    ~~ _ <$> u ...(subCong cong u)
  }

public export
All : Vect n (Property sig k) -> Property sig k
All ps = MkProp (\f => All (\p => p.Prop f) ps) (\cong => mapRel (\p => p.cong cong))

-- Equivalence -----------------------------------------------------------------

public export 0
(<=>) : Property sig k -> Property sig k -> Type
p <=> q = forall n. (f : Fin k -> Term sig n) -> p.Prop f <=> q.Prop f

-- Properties

export
unifiesSym : Unifies t u <=> Unifies u t
unifiesSym f = MkEquivalence (\prf => sym prf) (\prf => sym prf)

export
unifiesOp :
  {k : Nat} ->
  {0 op : sig.Operator k} ->
  {0 ts, us : Vect k (Term sig j)} ->
  Unifies (Op op ts) (Op op us) <=> All (tabulate (\i => Unifies (index i ts) (index i us)))
unifiesOp f = MkEquivalence
  { leftToRight = \prf => tabulate (\i => Unifies (index i ts) (index i us)) (leftToRight prf)
  , rightToLeft = \prf => irrelevantEq $ cong (Op op) $ rightToLeft prf
  }
  where
  leftToRight :
    (Unifies (Op op ts) (Op op us)).Prop f ->
    (i : _) ->
    f <$> index i ts = f <$> index i us
  leftToRight prf i =
    Calc $
      |~ f <$> index i ts
      ~~ index i (map (f <$>) ts) ...(sym $ indexNaturality i (f <$>) ts)
      ~~ index i (map (f <$>) us) ...(cong (index i) $ opInjectiveTs' prf)
      ~~ f <$> index i us         ...(indexNaturality i (f <$>) us)

  rightToLeft :
    {k : Nat} ->
    {ts, us : Vect k (Term sig j)} ->
    (All (tabulate (\i => Unifies (index i ts) (index i us)))).Prop f ->
    map (f <$>) ts = map (f <$>) us
  rightToLeft {ts = [], us = []} [] = Refl
  rightToLeft {ts = t :: ts, us = u :: us} (prf :: prfs) = cong2 (::) prf (rightToLeft prfs)

-- Nothing ---------------------------------------------------------------------

public export 0
Nothing : Property sig k -> Type
Nothing p = forall n. (f : Fin k -> Term sig n) -> Not (p.Prop f)

-- Properties

export
nothingEquiv : p <=> q -> Nothing p -> Nothing q
nothingEquiv eq absurd f x = absurd f ((eq f).rightToLeft x)

-- Extensions ------------------------------------------------------------------

public export
Extension : Property sig k -> (Fin k -> Term sig n) -> Property sig n
Extension p f = MkProp
  { Prop = \g => p.Prop (g . f)
  , cong = \prf => p.cong (\i => subCong prf (f i))
  }

-- Properties

export
nothingExtends : Nothing p -> Nothing (Extension p f)
nothingExtends absurd g x = void $ absurd (g . f) x

export
extendUnit : (p : Property sig k) -> p <=> Extension p Var
extendUnit p f = MkEquivalence id id

export
extendAssoc :
  (p : Property sig j) ->
  (f : Fin j -> Term sig k) ->
  (g : Fin k -> Term sig m) ->
  Extension (Extension p f) g <=> Extension p (g . f)
extendAssoc p f g h =
  MkEquivalence
    (p.cong (\i => subAssoc h g (f i)))
    (p.cong (\i => sym $ subAssoc h g (f i)))

export
extendUnify :
  (t, u : Term sig j) ->
  (f : Fin j -> Term sig k) ->
  (g : Fin k -> Term sig m) ->
  Extension (Unifies t u) (g . f) <=> Extension (Unifies (f <$> t) (f <$> u)) g
extendUnify t u f g h =
  MkEquivalence
    (\prf => Calc $
      |~ (h . g) <$> (f <$> t)
      ~~ ((h . g) . f) <$> t   ...(sym $ subAssoc (h . g) f t)
      ~~ (h . (g . f)) <$> t   ...(subCong (\i => subAssoc h g (f i)) t)
      ~~ (h . (g . f)) <$> u   ...(prf)
      ~~ ((h . g) . f) <$> u   ...(sym $ subCong (\i => subAssoc h g (f i)) u)
      ~~ (h . g) <$> (f <$> u) ...(subAssoc (h . g) f u))
    (\prf => Calc $
      |~ (h . (g . f)) <$> t
      ~~ ((h . g) . f) <$> t   ...(sym $ subCong (\i => subAssoc h g (f i)) t)
      ~~ (h . g) <$> (f <$> t) ...(subAssoc (h . g) f t)
      ~~ (h . g) <$> (f <$> u) ...(prf)
      ~~ ((h . g) . f) <$> u   ...(sym $ subAssoc (h . g) f u)
      ~~ (h . (g . f)) <$> u   ...(subCong (\i => subAssoc h g (f i)) u))

-- Ordering --------------------------------------------------------------------

public export
record (<=) (f : Fin k -> Term sig m) (g : Fin k -> Term sig n) where
  constructor MkLte
  sub : Fin n -> Term sig m
  prf : f .=. sub . g

%name Property.(<=) prf

-- Properties

export
lteCong : f .=. f' -> g .=. g' -> f <= g -> f' <= g'
lteCong prf1 prf2 prf3 = MkLte
  { sub = prf3.sub
  , prf = \i => Calc $
    |~ f' i
    ~~ f i               ...(sym $ prf1 i)
    ~~ prf3.sub <$> g i  ...(prf3.prf i)
    ~~ prf3.sub <$> g' i ...(cong (prf3.sub <$>) $ prf2 i)
  }

export
Reflexive (Fin k -> Term sig m) (<=) where
  reflexive = MkLte Var (\i => sym $ subUnit _)

export
transitive : f <= g -> g <= h -> f <= h
transitive prf1 prf2 = MkLte
  { sub = prf1.sub . prf2.sub
  , prf = \i => Calc $
    |~ f i
    ~~ prf1.sub <$> g i              ...(prf1.prf i)
    ~~ prf1.sub <$> prf2.sub <$> h i ...(cong (prf1.sub <$>) $ prf2.prf i)
    ~~ (prf1.sub . prf2.sub) <$> h i ...(sym $ subAssoc prf1.sub prf2.sub (h i))
  }

export
varMax : (f : Fin k -> Term sig m) -> f <= Var
varMax f = MkLte f (\i => Refl)

export
compLte : f <= g -> (h : Fin k -> Term sig m) -> f . h <= g . h
compLte prf h = MkLte
  { sub = prf.sub
  , prf = \i => Calc $
    |~ f <$> h i
    ~~ (prf.sub . g) <$> h i ...(subCong prf.prf (h i))
    ~~ prf.sub <$> g <$> h i ...(subAssoc prf.sub g (h i))
  }

export
lteExtend : 
  {p : Property sig k} ->
  {f : Fin k -> Term sig m} ->
  {g : Fin k -> Term sig n} ->
  (prf : f <= g) -> 
  p.Prop f -> 
  (Extension p g).Prop prf.sub
lteExtend prf x = p.cong prf.prf x

-- Maximal ---------------------------------------------------------------------

public export
Max : Property sig k -> Property sig k
Max p = MkProp
  { Prop = \f => (p.Prop f, forall n. (g : Fin k -> Term sig n) -> p.Prop g -> g <= f)
  , cong = \cong, (x, max) => (p.cong cong x, \g, y => lteCong (\_ => Refl) cong (max g y))
  }

export
maxCong : p <=> q -> Max p <=> Max q
maxCong prf f = MkEquivalence
  { leftToRight = \(x, max) => ((prf f).leftToRight x, \g, y => max g ((prf g).rightToLeft y))
  , rightToLeft = \(x, max) => ((prf f).rightToLeft x, \g, y => max g ((prf g).leftToRight y))
  }

-- Downward Closed -------------------------------------------------------------

public export 0
DClosed : Property sig k -> Type
DClosed p =
  forall m, n.
  (f : Fin k -> Term sig m) ->
  (g : Fin k -> Term sig n) ->
  f <= g ->
  p.Prop g ->
  p.Prop f

-- Properties

unifiesDClosed : (t, u : Term sig k) -> DClosed (Unifies t u)
unifiesDClosed t u f g prf1 prf2 = Calc $
  |~ f <$> t
  ~~ (prf1.sub . g) <$> t ...(subCong prf1.prf t)
  ~~ prf1.sub <$> g <$> t ...(subAssoc prf1.sub g t)
  ~~ prf1.sub <$> g <$> u ...(cong (prf1.sub <$>) prf2)
  ~~ (prf1.sub . g) <$> u ..<(subAssoc prf1.sub g u)
  ~~ f <$> u              ..<(subCong prf1.prf u)

optimistLemma :
  {q : Property sig j} ->
  {a : Fin j -> Term sig k} ->
  {f : Fin k -> Term sig m} ->
  {g : Fin m -> Term sig n} ->
  DClosed p ->
  (Max (Extension p a)).Prop f ->
  (Max (Extension q (f . a))).Prop g ->
  (Max (Extension (All [p, q]) a)).Prop (g . f)
optimistLemma closed (x, max1) (y, max2) =
  ([ closed ((g . f) . a) (f . a) (compLte (MkLte g (\i => Refl)) a) x
  ,  q.cong (\i => sym $ subAssoc g f (a i)) y
  ], \h, [x', y'] =>
    let prf1 = max1 h x' in
    let y'' = q.cong (\i => trans (subCong prf1.prf (a i)) (subAssoc prf1.sub f (a i))) y' in
    let prf2 = max2 prf1.sub y'' in
    MkLte
      { sub = prf2.sub
      , prf = \i => Calc $
        |~ h i
        ~~ prf1.sub <$> f i         ...(prf1.prf i)
        ~~ (prf2.sub . g) <$> f i   ...(subCong (prf2.prf) (f i))
        ~~ prf2.sub <$> (g <$> f i) ...(subAssoc prf2.sub g (f i))
      }
  )