module Quox.Syntax.DimEq

import public Quox.Syntax.Var
import public Quox.Syntax.Dim
import public Quox.Syntax.Subst
import public Quox.Context
import Quox.Pretty
import Quox.Name
import Quox.Thin
import Quox.FinExtra

import Data.Maybe
import Data.Nat
import Data.DPair
import Data.Fun.Graph
import Data.SnocVect
import Decidable.Decidable
import Decidable.Equality
import Derive.Prelude

%language ElabReflection
%default total


public export
DimEq' : Nat -> Type
DimEq' = Context (Maybe . DimT)


public export
data DimEq : Nat -> Type where
  ZeroIsOne : DimEq d
  C         : (eqs : DimEq' d) -> DimEq d
%name DimEq eqs
%runElab deriveIndexed "DimEq" [Eq]

export
Show (DimEq d) where
  showPrec d ZeroIsOne = "ZeroIsOne"
  showPrec d (C eq')   = showCon d "C" $ showArg eq' @{ShowTelRelevant}


public export
consistent : DimEq d -> Bool
consistent ZeroIsOne = False
consistent (C eqs)   = True


public export
data IfConsistent : DimEq d -> Type -> Type where
  Nothing :      IfConsistent ZeroIsOne a
  Just    : a -> IfConsistent (C eqs)   a

export
Functor (IfConsistent eqs) where
  map f Nothing  = Nothing
  map f (Just x) = Just (f x)

export
Foldable (IfConsistent eqs) where
  foldr f z Nothing  = z
  foldr f z (Just x) = f x z

export
Traversable (IfConsistent eqs) where
  traverse f Nothing  = pure Nothing
  traverse f (Just x) = Just <$> f x

public export
ifConsistent : Applicative f => (eqs : DimEq d) -> f a -> f (IfConsistent eqs a)
ifConsistent ZeroIsOne act = pure Nothing
ifConsistent (C _)     act = Just <$> act

public export
toMaybe : IfConsistent eqs a -> Maybe a
toMaybe Nothing  = Nothing
toMaybe (Just x) = Just x


export
fromGround' : Context' DimConst d -> DimEq' d
fromGround' [<]        = [<]
fromGround' (ctx :< e) = fromGround' ctx :< Just (KT e noLoc)

export
fromGround : Context' DimConst d -> DimEq d
fromGround = C . fromGround'


public export %inline
zeroEq : DimEq 0
zeroEq = C [<]

public export %inline
new' : {d : Nat} -> DimEq' d
new' {d = 0}   = [<]
new' {d = S d} = new' :< Nothing

public export %inline
new : {d : Nat} -> DimEq d
new = C new'


public export %inline
get' : DimEq' d -> Fin d -> Maybe (DimT d)
get' = getWith $ \p, by => map (// by) p

public export %inline
getShift' : Shift len out -> DimEq' len -> Fin len -> Maybe (DimT out)
getShift' = getShiftWith $ \p, by => map (// by) p

public export %inline
get : {d : Nat} -> DimEq' d -> DimT d -> DimT d
get eqs p@(Th _ (K {})) = p
get eqs p@(Th i (B _))  = fromMaybe p $ get' eqs i.fin


public export %inline
equal : {d : Nat} -> DimEq d -> (p, q : DimT d) -> Bool
equal ZeroIsOne p q = True
equal (C eqs)   p q = get eqs p == get eqs q


infixl 7 :<?
export %inline
(:<?) : {d : Nat} -> DimEq d -> Maybe (DimT d) -> DimEq (S d)
ZeroIsOne :<? d = ZeroIsOne
C eqs     :<? d = C $ eqs :< map (get eqs) d


private %inline
isVar : {d : Nat} -> Fin d -> DimT d -> Bool
isVar i (Th j (B _))  = i == j.fin
isVar i (Th _ (K {})) = False

private %inline
ifVar : {d : Nat} -> Fin d -> DimT d -> Maybe (DimT d) -> Maybe (DimT d)
ifVar i p = map $ \q => if isVar i q then p else q

-- (using decEq instead of (==) because of the proofs below)
private %inline
checkConst : (e, f : DimConst) -> (eqs : Lazy (DimEq' d)) -> DimEq d
checkConst e f eqs = if isYes $ e `decEq` f then C eqs else ZeroIsOne


export
setConst : {d : Nat} -> Fin d -> DimConst -> Loc -> DimEq' d -> DimEq d
setConst FZ e loc (eqs :< Nothing) =
  C $ eqs :< Just (KT e loc)
setConst FZ e _ (eqs :< Just (Th _ (K f loc))) =
  checkConst e f $ eqs :< Just (KT f loc)
setConst FZ e loc (eqs :< Just (Th j (B _))) =
  setConst j.fin e loc eqs :<? Just (KT e loc)
setConst (FS i) e loc (eqs :< p) =
  setConst i e loc eqs :<? ifVar i (KT e loc) p

mutual
  private
  setVar' : {d : Nat} ->
            (i, j : Fin d) -> (0 _ : i `LT` j) -> Loc -> DimEq' d -> DimEq d
  setVar' FZ (FS i) LTZ loc (eqs :< Nothing) =
    C eqs :<? Just (BV i loc)
  setVar' FZ (FS i) LTZ loc (eqs :< Just (Th _ (K e eloc))) =
    setConst i e loc eqs :<? Just (KT e eloc)
  setVar' FZ (FS i) LTZ loc (eqs :< Just (Th j (B jloc))) =
    let j = j.fin in
    setVar i j loc jloc eqs :<? Just (if j > i then BV j jloc else BV i loc)
  setVar' (FS i) (FS j) (LTS lt) loc (eqs :< p) =
    setVar' i j lt loc eqs :<? ifVar i (BV j loc) p

  export %inline
  setVar : {d : Nat} -> (i, j : Fin d) -> Loc -> Loc -> DimEq' d -> DimEq d
  setVar i j li lj eqs with (compareP i j)
    setVar i j li lj eqs | IsLT lt = setVar' i j lt lj eqs
    setVar i i li lj eqs | IsEQ    = C eqs
    setVar i j li lj eqs | IsGT gt = setVar' j i gt li eqs


export %inline
set : {d : Nat} -> (p, q : DimT d) -> DimEq d -> DimEq d
set _ _ ZeroIsOne = ZeroIsOne
set (Th _ (K e _))  (Th _ (K f _))  (C eqs) = checkConst e f eqs
set (Th _ (K e el)) (Th j (B _))    (C eqs) = setConst j.fin e el eqs
set (Th i (B _))    (Th _ (K e el)) (C eqs) = setConst i.fin e el eqs
set (Th i (B il))   (Th j (B jl))   (C eqs) = setVar i.fin j.fin il jl eqs


public export %inline
Split : Nat -> Type
Split d = (DimEq' d, DSubst (S d) d)

export %inline
split1 : {d : Nat} -> DimConst -> Loc -> DimEq' (S d) -> Maybe (Split d)
split1 e loc eqs = case setConst 0 e loc eqs of
  ZeroIsOne    => Nothing
  C (eqs :< _) => Just (eqs, id (B loc) :< KT e loc)

export %inline
split : {d : Nat} -> Loc -> DimEq' (S d) -> List (Split d)
split loc eqs = toList (split1 Zero loc eqs) <+> toList (split1 One loc eqs)

export
splits' : {d : Nat} -> Loc -> DimEq' d -> List (DSubst d 0)
splits' _   [<]          = [[<]]
splits' loc eqs@(_ :< _) =
  [th . ph | (eqs', th) <- split loc eqs, ph <- splits' loc eqs']

||| the Loc is put into each of the DimConsts
export %inline
splits : {d : Nat} -> Loc -> DimEq d -> List (DSubst d 0)
splits _   ZeroIsOne = []
splits loc (C eqs)   = splits' loc eqs


-- private
-- 0 newGetShift : (d : Nat) -> (i : Fin d) -> (by : Shift d d') ->
--                 getShift' by (new' {d}) i = Nothing
-- newGetShift (S d) FZ     by = Refl
-- newGetShift (S d) (FS i) by = newGetShift d i (ssDown by)

-- export
-- 0 newGet' : (d : Nat) -> (i : Fin d) -> get' (new' {d}) i = Nothing
-- newGet' d i = newGetShift d i SZ

-- export
-- 0 newGet : (d : Nat) -> (p : Dim d) -> get (new' {d}) p = p
-- newGet d (K e _) = Refl
-- newGet d (B i _) = rewrite newGet' d i in Refl


-- export
-- 0 setSelf : (p : Dim d) -> (eqs : DimEq d) -> set p p eqs = eqs
-- setSelf p ZeroIsOne = Refl
-- setSelf (K Zero _) (C eqs) = Refl
-- setSelf (K One _)  (C eqs) = Refl
-- setSelf (B i _)    (C eqs) with (compareP i i) | (compare i.nat i.nat)
--   _ | IsLT lt | LT = absurd lt
--   _ | IsEQ    | EQ = Refl
--   _ | IsGT gt | GT = absurd gt


parameters {opts : LayoutOpts}
  private
  prettyDVars : {d : Nat} -> BContext d -> Eff Pretty (SnocList (Doc opts))
  prettyDVars = traverse prettyDBind . toSnocList'

  private
  prettyCst : {d : Nat} -> BContext d -> DimT d -> DimT d -> Eff Pretty (Doc opts)
  prettyCst dnames p q =
    hsep <$> sequence [prettyDim dnames p, cstD, prettyDim dnames q]

  private
  prettyCsts : {d : Nat} -> BContext d -> DimEq' d ->
               Eff Pretty (SnocList (Doc opts))
  prettyCsts [<] [<] = pure [<]
  prettyCsts dnames (eqs :< Nothing) = prettyCsts (tail dnames) eqs
  prettyCsts dnames (eqs :< Just q)  =
    [|prettyCsts (tail dnames) eqs :<
      prettyCst dnames (BV 0 noLoc) (weak 1 q)|]

  export
  prettyDimEq' : {d : Nat} -> BContext d -> DimEq' d -> Eff Pretty (Doc opts)
  prettyDimEq' dnames eqs = do
    vars <- prettyDVars dnames
    eqs  <- prettyCsts dnames eqs
    let prec = if length vars <= 1 && null eqs then Arg else Outer
    parensIfM prec $ fillSeparateTight !commaD $ toList vars ++ toList eqs

  export
  prettyDimEq : {d : Nat} -> BContext d -> DimEq d -> Eff Pretty (Doc opts)
  prettyDimEq dnames ZeroIsOne = do
    vars <- prettyDVars dnames
    cst  <- prettyCst [<] (KT Zero noLoc) (KT One noLoc)
    pure $ separateTight !commaD $ vars :< cst
  prettyDimEq dnames (C eqs) = prettyDimEq' dnames eqs


public export
wf' : {d : Nat} -> DimEq' d -> Bool
wf' [<]                         = True
wf' (eqs :< Nothing)            = wf' eqs
wf' (eqs :< Just (Th _ (K {}))) = wf' eqs
wf' (eqs :< Just (Th i (B _)))  = isNothing (get' eqs i.fin) && wf' eqs

public export
wf : {d : Nat} -> DimEq d -> Bool
wf ZeroIsOne = True
wf (C eqs)   = wf' eqs