module Data.Term.Unify

import Data.Fin.Occurs
import Data.Term
import Data.Term.Property
import Data.Term.Zipper
import Data.Vect.Quantifiers.Extra

import Decidable.Equality

import Syntax.PreorderReasoning

-- Check -----------------------------------------------------------------------

export
check : {k : Nat} -> Fin (S k) -> Term sig (S k) -> Maybe (Term sig k)
checkAll : {k : Nat} -> Fin (S k) -> Vect n (Term sig (S k)) -> Maybe (Vect n (Term sig k))

check x (Var y) = Var <$> thick x y
check x (Op op ts) = Op op <$> checkAll x ts

checkAll x [] = Just []
checkAll x (t :: ts) = [| check x t :: checkAll x ts |]

-- Properties

public export
data CheckProof : Fin (S k) -> Term sig (S k) -> Maybe (Term sig k) -> Type where
  Occurs : (zip : Zipper sig (S k)) -> res = Nothing -> CheckProof x (zip + Var x) res
  Stronger : (t : Term sig k) -> res = Just t -> CheckProof x (pure (thin x) <$> t) res

data CheckAllProof : Fin (S k) -> Vect n (Term sig (S k)) -> Maybe (Vect n (Term sig k)) -> Type where
  OccursAt :
    (i : Fin (S n)) ->
    (ts : Vect n (Term sig (S k))) ->
    (zip : Zipper sig (S k)) ->
    res = Nothing ->
    CheckAllProof x (insertAt i (zip + Var x) ts) res
  AllStronger :
    (ts : Vect n (Term sig k)) ->
    res = Just ts ->
    CheckAllProof x (map (pure (thin x) <$>) ts) res

export
checkProof : {k : Nat} -> (x : Fin (S k)) -> (t : Term sig (S k)) -> CheckProof x t (check x t)
checkAllProof :
  {k : Nat} ->
  (x : Fin (S k)) ->
  (ts : Vect n (Term sig (S k))) ->
  CheckAllProof x ts (checkAll x ts)

checkProof x (Var y) with (thickProof x y)
  checkProof x (Var x) | Equal prf = Occurs Top (cong (Var <$>) prf)
  checkProof x (Var _) | Thinned y prf = Stronger (Var y) (cong (Var <$>) prf)
checkProof x (Op op ts) with (checkAllProof x ts)
  checkProof x (Op op _) | OccursAt i ts zip prf = Occurs (Op op i ts zip) (cong (Op op <$>) prf)
  checkProof x (Op op _) | AllStronger ts prf = Stronger (Op op ts) (cong (Op op <$>) prf)

