{-# 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     
     
     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 (only needed for
                           -- the 1st projection: 2nd projection does
                           -- not change relevance)
              Neutral -> Bool -> -- ^ True for 1st projection; False for 2nd.
              Irr String -> Neutral 
     
     
     OfParam :: Ident -> NF -> Neutral

     Destr :: Int -> Variable -> Variable -- argument: depth where destruction occurs.
     Param :: Variable -> Variable 
     V :: Sort -> Int -> Variable -- shift, deBruijn 
     Hole :: String -> Variable

etaExpand :: Relevance -> Neutral -> NF -> NF
etaExpand o' v (Pi    o i a b) = Lam  o i a (etaExpand o' (App o (wkne 1 v) 
                                                           $ etaExpand o (var' 0) a) b)
etaExpand o' v (Sigma o i a b) = Pair o i   (etaExpand o  (Proj o' v True  (Irr $ idString i)) a) 
                                            (etaExpand o' (Proj o' v False (Irr $ idString i)) b)
etaExpand o' v _ = Neu v



type Subst = [NF]

deriving instance Eq (Term n)
deriving instance Show (Term n)

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

var'' = V (Sort 0)

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


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

showShift (Sort l) = replicate l '^' 

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))
  
  Destr d x -> destroy d (s x)
  Hole x -> Neu $ Var $ Hole x
  V s x -> shift s (f !! x)
  Param x -> param (s x)
 where s' = subst (var 0 : map wk f)
       s  = subst f


-- Double renaming substitution
-- 1st component: regular; 2nd component: param
subst2 :: [(NF,NF)] -> Term n -> NF
subst2 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 (fst $ f !! x)
  Param (V s x) -> shift s (snd $ f !! x)
  Destr d x -> destroy d (s x)
  Param x -> param (s x)
 where s' = subst2 ((var 0, param $ var 0) : map (both wk) f)
       s  = subst2 f


subst2d :: Int -> (NF,NF) -> Term n -> NF
subst2d d u = subst2 $ [(var i,param $ var i) | i <- [0..d-1]] ++ u : 
                       [(var i,param $ var i) | i <- [d..]]


{-
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 (V s x) -> shift s (snd $ f !! x)
  Param x -> Param (s x)
 where s' o = subst' (p f)
       s  = subst' f
       p xs = (V zero 0, Param $ 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 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 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 -> OfParam (modId (++showShift d) i) (s x)

  Hole x -> Hole x
  Param x -> Param (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..])

wkdn :: Int -> Int -> NF -> NF
wkdn d n = subst (map var [0..d-1] ++ map var [d+n..])

wk = wkn 1
str = subst0 (Neu $ Var $ Hole "str: oops!")

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

wkne :: Int -> Neutral -> Neutral
wkne n (Var x) = Var (wkv n x)
wkne n (App o a b) = App o (wkne n a) (wkn n b)
wkne n (Proj o a k f) = Proj o (wkne n a) k f
wkne n (OfParam i a) = OfParam i (wkn n a)


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


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

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

freeVars :: Term n -> [Int]
freeVars (Var x) = freeVars x
freeVars (Destr _ 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

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

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 (Param x) = cPrint p ii x <> "!"
cPrint p ii (Destr d x) = cPrint p ii x <> "%" <> pretty d
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 o@(Sort l) k) 
  | k < 0 || k >= length ii  = text "<deBrujn index" <+> pretty k <+> text "out of range>"
  | otherwise = pretty (ii `index` k)  <> shft
  where shft = text (showShift o)
cPrint p ii (Proj o x k (Irr f))     = cPrint p ii x <> sss (pretty o) <> (if k then "." <> text f else "/")
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 arrow ii mempty $ nestedPis t)
cPrint p ii t@(Sigma _ _ _ _) = parensIf (p > 1) (printBinders cross 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])

cross Ir = "⤬" -- ⚔⤬⤫⨯
cross Re = "×" -- ×⨯

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 :: (Relevance -> 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 <+> sss (pretty o) <> sep o)) (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 <+> colon o <+> 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 (pretty name <+> colon o <+> 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)

-- paramShift = if collapseRelevance options then zero else oneRel
              -- TODO: have this as an argument to
              -- Param. Alternatively, add a construct to collapse
              -- levels.

next :: Relevance -> Relevance
next _ = Ir -- (+ (sortRelevance paramShift))


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

-- 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
transV    :: Int -> Variable -> Variable

transV  d (V o x) = Param $ V o x
transV  d (Param x) = Param $ transV d x
transV  d (Hole s) = Hole (s ++ "!")

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

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

trans'  d ty = Lam Ir (synthId "z") ty (zerInRel d ty)

-- | Build the relation x ∈ ⟦ty⟧. (where 'x' is 0; but not bound in 'ty'.)
zerInRel d ty = inTrans (d + 1) (wk ty) (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 -> NF -> NF
inTrans d (Star  s)       z = (Pi Ir dummyId z (Star s))
inTrans d (Pi    o i a b) z = transBind d Pi o i a (inTrans (d + 1) b (app o (wk z) (var 0)))
inTrans d (Sigma o i a b) z = Sigma o (re i) (inTrans d a (proj o z True f)) $
                              subst2d 1 (wk $ proj o z True f, var 0) $ wk $
                              inTrans (1 + d) b (proj o (wk z) False f) -- TEST: is depth ok?
 where (Irr (Identifier (_,nam))) = i
       f = Irr nam
inTrans d t z = app Ir (transNF d t) 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 Re i a rest = binder Ir i a $ 
                                 binder Re (re i) (zerInRel d a) $ 
                                 subst2d 2 (var 1,var 0) $ wkn 2 rest

transBind d binder Ir i a rest = binder Ir i a rest

-- Invariant: the whole term is not destroyed.
destroy :: Int -> 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 x | d == 0 -> x
          | otherwise -> Destr d $ Param 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' o a) (pr b) 
  (App o a b)  -> case isDestroyed o of
                   True -> pr a
                   False -> App o (pr a) (pr' o 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 -- FIXME: hmmm, here we should probably use pr' (symmetry)
  (OfParam n x) -> OfParam (modId (++ "%" ++ show d) n) $ pr x
  
 where 
   isDestroyed o = d == 0 && o == Ir
   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' o  a) (pr b)
   pr x = destroy d x
   pr' Ir x = destroy (d-1) x
   pr' Re x = pr x