{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE DeriveDataTypeable #-}
module Syntax where

import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Map.Strict (Map)
import Data.Sequence (Seq)
import Data.Function (on)
import Data.IORef (IORef)
import Data.Text (Text)
import Data.Set (Set)
import Data.Data

import Presyntax.Presyntax (Plicity(..), Posn)

data WiredIn
  = WiType
  | WiPretype
  | WiInterval
  | WiI0
  | WiI1
  | WiIAnd
  | WiIOr
  | WiINot
  | WiPathP

  | WiPartial  -- (φ : I) -> Type           -> Typeω
  | WiPartialP -- (φ : I) -> Partial r Type -> Typeω
  | WiPOr      -- (A : Type) (φ ψ : I) -> Partial φ A -> Partial ψ A -> Partial (ior φ ψ) A

  | WiSub      -- (A : Type) (φ : I) -> Partial φ A -> Typeω
  | WiInS      -- {A : Type} {φ : I} (u : A) -> Sub A φ (λ x -> u)
  | WiOutS     -- {A : Type} {φ : I} {u : Partial φ A} -> Sub A φ u -> A

  | WiComp  --    {A : I -> Type} {phi : I}
            -- -> ((i : I) -> Partial phi (A i)
            -- -> (A i0)[phi -> u i0] -> A i1

  | WiGlue      -- (A : Type) {phi : I} (T : Partial phi Type) (e : PartialP phi (\o -> Equiv (T o) A)) -> Type
  | WiGlueElem  -- {A : Type} {phi : I} {T : Partial phi Type} {e : PartialP phi (\o -> Equiv (T o) A)} -> (t : PartialP phi T) -> Sub A phi (\o -> e o (t o)) -> Glue A phi T e
  | WiUnglue    -- {A : Type} {phi : I} {T : Partial phi Type} {e : PartialP phi (\o -> Equiv (T o) A)} -> Glue A phi T e -> A

  | WiLineToEquiv -- {A : I -> Type} -> Equiv (P i0) (P i1)

  -- Two-level
  | WiSEq   -- Eq_s   : {A : Pretype} (x y : A) -> Pretype
  | WiSRefl -- refl_s : {A : Pretype} {x : A} -> EqS {A} x x
  | WiSK    -- K_s    : {A : Pretype} {x : A} (P : x = x -> Pretype) -> P refl -> (p : x = x) -> P p
  | WiSJ    -- J_s   : {A : Pretype} {x : A} (P : (y : A) -> x = y -> Pretype) -> P x refl -> {y : A} -> (p : x = y) -> P y p
  deriving (Eq, Show, Ord)

data Term
  = Ref  Name
  | Con  Name
  | PCon Term Name
  | Data Bool Name
  | App  Plicity Term Term
  | Lam  Plicity Name Term
  | Pi   Plicity Name Term Term
  | Let  [(Name, Term, Term)] Term
  | Meta MV
  | Type
  | Typeω

  | Sigma Name Term Term
  | Pair Term Term
  | Proj1 Term
  | Proj2 Term

  | I
  | I0 | I1
  | IAnd Term Term
  | IOr  Term Term
  | INot Term

  | PathP     Term Term Term
  -- ^ PathP : (A : I -> Type) -> A i0 -> A i1 -> Type
  | IElim     Term Term Term Term Term
  -- ^ IElim : {A : I -> Type} {x : A i0} {y : A i1} (p : PathP A x y) (i : I) -> A i
  | PathIntro Term Term Term Term
  -- ^ PathIntro : (A : I -> Type) (f : (i : I) -> A i) -> PathP A (f i0) (f i1)
  --   ~~~~~~~~~ not printed at all

  | Partial  Term Term
  | PartialP Term Term

  | System (Map Term Term)

  | Sub Term Term Term
  | Inc Term Term Term
  | Ouc Term Term Term Term

  | Comp Term Term Term Term
  | HComp Term Term Term Term

  | GlueTy Term Term Term Term
  | Glue Term Term Term Term Term Term
  | Unglue Term Term Term Term Term

  | Case Term Term [(Term, Int, Term)]

  | EqS  Term Term Term
  | Refl Term Term
  | AxK  Term Term Term Term Term
  | AxJ  Term Term Term Term Term Term
  deriving (Eq, Show, Ord, Data)

data MV =
  MV { mvName :: Text
     , mvCell :: {-# UNPACK #-} !(IORef (Maybe Value))
     , mvType :: NFType
     , mvSpan :: Maybe (Text, Posn, Posn)
     , mvInteraction :: Bool
     , mvContext :: Map Name NFType
     }

instance Eq MV where
  (==) = (==) `on` mvName
instance Ord MV where
  (<=) = (<=) `on` mvName
instance Show MV where
  show x = show (mvName x, mvSpan x)

instance Data MV where
  toConstr _   = error "toConstr"
  gunfold _ _  = error "gunfold"
  dataTypeOf _ = mkNoRepType "MV"


data Name
  = Bound   {getNameText :: Text, getNameNum :: !Int}
  | Defined {getNameText :: Text, getNameNum :: !Int}
  | ConName {getNameText :: Text, getNameNum :: !Int, conSkip :: !Int, conArity :: !Int}
  deriving (Show, Data)

instance Eq Name where
  x == y = getNameText x == getNameText y && getNameNum x == getNameNum y

instance Ord Name where
  compare x y = getNameText x `compare` getNameText y <> getNameNum x `compare` getNameNum y

type NFType = Value
type NFEndp = Value
type NFLine = Value
type NFSort = Value
type NFPartial = Value

data Value
  = VNe    Head (Seq Projection)
  | VLam   Plicity Closure
  | VPi    Plicity Value Closure
  | VSigma Value Closure
  | VPair  Value Value

  | GluedVl Head (Seq Projection) Value

  | VType | VTypeω

  | VI
  | VI0 | VI1
  | VIAnd NFEndp NFEndp
  | VIOr  NFEndp NFEndp
  | VINot NFEndp

  | VPath NFLine Value Value
  | VLine NFLine Value Value Value

  | VPartial  NFEndp Value
  | VPartialP NFEndp Value
  | VSystem (Map Value Value)

  | VSub NFSort NFEndp Value
  | VInc NFSort NFEndp Value

  | VComp NFLine NFEndp Value Value
  | VHComp NFSort NFEndp Value Value

  | VGlueTy NFSort NFEndp NFPartial NFPartial
  | VGlue   NFSort NFEndp NFPartial NFPartial NFPartial Value
  | VUnglue NFSort NFEndp NFPartial NFPartial Value

  | VCase (Map.Map Name (NFType, Value)) Value Value [(Term, Int, Term)]

  | VEqStrict   NFSort Value Value
  | VReflStrict NFSort Value
  deriving (Eq, Show, Ord)

pattern VVar :: Name -> Value
pattern VVar x = VNe (HVar x) Seq.Empty

quoteWith :: Bool -> Set Name -> Value -> Term
quoteWith short names (VNe h sp) = foldl goSpine (goHead h) sp where
  goHead (HVar v)    = Ref v
  goHead (HMeta m)   = Meta m
  goHead (HCon _ v)  = Con v
  goHead (HPCon sys _ v) =
    case sys of
      VSystem f ->
        case Map.lookup VI1 f of
          Just vl -> constantly (length sp) vl
          _ -> PCon (quote sys) v
      VLam{} -> PCon (quote sys) v
      s -> constantly (length sp) s
  goHead (HData x v) = Data x v

  goSpine t (PApp p v) = App p t (quoteWith short names v)
  goSpine t (PIElim l x y i) = IElim (quote l) (quote x) (quote y) t (quote i)
  goSpine t (PK l x y i) = AxK (quote l) (quote x) (quote y) (quote i) t
  goSpine t (PJ l x y i f) = AxJ (quote l) (quote x) (quote y) (quote i) (quote f) t
  goSpine t PProj1 = Proj1 t
  goSpine t PProj2 = Proj2 t
  goSpine t (POuc a phi u) = Ouc (quote a) (quote phi) (quote u) t

  constantly 0 n = quoteWith short names n
  constantly k x = Lam Ex (Bound (T.pack "x") (negate 1)) $ constantly (k - 1) x

quoteWith short names (GluedVl _ Seq.Empty x) = quoteWith short names x

quoteWith short names (GluedVl h sp (VLam p (Closure n k))) =
  quoteWith short names (VLam p (Closure n (\a -> GluedVl h (sp Seq.:|> PApp p a) (k a))))

quoteWith short names (GluedVl h sp (VLine ty x y (VLam p (Closure n k)))) =
  quoteWith short names (VLine ty x y (VLam p (Closure n (\a -> GluedVl h (sp Seq.:|> PIElim ty x y a) (k a)))))

quoteWith short names (GluedVl h sp vl)
  | GluedVl _ _ inner <- vl = quoteWith short names (GluedVl h sp inner)
  | short || alwaysShort vl = quoteWith short names vl
  | _ Seq.:|> PIElim _ x y i <- sp =
    case i of
      VI0 -> quoteWith short names x
      VI1 -> quoteWith short names y
      _ -> quoteWith short names (VNe h sp)
  | otherwise = quoteWith short names (VNe h sp)

quoteWith short names (VLam p (Closure n k)) =
  let n' = refresh Nothing names n
   in Lam p n' (quoteWith short (Set.insert n' names) (k (VVar n')))

quoteWith short names (VPi p d (Closure n k)) =
  let n' = refresh (Just d) names n
   in Pi p n' (quoteWith short names d) (quoteWith short (Set.insert n' names) (k (VVar n')))

quoteWith short names (VSigma d (Closure n k)) =
  let n' = refresh (Just d) names n
   in Sigma n' (quoteWith short names d) (quoteWith short (Set.insert n' names) (k (VVar n')))

quoteWith short names (VPair a b) = Pair (quoteWith short names a) (quoteWith short names b)
quoteWith _ _ VType = Type
quoteWith _ _ VTypeω = Typeω

quoteWith _ _ VI     = I
quoteWith _ _ VI0    = I0
quoteWith _ _ VI1    = I1
quoteWith short names (VIAnd x y) = IAnd (quoteWith short names x) (quoteWith short names y)
quoteWith short names (VIOr x y) = IOr (quoteWith short names x) (quoteWith short names y)
quoteWith short names (VINot x) = INot (quoteWith short names x)

quoteWith short names (VPath line x y) = PathP (quoteWith short names line) (quoteWith short names x) (quoteWith short names y)
quoteWith short names (VLine p x y f)  = PathIntro (quoteWith short names p) (quoteWith short names x) (quoteWith short names y) (quoteWith short names f)

quoteWith short names (VPartial x y)  = Partial (quoteWith short names x) (quoteWith short names y)
quoteWith short names (VPartialP x y) = PartialP (quoteWith short names x) (quoteWith short names y)
quoteWith short names (VSystem fs) = System (Map.fromList (map (\(x, y) -> (quoteWith short names x, quoteWith short names y)) (Map.toList fs)))
quoteWith short names (VSub a b c) = Sub (quoteWith short names a) (quoteWith short names b) (quoteWith short names c)
quoteWith short names (VInc a b c) = Inc (quoteWith short names a) (quoteWith short names b) (quoteWith short names c)
quoteWith short names (VComp a phi u a0) = Comp (quoteWith short names a) (quoteWith short names phi) (quoteWith short names u) (quoteWith short names a0)
quoteWith short names (VHComp a phi u a0) = HComp (quoteWith short names a) (quoteWith short names phi) (quoteWith short names u) (quoteWith short names a0)

quoteWith short names (VGlueTy a phi t e)    = GlueTy (quoteWith short names a) (quoteWith short names phi) (quoteWith short names t) (quoteWith short names e)
quoteWith short names (VGlue a phi ty e t x) = Glue (quoteWith short names a) (quoteWith short names phi) (quoteWith short names ty) (quoteWith short names e) (quoteWith short names t) (quoteWith short names x)
quoteWith short names (VUnglue a phi ty e x) = Unglue (quoteWith short names a) (quoteWith short names phi) (quoteWith short names ty) (quoteWith short names e) (quoteWith short names x)

quoteWith short names (VCase _ rng v xs) = Case (quoteWith short names rng) (quoteWith short names v) xs

quoteWith short names (VEqStrict a x y) = EqS (quoteWith short names a) (quoteWith short names x) (quoteWith short names y)
quoteWith short names (VReflStrict a x) = Syntax.Refl (quoteWith short names a) (quoteWith short names x)

alwaysShort :: Value -> Bool
alwaysShort (VNe HCon{} _) = True
alwaysShort (VNe HPCon{} _) = True
alwaysShort VVar{} = True
alwaysShort (VLine _ _ _ v) = alwaysShort v
alwaysShort (VLam _ (Closure n k)) = alwaysShort (k (VVar n))
alwaysShort VHComp{} = True
alwaysShort _ = False

refresh :: Maybe Value -> Set Name -> Name -> Name
refresh (Just VI) n _ = refresh Nothing n (Bound (T.pack "phi") 0)
refresh x s n
  | T.singleton '_' == getNameText n = n
  | n `Set.notMember` s = n
  | otherwise = refresh x s (Bound (getNameText n <> T.singleton '\'') 0)

quote :: Value -> Term
quote = quoteWith True mempty

data Closure
  = Closure
    { clArgName :: Name
    , clCont    :: Value -> Value
    }

instance Show Closure where
  show (Closure n c) = "Closure \\" ++ show n ++ " -> " ++ show (c (VVar n))

instance Eq Closure where
  Closure _ k == Closure _ k' =
    k (VVar (Bound (T.pack "_") 0)) == k' (VVar (Bound (T.pack "_") 0))

instance Ord Closure where
  Closure _ k <= Closure _ k' =
    k (VVar (Bound (T.pack "_") 0)) <= k' (VVar (Bound (T.pack "_") 0))

data Head
  = HVar Name
  | HCon Value Name
  | HPCon Value Value Name
  | HMeta MV
  | HData Bool Name
  deriving (Ord, Show)

instance Eq Head where
  HVar x == HVar y = x == y
  HCon _ x == HCon _ y = x == y
  HPCon _ _ x == HPCon _ _ y = x == y
  HMeta x == HMeta y = x == y
  HData x y == HData x' y' = x == x' && y == y'
  _ == _ = False

data Projection
  = PApp   Plicity Value
  | PIElim Value Value Value NFEndp
  | PProj1
  | PProj2
  | POuc  NFSort NFEndp Value
  | PK    NFSort Value NFSort Value
  | PJ    NFSort Value NFSort Value Value
  deriving (Eq, Show, Ord)

data Boundary = Boundary { getBoundaryNames :: [Name], getBoundaryMap :: Value }
  deriving (Eq, Show, Ord)

unPi :: Value -> ([(Plicity, Value)], Value)
unPi (VPi p d (Closure n k)) =
  let (a, x) = unPi (k (VVar n))
   in ((p, d):a, x)
unPi x = ([], x)