module Data.Term.Property

import Data.Term

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

import Syntax.PreorderReasoning

%prefix_record_projections off

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

public export
record Property (sig : Signature) (k : Nat) where
  constructor MkProp
  0 Prop : forall n. (SFin k -> Term sig n) -> Type
  cong : forall n. {f, g : SFin 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))

public export
UnifiesAll : Vect n (Term sig k) -> Vect n (Term sig k) -> Property sig k
UnifiesAll ts us = All (zipWith Unifies ts us)

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

public export
record (<=>) (p, q : Property sig k) where
  constructor MkEquivalence
  leftToRight : forall n. (f : SFin k -> Term sig n) -> p.Prop f -> q.Prop f
  rightToLeft : forall n. (f : SFin k -> Term sig n) -> q.Prop f -> p.Prop f

-- Properties

export
Reflexive (Property sig k) (<=>) where
  reflexive = MkEquivalence (\_ => id) (\_ => id)

export
Symmetric (Property sig k) (<=>) where
  symmetric prf = MkEquivalence prf.rightToLeft prf.leftToRight

export
unifiesSym : (0 t, u : Term sig k) -> Unifies t u <=> Unifies u t
unifiesSym t u = MkEquivalence (\_, prf => sym prf) (\_, prf => sym prf)

export
unifiesOp :
  {0 k : Nat} ->
  (0 op : sig.Operator k) ->
  (ts, us : Vect k (Term sig j)) ->
  Unifies (Op op ts) (Op op us) <=> UnifiesAll ts us
unifiesOp op ts us = MkEquivalence
  { leftToRight = \f, prf => leftToRight ts us f (opInjectiveTs' prf)
  , rightToLeft = \f, prf => cong (Op op) $ rightToLeft ts us f prf
  }
  where
  leftToRight :
    forall k.
    (ts, us : Vect k (Term sig j)) ->
    (0 f : (SFin j -> Term sig n)) ->
    map (f <$>) ts = map (f <$>) us ->
    (UnifiesAll ts us).Prop f
  leftToRight [] [] f prf = []
  leftToRight (t :: ts) (u :: us) f prf =
    fst (biinj (::) prf) :: leftToRight ts us f (snd $ biinj (::) prf)

  rightToLeft :
    forall k.
    (ts, us : Vect k (Term sig j)) ->
    (0 f : (SFin j -> Term sig n)) ->
    (UnifiesAll ts us).Prop f ->
    map (f <$>) ts = map (f <$>) us
  rightToLeft [] [] f [] = Refl
  rightToLeft (t :: ts) (u :: us) f (prf :: prfs) = cong2 (::) prf (rightToLeft ts us f prfs)

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

public export 0
Nothing : Property sig k -> Type
Nothing p = forall n. (f : SFin 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.rightToLeft f x)

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

public export
Extension : Property sig k -> (SFin 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
extendCong2 :
  {f, g : SFin n -> Term sig k} ->
  {p, q : Property sig n} ->
  f .=. g ->
  p <=> q ->
  Extension p f <=> Extension q g
extendCong2 prf1 prf2 = MkEquivalence
  (\h, x => prf2.leftToRight (h . g) $ p.cong (\i => cong (h <$>) $ prf1 i) x)
  (\h, x => prf2.rightToLeft (h . f) $ q.cong (\i => sym $ cong (h <$>) $ prf1 i) x)

export
extendCong :
  (f : SFin n -> Term sig k) ->
  p <=> q ->
  Extension p f <=> Extension q f
extendCong f prf = MkEquivalence
  (\g => prf.leftToRight (g . f))
  (\g => prf.rightToLeft (g . f))

export
extendUnit : (p : Property sig k) -> p <=> Extension p Var'
extendUnit p = MkEquivalence (\_, x => x) (\_, x => x)

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

export
extendUnify :
  (t, u : Term sig j) ->
  (f : SFin j -> Term sig k) ->
  (g : SFin k -> Term sig m) ->
  Extension (Unifies t u) (g . f) <=> Extension (Unifies (f <$> t) (f <$> u)) g
extendUnify t u f g =
  MkEquivalence
    (\h, 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))
    (\h, 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))

export
extendUnifyAll :
  (ts, us : Vect n (Term sig j)) ->
  (f : SFin j -> Term sig k) ->
  (g : SFin k -> Term sig m) ->
  Extension (UnifiesAll ts us) (g . f) <=>
  Extension (UnifiesAll (map (f <$>) ts) (map (f <$>) us)) g
extendUnifyAll [] [] f g = MkEquivalence (\h, [] => []) (\h, [] => [])
extendUnifyAll (t :: ts) (u :: us) f g =
  let head = extendUnify t u f g in
  let tail = extendUnifyAll ts us f g in
  MkEquivalence
    (\h, (x :: xs) => head.leftToRight h x :: tail.leftToRight h xs)
    (\h, (x :: xs) => head.rightToLeft h x :: tail.rightToLeft h xs)

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

public export
record (<=) (f : SFin k -> Term sig m) (g : SFin k -> Term sig n) where
  constructor MkLte
  sub : SFin 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 (SFin 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 : SFin k -> Term sig m) -> f <= Var'
varMax f = MkLte f (\i => Refl)

export
compLte : f <= g -> (h : SFin 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 : SFin k -> Term sig m} ->
  {g : SFin 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 : SFin 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 = MkEquivalence
  { leftToRight = \f, (x, max) => (prf.leftToRight f x, \g, y => max g (prf.rightToLeft g y))
  , rightToLeft = \f, (x, max) => (prf.rightToLeft f x, \g, y => max g (prf.leftToRight g y))
  }

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

public export
record DClosed (p : Property sig k) where
  constructor MkDClosed
  closed :
    forall m, n.
    (f : SFin k -> Term sig m) ->
    (g : SFin k -> Term sig n) ->
    f <= g ->
    p.Prop g ->
    p.Prop f

-- Properties

export
unifiesDClosed : (t, u : Term sig k) -> DClosed (Unifies t u)
unifiesDClosed t u = MkDClosed (\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))

export
optimistLemma :
  {ps : Vect _ (Property sig j)} ->
  {a : SFin j -> Term sig k} ->
  {f : SFin k -> Term sig m} ->
  {g : SFin m -> Term sig n} ->
  DClosed p ->
  (Max (Extension p a)).Prop f ->
  (Max (Extension (All ps) (f . a))).Prop g ->
  (Max (Extension (All (p :: ps)) a)).Prop (g . f)
optimistLemma prf (x, max1) (y, max2) =
  ( ( prf.closed ((g . f) . a) (f . a) (compLte (MkLte g (\i => Refl)) a) x
    :: (All ps).cong (\i => sym $ subAssoc g f (a i)) y
    )
  , \h, (x' :: y') =>
    let prf1 = max1 h x' in
    let y'' = (All ps).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))
      }
  )

export
failHead : Nothing (Extension p a) -> Nothing (Extension (All (p :: ps)) a)
failHead absurd f (x :: xs) = absurd f x

export
failTail :
  {ps : Vect _ (Property sig j)} ->
  {a : SFin j -> Term sig k} ->
  {f : SFin k -> Term sig m} ->
  (Max (Extension p a)).Prop f ->
  Nothing (Extension (All ps) (f . a)) ->
  Nothing (Extension (All (p :: ps)) a)
failTail (x, max) absurd g (y :: ys) =
  let prf = max g y in
  absurd prf.sub $ (All ps).cong (compLte prf a).prf ys