module Data.Term.Unify

import Data.DPair
import Data.Fin.Occurs
import Data.Maybe.Properties
import Data.Term
import Data.Term.Zipper

import Decidable.Equality

import Syntax.PreorderReasoning

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

export
check : Fin k -> Term sig k -> Maybe (Term sig (pred k))
checkAll : Fin k -> Vect n (Term sig k) -> Maybe (Vect n (Term sig (pred 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

export
checkOccurs : (x : Fin k) -> (zip : Zipper sig k) -> IsNothing (check x (zip + Var x))
checkAllOccurs :
  (x : Fin k) ->
  (i : Fin (S n)) ->
  (ts : Vect n (Term sig k)) ->
  (zip : Zipper sig k) ->
  IsNothing (checkAll x (insertAt i (zip + Var x) ts))

checkOccurs x Top = mapNothing Var (thickRefl x)
checkOccurs x (Op op i ts zip) = mapNothing (Op op) (checkAllOccurs x i ts zip)

checkAllOccurs x FZ ts zip =
  appLeftNothingIsNothing
    (appRightNothingIsNothing (Just (::)) (checkOccurs x zip))
    (checkAll x ts)
checkAllOccurs x (FS i) (t :: ts) zip =
  appRightNothingIsNothing
    (Just (::) <*> check x t)
    (checkAllOccurs x i ts zip)

export
checkNothing :
  (x : Fin k) ->
  (t : Term sig k) ->
  (0 _ : IsNothing (check x t)) ->
  (zip : Zipper sig k ** t = zip + Var x)
checkAllNothing :
  (x : Fin k) ->
  (t : Term sig k) ->
  (ts : Vect n (Term sig k)) ->
  (0 _ : IsNothing (checkAll x (t :: ts))) ->
  (i ** ts' ** zip : Zipper sig k ** t :: ts = insertAt i (zip + Var x) ts')

checkNothing x (Var y) prf =
  (Top ** sym (cong Var (thickNothing x y (mapNothingInverse Var (thick x y) prf))))
checkNothing x (Op op (t :: ts)) prf =
  let (i ** ts' ** zip ** prf) = checkAllNothing x t ts (mapNothingInverse (Op op) _ prf) in
  (Op op i ts' zip ** cong (Op op) prf)

checkAllNothing x t ts prf with (appNothingInverse (Just (::) <*> check x t) (checkAll x ts) prf)
  _ | Left prf' = case appNothingInverse (Just (::)) (check x t) prf' of
      Right prf =>
        let (zip ** prf) = checkNothing x t prf in
        (FZ ** ts ** zip ** cong (:: ts) prf)
  checkAllNothing x t (t' :: ts) prf | Right prf' =
    let (i ** ts ** zip ** prf) = checkAllNothing x t' ts prf' in
    (FS i ** t :: ts ** zip ** cong (t ::) prf)

export
checkThin : (x : Fin k) -> (t : Term sig (pred k)) -> IsJust (check x (pure (thin x) <$> t))
checkAllThin :
  (x : Fin k) ->
  (ts : Vect n (Term sig (pred k))) ->
  IsJust (checkAll x (map (pure (thin x) <$>) ts))

checkThin x (Var y) = mapIsJust Var (thickNeq x (thin x y) (\prf => thinSkips x y $ sym prf))
checkThin x (Op op ts) = mapIsJust (Op op) (checkAllThin x ts)

checkAllThin x [] = ItIsJust
checkAllThin x (t :: ts) =
  appIsJust
    (appIsJust ItIsJust (checkThin x t))
    (checkAllThin x ts)

export
checkJust :
  (x : Fin k) ->
  (t : Term sig k) ->
  (0 _ : check x t = Just u) ->
  t = pure (thin x) <$> u
checkAllJust :
  (x : Fin k) ->
  (ts : Vect n (Term sig k)) ->
  (0 _ : checkAll x ts = Just us) ->
  ts = map (pure (thin x) <$>) us

checkJust x (Var y) prf =
  let 0 z = mapJustInverse Var (thick x y) prf in
  let 0 prf = thickJust x y (fst z.snd) in
  sym $ Calc $
    |~ pure (thin x) <$> u
    ~~ pure (thin x) <$> Var z.fst ...(cong (pure (thin x) <$>) $ snd z.snd)
    ~~ Var y                       ...(cong Var prf)
checkJust x (Op op ts) prf =
  let 0 z = mapJustInverse (Op op) (checkAll x ts) prf in
  let 0 prf = checkAllJust x ts (fst z.snd) in
  Calc $
    |~ Op op ts
    ~~ pure (thin x) <$> Op op z.fst ...(cong (Op op) prf)
    ~~ pure (thin x) <$> u           ...(sym $ cong (pure (thin x) <$>) $ snd z.snd)

checkAllJust x [] Refl = Refl
checkAllJust x (t :: ts) prf =
  let 0 z = appJustInverse (Just (::) <*> check x t) (checkAll x ts) prf in
  let 0 w = appJustInverse (Just (::)) (check x t) (fst z.snd.snd) in
  Calc $
    |~ t :: ts
    ~~ map (pure (thin x) <$>) (w.snd.fst :: z.snd.fst)    ...(cong2 (::) (checkJust x t (fst $ snd w.snd.snd)) (checkAllJust x ts (fst $ snd z.snd.snd)))
    ~~ map (pure (thin x) <$>) (w.fst w.snd.fst z.snd.fst) ...(cong (\f => map (pure (thin x) <$>) (f w.snd.fst z.snd.fst)) $ injective $ fst w.snd.snd)
    ~~ map (pure (thin x) <$>) (z.fst z.snd.fst)           ...(sym $ cong (\f => map (pure (thin x) <$>) (f z.snd.fst)) $ snd $ snd w.snd.snd)
    ~~ map (pure (thin x) <$>) us                          ...(sym $ cong (map (pure (thin x) <$>)) $ snd $ snd z.snd.snd)

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

export
for : Term sig (pred k) -> Fin k -> Fin k -> Term sig (pred k)
(t `for` x) y = maybe t Var (thick x y)

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

export
forUnifies :
  (x : Fin k) ->
  (t : Term sig k) ->
  (0 _ : check x t = Just u) ->
  (u `for` x) <$> t = (u `for` x) <$> Var x
forUnifies x t prf = Calc $
  |~ (u `for` x) <$> t
  ~~ (u `for` x) <$> pure (thin x) <$> u ...(cong ((u `for` x) <$>) $ checkJust x t prf)
  ~~ (u `for` x) . thin x <$> u          ...(sym $ subAssoc (u `for` x) (pure (thin x)) u)
  ~~ Var <$> u                           ...(subCong (forThin u x) u)
  ~~ u                                   ...(subUnit u)
  ~~ (u `for` x) <$> Var x               ...(sym $ cong (maybe u Var) $ extractIsNothing $ thickRefl x)

-- 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
  (:<) : Exists (AList sig n) -> (Term sig n, Fin (S n)) -> Exists (AList sig (S n))
  Evidence _ sub :< tx = Evidence _ (sub :< tx)

export
eval : 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 ...(subCong (evalHomo sub2 sub1) ((t `for` x) i))
  ~~ eval sub2 <$> eval sub1 <$> (t `for` x) i ...(subAssoc (eval sub2) (eval sub1) ((t `for` x) i))

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

flexFlex : (x, y : Fin n) -> Exists (AList sig n)
flexFlex x y =
  rewrite (snd $ indexIsSuc x) in
  case thick x' y' of
    Just z => Evidence _ [<(Var z, x')]
    Nothing => Evidence _ [<]
  where
  x', y' : Fin (S $ fst $ indexIsSuc x)
  x' = replace {p = Fin} (snd $ indexIsSuc x) x
  y' = replace {p = Fin} (snd $ indexIsSuc x) y

flexRigid : Fin n -> Term sig n -> Maybe (Exists (AList sig n))
flexRigid x t =
  rewrite (snd $ indexIsSuc x) in
  case check x' t' of
    Just u => Just (Evidence _ [<(u, x')])
    Nothing => Nothing
  where
  x' : Fin (S $ fst $ indexIsSuc x)
  x' = replace {p = Fin} (snd $ indexIsSuc x) x
  t' : Term sig (S $ fst $ indexIsSuc x)
  t' = replace {p = Term sig} (snd $ indexIsSuc x) t

export
amgu :
  DecEq (Exists sig.Operator) =>
  (t, u : Term sig n) ->
  Exists (AList sig n) ->
  Maybe (Exists (AList sig n))
amguAll :
  DecEq (Exists sig.Operator) =>
  (ts, us : Vect k (Term sig n)) ->
  Exists (AList sig n) ->
  Maybe (Exists (AList sig n))

amgu (Op op ts) (Op op' us) acc =
  case decEq {t = Exists sig.Operator} (Evidence _ op) (Evidence _ op') of
    Yes prf => amguAll ts (replace {p = \k => Vect k (Term sig n)} (sym $ cong fst prf) us) acc
    No neq => Nothing
amgu (Var x) (Var y) (Evidence _ [<]) = Just (flexFlex x y)
amgu (Var x) (Op op' us) (Evidence _ [<]) = flexRigid x (Op op' us)
amgu (Op op ts) (Var y) (Evidence _ [<]) = flexRigid y (Op op ts)
amgu t@(Var _) u@(Var _) (Evidence _ (sub :< (v, z))) =
  (:< (v, z)) <$> amgu ((v `for` z) <$> t) ((v `for` z) <$> u) (Evidence _ sub)
amgu t@(Var _) u@(Op _ _) (Evidence _ (sub :< (v, z))) =
  (:< (v, z)) <$> amgu ((v `for` z) <$> t) ((v `for` z) <$> u) (Evidence _ sub)
amgu t@(Op _ _) u@(Var _) (Evidence _ (sub :< (v, z))) =
  (:< (v, z)) <$> amgu ((v `for` z) <$> t) ((v `for` z) <$> u) (Evidence _ sub)

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

export
mgu : DecEq (Exists sig.Operator) => (t, u : Term sig n) -> Maybe (Exists (AList sig n))
mgu t u = amgu t u (Evidence _ [<])