checkAllProof x [] = AllStronger [] Refl
checkAllProof x (t :: ts) with (checkProof x t)
  checkAllProof x (_ :: ts) | Occurs zip prf =
    OccursAt FZ ts zip (cong (\t => [| t :: checkAll x ts |]) prf)
  checkAllProof x (_ :: ts) | Stronger u prf with (checkAllProof x ts)
    checkAllProof x (_ :: _) | Stronger u prf | OccursAt i ts zip prf' =
      OccursAt (FS i) (_ :: ts) zip (cong2 (\t, ts => [| t :: ts |]) prf prf')
    checkAllProof x (_ :: _) | Stronger u prf | AllStronger us prf' =
      AllStronger (u :: us) (cong2 (\t, ts => [| t :: ts |]) prf prf')

-- Single Variable Substitution ------------------------------------------------

export
for : {k : Nat} -> Term sig k -> Fin (S k) -> Fin (S k) -> Term sig k
(t `for` x) y = maybe t Var (thick x y)

export
forRefl :
  (0 u : Term sig k) ->
  (x : Fin (S k)) ->
  (u `for` x) x = u
forRefl u x = cong (maybe u Var) $ thickRefl x

export
forThin :
  (0 t : Term sig k) ->
  (x : Fin (S k)) ->
  (t `for` x) . thin x .=. Var
forThin t x i = cong (maybe t Var) (thickThin x i)

export
forUnifies :
  (x : Fin (S k)) ->
  (t : Term sig k) ->
  (t `for` x) <$> pure (thin x) <$> t = (t `for` x) <$> Var x
forUnifies x t = Calc $
  |~ (t `for` x) <$> pure (thin x) <$> t
  ~~ (t `for` x) . thin x <$> t          ...(sym $ subAssoc (t `for` x) (pure (thin x)) t)
  ~~ Var <$> t                           ...(subExtensional (forThin t x) t)
  ~~ t                                   ...(subUnit t)
  ~~ (t `for` x) <$> Var x               ...(sym $ forRefl t x)

export
varElim :
  {n : Nat} ->
  (x : Fin (S n)) ->
  (t : Term sig n) ->
  MostGeneral (Unifies (Var x) (pure (thin x) <$> t)) (t `for` x)
varElim x t =
  MkMostGeneral
    { valid = sym $ forUnifies x t
    , universal = \g, prf => g . thin x
    , factors = factors
    , unique = unique
    }
  where
  factors :
    forall k.
    (g : Fin (S n) -> Term sig k) ->
    (Unifies (Var x) (pure (thin x) <$> t)).Prop g ->
    g .=. (g . thin x) . (t `for` x)
  factors g prf i with (thickProof x i)
    factors g prf _ | Equal prf' = Calc $
      |~ g x
      ~~ g <$> pure (thin x) <$> t    ...(prf)
      ~~ g . thin x <$> t             ..<(subAssoc g (pure (thin x)) t)
      ~~ g . thin x <$> (t `for` x) x ..<(cong ((g . thin x <$>) . maybe t Var) prf')
    factors g prf _ | Thinned i prf' = sym $ cong (g . thin x <$>) (forThin t x i)

  unique :
    forall k.
    (g : Fin (S n) -> Term sig k) ->
    (Unifies (Var x) (pure (thin x) <$> t)).Prop g ->
    (h : Fin n -> Term sig k) ->
    g .=. h . (t `for` x) ->
    h .=. g . thin x
  unique g prf h prf' i = Calc $
    |~ h i
    ~~ h <$> (t `for` x) (thin x i) ..<(cong (h <$>) $ forThin t x i)
    ~~ g (thin x i)                 ..<(prf' (thin x i))

-- Substitution List -----------------------------------------------------------

public export
data AList : Signature -> Nat -> Nat -> Type where
  Lin : AList sig n n
  (:<) : AList sig k n -> (Term sig k, Fin (S k)) -> AList sig (S k) n

%name AList sub

namespace Exists
  public export
  (:<) : (k ** AList sig n k) -> (Term sig n, Fin (S n)) -> (k ** AList sig (S n) k)
  (_ ** sub) :< tx = (_ ** sub :< tx)

export
eval : {k : Nat} -> AList sig k n -> Fin k -> Term sig n
eval [<] = Var
eval (sub :< (t, x)) = eval sub . (t `for` x)

export
(++) : AList sig k n -> AList sig j k -> AList sig j n
sub ++ [<] = sub
sub ++ sub1 :< tx = (sub ++ sub1) :< tx

-- Properties

export
evalHomo :
  (0 sub2 : AList sig k n) ->
  (sub1 : AList sig j k) ->
  eval (sub2 ++ sub1) .=. eval sub2 . eval sub1
evalHomo sub2 [<] i = Refl
evalHomo sub2 (sub1 :< (t, x)) i = Calc $
  |~ eval (sub2 ++ sub1) <$> (t `for` x) i
  ~~ (eval sub2 . eval sub1) <$> (t `for` x) i ...(subExtensional (evalHomo sub2 sub1) ((t `for` x) i))
  ~~ eval sub2 <$> eval sub1 <$> (t `for` x) i ...(subAssoc (eval sub2) (eval sub1) ((t `for` x) i))

export
appendUnitLeft : (sub : AList sig k n) -> [<] ++ sub = sub
appendUnitLeft [<] = Refl
appendUnitLeft (sub :< tx) = cong (:< tx) (appendUnitLeft sub)

export
appendAssoc :
  (sub3 : AList sig _ _) ->
  (sub2 : AList sig _ _) ->
  (sub1 : AList sig _ _) ->
  sub3 ++ (sub2 ++ sub1) = (sub3 ++ sub2) ++ sub1
appendAssoc sub3 sub2 [<] = Refl
appendAssoc sub3 sub2 (sub1 :< tx) = cong (:< tx) (appendAssoc sub3 sub2 sub1)

-- Unification -----------------------------------------------------------------

coerce : {0 op, op' : sig.Op} -> op = op' -> Vect op'.fst a -> Vect op.fst a
coerce Refl = id

flexFlex : {n : Nat} -> Fin n -> Fin n -> (k ** AList sig n k)
flexFlex {n = S n} x y = case thick x y of
  Just z => (_ **[<(Var z, x)])
  Nothing => (_ ** [<])

flexRigid : {n : Nat} -> Fin n -> Term sig n -> Maybe (k ** AList sig n k)
flexRigid {n = S n} x t = case check x t of
  Just u => Just (_ ** [<(u, x)])
  Nothing => Nothing

export
amgu :
  DecOp sig =>
  {k, n : Nat} ->
  (t, u : Term sig n) ->
  AList sig n k ->
  Maybe (k ** AList sig n k)
amguAll :
  DecOp sig =>
  {k, n : Nat} ->
  (ts, us : Vect j (Term sig n)) ->
  AList sig n k ->
  Maybe (k ** AList sig n k)

amgu (Op op ts) (Op op' us) acc with (decOp (Evidence _ op) (Evidence _ op'))
  amgu (Op op ts) (Op op us) acc | Yes Refl = amguAll ts us acc
  _ | No neq = Nothing
amgu (Var x) (Var y) [<] = Just (flexFlex x y)
amgu (Var x) (Op op' us) [<] = flexRigid x (Op op' us)
amgu (Op op ts) (Var y) [<] = flexRigid y (Op op ts)
amgu t@(Var _) u@(Var _) (sub :< (v, z)) =
  (:< (v, z)) <$> amgu ((v `for` z) <$> t) ((v `for` z) <$> u) sub
amgu t@(Var _) u@(Op _ _) (sub :< (v, z)) =
  (:< (v, z)) <$> amgu ((v `for` z) <$> t) ((v `for` z) <$> u) sub
amgu t@(Op _ _) u@(Var _) (sub :< (v, z)) =
  (:< (v, z)) <$> amgu ((v `for` z) <$> t) ((v `for` z) <$> u) sub

amguAll [] [] acc = Just (_ ** acc)
amguAll (t :: ts) (u :: us) acc = do
  acc <- amgu t u acc
  amguAll ts us acc.snd

export
mgu : DecOp sig => {n : Nat} -> (t, u : Term sig n) -> Maybe (k ** AList sig n k)
mgu t u = amgu t u [<]

-- Properties

export
trivialUnify :
  (0 f : Fin n -> Term sig k) ->
  (t : Term sig n) ->
  MostGeneral (Extension (Unifies t t) f) Var
trivialUnify f t = varMostGeneral Refl

flexFlexUnifies :
  {n : Nat} ->
  (x, y : Fin n) ->
  MostGeneral {sig} (Unifies (Var x) (Var y)) (eval (flexFlex x y).snd)
flexFlexUnifies {n = S n} x y with (thickProof x y)
  flexFlexUnifies {n = S n} x _ | Equal prf =
    rewrite prf in
    trivialUnify Var (Var x)
  flexFlexUnifies {n = S n} x _ | Thinned y prf =
    rewrite prf in
    (UniqueMax (Unifies (Var x) (Var _))).cong _ _
      (\i => sym $ subUnit ((Var y `for` x) i))
      (varElim x (Var y))

data FlexRigidProof : Fin n -> Term sig n -> Maybe (k ** AList sig n k) -> Type where
  NoUnifier : Nothing (Unifies (Var x) t) -> res = Nothing -> FlexRigidProof x t res
  ElimVar :
    {n : Nat} ->
    (sub : AList sig k n) ->
    (mgu : MostGeneral (Unifies (Var x) t) (eval sub)) ->
    res = Just (_ ** sub) ->
    FlexRigidProof x t res

flexRigidProof :
  {n : Nat} ->
  (x : Fin n) ->
  (op : sig.Operator k) ->
  (ts : Vect k (Term sig n)) ->
  FlexRigidProof x (Op op ts) (flexRigid x (Op op ts))
flexRigidProof {n = S n} x op ts with (checkProof x (Op op ts))
  flexRigidProof x op _ | Occurs zip@(Op op i ts zip') prf =
    rewrite prf in
    NoUnifier
      (\f, prf' =>
        let
          cycle : (f x = (f <$> zip) + f x)
          cycle = Calc $
           |~ f x
           ~~ f <$> Op op _     ...(prf')
           ~~ (f <$> zip) + f x ...(actionHomo f zip (Var x))

          zipIsTop : (zip = Top)
          zipIsTop = invertActionTop zip $ noCycles (f <$> zip) (f x) (sym cycle)

          opIsVar : Op op (insertAt i (zip' + Var x) ts) = Var x
          opIsVar = cong (+ Var x) zipIsTop
        in
        absurd opIsVar)
      Refl
  flexRigidProof x op _ | Stronger (Op op us) prf =
    rewrite prf in
    ElimVar
      [<(Op op us, x)]
      ((UniqueMax (Unifies (Var x) (Op op _))).cong _ _
        (\i => sym $ subUnit ((Op op us `for` x) i))
        (varElim x (Op op us)))
      Refl

stepEquiv :
  {k : Nat} ->
  (t, u : Term sig (S k)) ->
  (sub : AList sig k j) ->
  (v : Term sig k) ->
  (x : Fin (S k)) ->
  Extension (Unifies t u) (eval (sub :< (v, x))) <=>
  Extension (Unifies ((v `for` x) <$> t) ((v `for` x) <$> u)) (eval sub)
stepEquiv t u sub v x = extendUnify t u (v `for` x) (eval sub)

parameters {auto _ : DecOp sig}
  public export
  data AmguProof :
    Term sig k ->
    Term sig k ->
    AList sig k n ->
    Maybe (n ** AList sig k n) ->
    Type
    where
    Failure :
      Nothing (Extension (Unifies t u) (eval sub)) ->
      res = Nothing ->
      AmguProof t u sub res
    Success :
      {j : Nat} ->
      {0 sub : AList sig k n} ->
      (sub' : AList sig n j) ->
      MostGeneral (Extension (Unifies t u) (eval sub)) (eval sub') ->
      res = Just (_ ** sub' ++ sub) ->
      AmguProof t u sub res

  public export
  data AmguAllProof :
    Vect j (Term sig k) ->
    Vect j (Term sig k) ->
    AList sig k n ->
    Maybe (n ** AList sig k n) ->
    Type
    where
    FailureAll :
      Nothing (Extension (UnifiesAll ts us) (eval sub)) ->
      res = Nothing ->
      AmguAllProof ts us sub res
    SuccessAll :
      {j : Nat} ->
      {0 sub : AList sig k n} ->
      (sub' : AList sig n j) ->
      MostGeneral (Extension (UnifiesAll ts us) (eval sub)) (eval sub') ->
      res = Just (_ ** sub' ++ sub) ->
      AmguAllProof ts us sub res

export
amguProof :
  DecOp sig =>
  {k, n : Nat} ->
  (t, u : Term sig k) ->
  (sub : AList sig k n) ->
  AmguProof t u sub (amgu t u sub)
export
amguAllProof :
  DecOp sig =>
  {k, n : Nat} ->
  (ts, us : Vect j (Term sig k)) ->
  (sub : AList sig k n) ->
  AmguAllProof ts us sub (amguAll ts us sub)

amguProof (Op op ts) (Op op' us) sub with (decOp (Evidence _ op) (Evidence _ op')) proof prf
  amguProof (Op op ts) (Op op us) sub | Yes Refl
    with (amguAllProof ts us sub)
    _ | SuccessAll sub' val prf' =
      let
        cong :
          UniqueMax (Extension (Unifies (Op op ts) (Op op us)) (eval sub)) <=>
          UniqueMax (Extension (UnifiesAll ts us) (eval sub))
        cong = maxCong $ extendCong (eval sub) $ unifiesOp op ts us
      in
      Success sub'
        (cong.rightToLeft (eval sub') val)
        prf'
    _ | FailureAll absurd prf' =
      Failure
        (nothingEquiv (symmetric $ extendCong (eval sub) $ unifiesOp op ts us) absurd)
        prf'
  _ | No neq =
    Failure
      (\f, prf => neq $ opInjectiveOp prf)
      Refl
amguProof (Var x) (Var y) [<] =
  Success
    (flexFlex x y).snd
    (flexFlexUnifies x y)
    (cong Just $ ext _ (flexFlex x y))
  where
  ext : (0 b : a -> Type) -> (v : (x : a ** b x)) -> v = (v.fst ** v.snd)
  ext b (fst ** snd) = Refl
amguProof (Var x) (Op op' us) [<] with (flexRigidProof x op' us)
  _ | NoUnifier absurd prf = Failure absurd prf
  _ | ElimVar sub val prf = Success sub val prf
amguProof (Op op ts) (Var y) [<] with (flexRigidProof y op ts)
  _ | NoUnifier absurd prf =
    let
      cong :
        Nothing (Extension (Unifies (Var y) (Op op ts)) Var) ->
        Nothing (Extension (Unifies (Op op ts) (Var y)) Var)
      cong = nothingEquiv $ extendCong Var $ unifiesSym (Var y) (Op op ts)
    in
    Failure (cong absurd) prf
  _ | ElimVar sub val prf =
    let
      cong :
        UniqueMax (Extension (Unifies (Var y) (Op op ts)) Var) <=>
        UniqueMax (Extension (Unifies (Op op ts) (Var y)) Var)
      cong = maxCong $ extendCong Var $ unifiesSym (Var y) (Op op ts)
    in
    Success sub (cong.leftToRight _ val) prf
amguProof (Var x) (Var y) (sub :< (v, z))
  with (amguProof ((v `for` z) x) ((v `for` z) y) sub)
  _ | Success sub' val prf =
    let
      cong' :
        UniqueMax (Extension (Unifies (Var x) (Var y)) (eval sub . (v `for` z))) <=>
        UniqueMax (Extension (Unifies ((v `for` z) x) ((v `for` z) y)) (eval sub))
      cong' = maxCong $ stepEquiv (Var x) (Var y) sub v z
    in
    Success sub'
      (cong'.rightToLeft (eval sub') val)
      (cong (map (:< (v, z))) prf)
  _ | Failure absurd prf =
    Failure
      (nothingEquiv (symmetric $ stepEquiv (Var x) (Var y) sub v z) absurd)
      (rewrite prf in Refl)
amguProof (Var x) (Op op' us) (sub :< (v, z))
  with (amguProof ((v `for` z) x) ((v `for` z) <$> Op op' us) sub)
  _ | Success sub' val prf =
    let
      cong' :
        UniqueMax (Extension (Unifies (Var x) (Op op' us)) (eval sub . (v `for` z))) <=>
        UniqueMax (Extension (Unifies ((v `for` z) x) ((v `for` z) <$> Op op' us)) (eval sub))
      cong' = maxCong $ stepEquiv (Var x) (Op op' us) sub v z
    in
    Success sub'
      (cong'.rightToLeft (eval sub') val)
      (cong (map (:< (v, z))) prf)
  _ | Failure absurd prf =
    Failure
      (nothingEquiv (symmetric $ stepEquiv (Var x) (Op op' us) sub v z) absurd)
      (rewrite prf in Refl)
amguProof (Op op ts) (Var y) (sub :< (v, z))
  with (amguProof ((v `for` z) <$> Op op ts) ((v `for` z) y) sub)
  _ | Success sub' val prf =
    let
      cong' :
        UniqueMax (Extension (Unifies (Op op ts) (Var y)) (eval sub . (v `for` z))) <=>
        UniqueMax (Extension (Unifies ((v `for` z) <$> Op op ts) ((v `for` z) y)) (eval sub))
      cong' = maxCong $ stepEquiv (Op op ts) (Var y) sub v z
    in
    Success sub'
      (cong'.rightToLeft (eval sub') val)
      (cong (map (:< (v, z))) prf)
  _ | Failure absurd prf =
    Failure
      (nothingEquiv (symmetric $ stepEquiv (Op op ts) (Var y) sub v z) absurd)
      (rewrite prf in Refl)

amguAllProof [] [] sub =
  SuccessAll [<]
    (varMostGeneral [])
    (sym $ cong (\sub => Just (_ ** sub)) $ appendUnitLeft sub)
amguAllProof (t :: ts) (u :: us) sub with (amguProof t u sub)
  _ | Success sub' val prf with (amguAllProof ts us (sub' ++ sub))
    _ | SuccessAll sub'' val' prf' =
      let
        cong' =
          maxCong $
          extendCong2 (evalHomo sub' sub) (reflexive {x = UnifiesAll ts us})
        opt =
          optimistLemma (unifiesDClosed t u) val $
          cong'.leftToRight (eval sub'') val'
      in
      SuccessAll (sub'' ++ sub')
        ((UniqueMax (Extension (UnifiesAll (t :: ts) (u :: us)) (eval sub))).cong _ _
          (\i => sym $ evalHomo sub'' sub' i)
          opt)
        (rewrite prf in rewrite prf' in
          cong (\sub => Just (_ ** sub)) $
          appendAssoc sub'' sub' sub)
    _ | FailureAll absurd prf' =
      let
        cong' = extendCong2 (evalHomo sub' sub) (reflexive {x = UnifiesAll ts us})
      in
      FailureAll
        (failTail val $ nothingEquiv cong' absurd)
        (rewrite prf in prf')
  _ | Failure absurd prf =
    FailureAll
      (failHead absurd)
      (rewrite prf in Refl)

parameters {auto _ : DecOp sig}
  namespace MguProof
    public export
    data MguProof : (t, u : Term sig k) -> Maybe (n ** AList sig k n) -> Type where
      Success :
        {n : Nat} ->
        (sub : AList sig k n) ->
        MostGeneral (Unifies t u) (eval sub) ->
        res = Just (_ ** sub) ->
        MguProof t u res
      Failure :
        Nothing (Unifies t u) ->
        res = Nothing ->
        MguProof t u res

export
mguProof : DecOp sig => {k : Nat} -> (t, u : Term sig k) -> MguProof t u (mgu t u)
mguProof t u with (amguProof t u [<])
  _ | Success sub val prf = Success sub val prf
  _ | Failure absurd prf = Failure absurd prf

export
invertMgu :
  DecOp sig =>
  {0 t, u : Term sig k} ->
  MguProof t u res ->
  res = Just (_ ** sub) ->
  MostGeneral (Unifies t u) (eval sub)
invertMgu (Success sub mg Refl) Refl = mg
invertMgu (Failure _ prf) Refl = absurd prf