{-# LANGUAGE GADTs, KindSignatures, OverloadedStrings, EmptyDataDecls, StandaloneDeriving, TypeSynonymInstances, TypeFamilies, MultiParamTypeClasses #-}
module Normal where

import Prelude hiding (length,elem,foldl)
import Basics
import Display
import Data.Foldable
import Control.Arrow (first, second)
import Data.Sequence hiding (zip,replicate,reverse)
import Options

data No
data Ne
data Va

type NF = Term No
type Neutral = Term Ne
type Variable = Term Va
type NF' = (NF, NF) -- value, type.

data Term n :: * where
     Neu :: Neutral -> NF
     Var :: Variable -> Neutral
     
     Star :: Sort -> NF     
     
     -- FIXME: only "column" / relevance should be here.
     Pi  :: Relevance -> Ident -> NF -> NF -> NF
     Lam :: Relevance -> Ident -> NF -> NF -> NF 
     App :: Relevance -> Neutral -> NF -> Neutral -- The sort is that of the argument.
     
     Sigma :: Relevance -> Ident -> NF -> NF -> NF
     Pair  :: Relevance -> Ident -> NF -> NF -> NF  -- Pair does not bind any variable.
     Proj  :: Relevance -> -- ^ Sort of the argument
              Neutral -> Bool -> -- ^ True for 1st projection; False for 2nd.
              Irr String -> Neutral 
     
     
     OfParam :: Ident -> NF -> Neutral

     Destr :: Relevance -> Variable -> Variable -- argument: level destroyed
     Param :: Relevance -> Variable -> Variable 
     V :: Sort -> Int -> Variable -- shift, deBruijn 
     Hole :: String -> Variable

type Subst = [NF]

deriving instance Eq (Term n)

var :: Int -> NF
var x = Neu $ var' x

var' x = Var $ V (Sort 0 0) x


-- | Hereditary substitution
subst0 :: NF -> NF -> NF
subst0 u = subst (u:map (var) [0..])  


subst :: Subst -> Term n -> NF
subst f t = case t of
  Neu x -> s x
  Var x -> s x
  
  Star x -> Star x
  
  Lam o i ty bo -> Lam o i (s ty) (s' bo)
  (Pair o i x y) -> Pair o i (s x) (s y)
  Pi o i a b -> Pi o i (s a) (s' b)
  Sigma o i a b -> Sigma o i (s a) (s' b)
  (App o a b) -> app o (s a) (s b)
  (Proj o x k f) -> proj o (s x) k f

  OfParam i x -> Neu (OfParam i (s x))
  
  Hole x -> Neu $ Var $ Hole x
  V s x -> shift s (f !! x)
  Param o x -> param o (s x)
  Destr f x -> destroy f (s x)
 where s' = subst (var 0 : map wk f)
       s  = subst f

-- Double renaming substitution
-- 1st component: regular; 2nd component: param
subst' :: [(Variable,Variable)] -> Term n -> Term n
subst' f t = case t of
  Neu x -> Neu (s x)
  Var x -> Var (s x)
  
  Star x -> Star x
  
  Lam o i ty bo -> Lam o i (s ty) (s' o bo)
  (Pair o i x y) -> Pair o i (s x) (s y)
  Pi o i a b -> Pi o i (s a) (s' o b)
  Sigma o i a b -> Sigma o i (s a) (s' o  b)
  (App o a b) -> App o (s a) (s b)
  (Proj o x k f) -> Proj o (s x) k f

  OfParam i x -> OfParam i (s x)
  
  Hole x -> Hole x
  V s x -> shift s (fst $ f !! x)
  Param o (V s x) -> shift s (snd $ f !! x)
  Param o x -> Param o (s x)
  Destr f x -> Destr f (s x)
 where s' o = subst' (p o f)
       s  = subst' f
       p o xs = (V zero 0, Param o $ V zero 0) : map (both $ wkv 1) xs
       
both f (x,y) = (f x, f y)

shift' :: Int -> Sort -> Term n -> Term n
shift' n d@(Sort _ r) t = case t of
  Neu x -> Neu $ s x
  Var x -> Var (s x)
  
  Star o -> Star (o + d)
  
  Lam o i ty bo -> Lam (o +. d) i (s ty) (s' bo)
  (Pair o i x y) -> Pair (o +. d) i (s x) (s y)
  Pi o i a b -> Pi (o +. d) i (s a) (s' b)
  Sigma o i a b -> Sigma (o +. d) i (s a) (s' b)
  (App o a b) -> App (o +. d) (s a) (s b)
  (Proj o x k f) -> Proj (o +. d) (s x) k f
  
  OfParam i x -> OfParam i (s x)

  Hole x -> Hole x
  Param o x -> Param (o +. d) (s x)
  Destr f x -> Destr (f + r) (s x)
  V s x | x < n  -> V s x
        | x >= n -> V (s + d) x
 where s = shift' n d
       s' = shift' (1 + n) d

shift = shift' 0

-----------------------------
-- Hereditary operations
  
app :: Relevance -> NF -> NF -> NF 
app _ (Lam _ i _ bo) u = subst0 u bo
app o (Neu n)      u = Neu (App o n u)

proj :: Relevance -> NF -> Bool -> Irr String -> NF
proj _ (Pair _ _ x y) True f = x
proj _ (Pair _ _ x y) False f = y
proj o (Neu x) k f = Neu (Proj o x k f)


wkn :: Int -> NF -> NF
wkn n = subst (map var [n..])
wk = wkn 1
str = subst0 (Neu $ Var $ Hole "str: oops!")

wkv :: Int -> Variable -> Variable
wkv n (Destr d x) = Destr d (wkv n x)
wkv n (Param o x) = Param o (wkv n x)
wkv n (V s x) = V s (x + n)
wkv n (Hole x) = Hole x

param :: Relevance -> NF -> NF
param o t = transNF 0 t o


-----------------------------------
-- Display

dec xs = [ x - 1 | x <- xs, x > 0]

freeVars :: Term n -> [Int]
freeVars (Var x) = freeVars x
freeVars (Neu x) = freeVars x
freeVars (Pi _ _ a b) = freeVars a <> (dec $ freeVars b)
freeVars (Sigma _ _ a b) = freeVars a <> (dec $ freeVars b)
freeVars (V _ x) = [x]
freeVars (App _ a b) = freeVars a <> freeVars b
freeVars (Lam _ _ ty b) = freeVars ty <> (dec $ freeVars b)
freeVars (Star _) = mempty
freeVars (Hole _) = mempty
freeVars (Pair _ _ x y) = freeVars x <> freeVars y
freeVars (Proj _ x _ _) = freeVars x
freeVars (Param _ x) = freeVars x
freeVars (OfParam _ x) = freeVars x
freeVars (Destr _ x) = freeVars x

iOccursIn :: Int -> Term n -> Bool
iOccursIn x t = x `elem` (freeVars t)

prettyRel' = prettySortNam

cPrint :: Int -> DisplayContext -> Term n -> Doc
cPrint p ii (Var x) = cPrint p ii x
cPrint p ii (Neu x) = cPrint p ii x
cPrint p ii (Destr i x) = cPrint p ii x <> "%" <> pretty i
cPrint p ii (Param o x) = cPrint p ii x <> sss (pretty o) <> "!"
cPrint p ii (OfParam i x) = pretty i
                             -- "⌊" <> cPrint (-1) ii x <> "⌋"
cPrint p ii (Hole x) = text x
cPrint p ii (Star i) = pretty i
cPrint p ii (V s k) 
  | k < 0 || k >= length ii  = text "<deBrujn index" <+> pretty k <+> text "out of range>"
  | otherwise = pretty (ii `index` k)  <> shft
  where shft | s == Sort 0 0 = mempty
             | otherwise = "⇧" <> prettySortNam s
cPrint p ii (Proj o x k (Irr f))     = cPrint p ii x <> sss (pretty o) <> (if k then "#" else "/") <> text f
cPrint p ii t@(App _ _ _)     = let (fct,args) = nestedApp t in 
                                 parensIf (p > 3) (cPrint 3 ii fct <+> sep [ sss (pretty o <> "· ") <> cPrint 4 ii a | (o,a) <- args]) 
cPrint p ii t@(Pi _ _ _ _)    = parensIf (p > 1) (printBinders "→" ii mempty $ nestedPis t)
cPrint p ii t@(Sigma _ _ _ _) = parensIf (p > 1) (printBinders "×" ii mempty $ nestedSigmas t)
cPrint p ii (t@(Lam _ _ _ _))   = parensIf (p > 1) (nestedLams ii mempty t)
cPrint p ii (Pair _ name x y) = parensIf (p > (-1)) (sep [pretty name <+> text "=" <+> cPrint 0 ii x <> comma,
                                                          cPrint (-1) ii y])

nestedPis  :: NF -> ([(Ident,Bool,NF,Relevance)], NF)
nestedPis (Pi o i a b) = (first ([(i,0 `iOccursIn` b,a,o)] ++)) (nestedPis b)
nestedPis x = ([],x)

nestedSigmas  :: NF -> ([(Ident,Bool,NF,Relevance)], NF)
nestedSigmas (Sigma o i a b) = (first ([(i,0 `iOccursIn` b,a,o)] ++)) (nestedSigmas b)
nestedSigmas x = ([],x)

printBinders :: Doc -> DisplayContext -> Seq Doc -> ([(Ident,Bool,NF,Relevance)], NF) -> Doc
printBinders sep ii xs (((i,occurs,a,o):pis),b) = printBinders sep (i <| ii) (xs |> (printBind' ii i occurs a o <+> sep)) (pis,b)
printBinders _ ii xs ([],b)                 = sep $ toList $ (xs |> cPrint 1 ii b) 


nestedLams :: DisplayContext -> Seq Doc -> Term n -> Doc
nestedLams ii xs (Lam o x ty c) = nestedLams (x <| ii) (xs |> parens (sss (pretty o) <> pretty x <+> ":" <+> cPrint 0 ii ty)) c
nestedLams ii xs t         = (text "\\ " <> (sep $ toList $ (xs |> "->")) <+> nest 3 (cPrint 0 ii t))

printBind' ii name occurs d o = case not (isDummyId name) ||  occurs of
                  True -> parens (sss (pretty o) <> pretty name <+> text ":" <+> cPrint 0 ii d)
                  False -> cPrint 2 ii d

nestedApp :: Neutral -> (Neutral,[(Relevance, NF)])
nestedApp (App o f a) = (second (++ [(o,a)])) (nestedApp f)
nestedApp t = (t,[])

sss x = if showSorts options then x else mempty

prettyTerm = cPrint (-100)


instance Pretty (Term n) where
    pretty = prettyTerm mempty


mv :: Int -> Int -> Int
mv d x | x < d     = (arity + 1) * x + idx
       | otherwise = (x - d) + (arity + 1) * d
                     -- x + arity * d

mv' :: Int -> Int -> (Variable, Variable)
mv' d x | x < d     = let v = (arity + 1) * x 
                      in (V zero $ v + idx, V zero v)
        | otherwise = let v = V zero $ (x - d) + (arity + 1) * d 
                      in (v, Hole "does not appear!")
                          -- Param evil v)

evil = Sort 0 666


renam :: Int -> Int -> NF -> NF
renam d idx = subst [var $ mv d $ x | x <- [0..]] . shift' d oneRel

renam' d = subst' (map (mv' d) [0..])

re :: Ident -> Ident
re (Irr (Identifier (pos ,x)))  = (Irr (Identifier (pos,x++"₁")))

arity, idx :: Int
arity = 1
idx = 1


-- | Transform a term to its relational interpretation
-- NOTE: the level of the sort is incorrect! In fact only the rel. ever matters. (Destroy)
transV :: Int -> Variable -> Relevance -> Variable

transV d (V o x) o' | x < d =                V o $ (arity + 1) * x
                    | otherwise = Param o' $ V o $ (x - d) + (arity + 1) * d
transV d (Param o' x) o = Param o' $ transV d x o
transV d (Destr n x) o = destroy n (transV d x o)
transV d (Hole s) _ = Hole (s ++ "!")

transNe :: Int -> Neutral -> Relevance -> NF
transNe d (Var v) o = Neu $ Var $ transV d v o
transNe d (App o f a) o' = app o (app (next o) (transNe d f o') (renam d idx a)) (transNF d a o) 
transNe d (Proj o x k f) o' = proj o (transNe d x o) k f
transNe d (OfParam i t) o = app o (renam' d t) (renam d idx (Neu $ OfParam i t))

transNF :: Int -> NF -> Relevance -> NF
transNF d (Neu v) o = transNe d v o
transNF d (Lam o i ty bo) o' = transBind d Lam o i ty (transNF (d+1) bo o')
transNF d (Pair o i x y) o' = Pair o i (transNF d x o) (transNF d y o') 
transNF d ty@(Star  _) o = trans' d ty o
transNF d ty@(Pi    _ _ _ _) o = trans' d  ty o
transNF d ty@(Sigma _ _ _ _) o = trans' d ty o

trans' d ty o = Lam (next o) (synthId "z₁") (renam d idx ty) (zerInRel d ty o)

-- | Build a relation witnessing x ∈ ⟦ty⟧. (where 'x' is Bound 0 in 'ty'.)

-- In the translated context, 'z1', ... 'zn' are bound, but not
-- 'zR'. (we are going to bind it soon).  However, 'inTrans' assumes
-- it has "full" translated context.  So we weaken 'ty' (putting 'z' in scope) and apply the
-- translation as normal.  But the translation is well behaved, so it
-- does not use 'zR'. We substitute it with nothing when the job is
-- done.
zerInRel d ty o = str $ inTrans (d + 1) (wk ty) o (var 0)

-- | Build a relation z ∈ ⟦ty⟧.  z is a term that, after renaming,
-- gives the vector of terms member of the relation.  Note that
-- 'trans' is never applied to 'z', therefore 'zR' never occurs in the result.

inTrans :: Int -> NF -> Relevance -- ^ sort of the 1st argument
           -> NF -> NF
inTrans d (Star  s)       o z = (Pi (next o) dummyId (renam d idx z) (Star s))
inTrans d (Pi    o i a b) o' z = transBind d Pi o i a (inTrans (d + 1) b o' (app o (wk z) (var 0)))
inTrans d (Sigma o i a b) o' z = Sigma o i (inTrans d a o (proj o z True f)) 
                             (subst (var 0:wk (renam d idx (proj o z True f)):map var [1..]) $
                              inTrans (1 + d) b o' (proj o (wk z) False f)) -- TEST: is depth ok?
 where (Irr (Identifier (_,nam))) = i
       f = Irr nam
inTrans d t o z = app (next o) (transNF d t o) (renam d idx z) 


-- | Translate a binding (x : A) into (x₁ : A₁) (⟦x⟧ : ⟦A⟧ x₁)
transBind :: Int -> (Relevance -> Ident -> NF -> NF -> NF) -> Relevance -> Ident -> NF -> NF -> NF
transBind d binder o i a rest = binder (next o) (re i) (renam d idx a) $
                                binder o            i  (zerInRel d a o) $ 
                                rest


-- Invariant: the whole term is not destroyed.
destroy :: Relevance -> Term n -> Term n
destroy d t = case t of
  Var x -> Var $ pr x
  Neu x -> Neu $ pr x

  V o x -> V o x
  Hole x -> Hole x
  Destr d' t -> destroy (min d d') t -- coalesce
  Param r x | r+1 == d -> x
            | otherwise -> Destr d $ Param r x 

  (Star o) -> Star o
  (Pi o i a b)    -> mb Pi    o i a b 
  (Sigma o i a b) -> mb Sigma o i a b 
  (Lam o i ty bo) -> mb Lam   o i ty bo  
  (Pair o i a b) 
      | isDestroyed o -> pr b
      | otherwise -> Pair o i (pr a) (pr b) 
  (App o a b)  -> case isDestroyed o of
                   True -> pr a
                   False -> App o (pr a) (pr b)
  (Proj o x k f) -> case isDestroyed o of
    True -> pr x -- result of the projection is not destroyed (by
                  -- assumpt.) but the whole pair would be -> we must
                  -- keep the 1st component.
    False -> Proj o (pr x) k f
  (OfParam n x) -> OfParam (modId (++ "%" ++ show d) n) $ pr x
  
 where 
   isDestroyed o = d `destroys` o
   mb :: (Relevance -> Ident -> NF -> NF -> NF) -> Relevance -> Ident -> NF -> NF -> NF
   mb binder o i a b = case isDestroyed o of
                             True -> str (pr b)
                             False -> binder o i (pr a) (pr b)
   pr x = destroy d x


r `destroys` r' = r' >= r