{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE ViewPatterns #-}
module Elab.WiredIn where

import Control.Exception

import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Text as T
import Data.Map.Strict (Map)
import Data.Text (Text)
import Data.Typeable

import Elab.Eval

import GHC.Stack (HasCallStack)

import qualified Presyntax.Presyntax as P

import Syntax.Pretty (prettyTm)
import Syntax

import System.IO.Unsafe

wiType :: WiredIn -> NFType
wiType WiType     = VType
wiType WiPretype  = VTypeω

wiType WiInterval = VTypeω
wiType WiI0       = VI
wiType WiI1       = VI

wiType WiIAnd     = VI ~> VI ~> VI
wiType WiIOr      = VI ~> VI ~> VI
wiType WiINot     = VI ~> VI
wiType WiPathP    = dprod (VI ~> VType) \a -> a @@ VI0 ~> a @@ VI1 ~> VType

wiType WiIsOne    = VI ~> VTypeω
wiType WiItIsOne  = VIsOne VI1

wiType WiPartial  = VI ~> VType ~> VTypeω
wiType WiPartialP = dprod VI \x -> VPartial x VType ~> VTypeω

wiType WiSub      = dprod VType \a -> dprod VI \phi -> VPartial phi a ~> VTypeω
wiType WiInS      = forAll VType \a -> forAll VI \phi -> dprod a \u -> VSub a phi (fun (const u))
wiType WiOutS     = forAll VType \a -> forAll VI \phi -> forAll (VPartial phi a) \u -> VSub a phi u ~> a

wiType WiComp     = dprod' "A" (VI ~> VType) \a -> forAll VI \phi -> dprod (dprod VI \i -> VPartial phi (a @@ i)) \u -> VSub (a @@ VI0) phi (u @@ VI0) ~> a @@ VI1
-- (A : Type) {phi : I} (T : Partial phi Type) (e : PartialP phi (\o -> Equiv (T o) A)) -> Type
wiType WiGlue     = dprod' "A" VType \a -> forAll' "phi" VI \phi -> dprod' "T" (VPartial phi VType) \t -> VPartialP phi (fun \o -> equiv (t @@ o) a) ~> VType
-- {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
wiType WiGlueElem = forAll' "A" VType \a -> forAll' "phi" VI \phi -> forAll' "T" (VPartial phi VType) \ty -> forAll' "e" (VPartialP phi (fun \o -> equiv (ty @@ o) a)) \eqv ->
  dprod' "t" (VPartialP phi ty) \t -> VSub a phi (fun \o -> vProj1 (eqv @@ o) @@ (t @@ o)) ~> VGlueTy a phi ty eqv
-- {A : Type} {phi : I} {T : Partial phi Type} {e : PartialP phi (\o -> Equiv (T o) A)} -> Glue A phi T e -> A
wiType WiUnglue = forAll' "A" VType \a -> forAll' "phi" VI \phi -> forAll' "T" (VPartial phi VType) \ty -> forAll' "e" (VPartialP phi (fun \o -> equiv (ty @@ o) a)) \e -> VGlueTy a phi ty e ~> a

wiValue :: WiredIn -> Value
wiValue WiType     = VType
wiValue WiPretype  = VTypeω

wiValue WiInterval = VI
wiValue WiI0       = VI0
wiValue WiI1       = VI1

wiValue WiIAnd      = fun \x -> fun \y -> iand x y
wiValue WiIOr       = fun \x -> fun \y -> ior x y
wiValue WiINot      = fun inot
wiValue WiPathP     = fun \a -> fun \x -> fun \y -> VPath a x y

wiValue WiIsOne   = fun VIsOne
wiValue WiItIsOne = VItIsOne

wiValue WiPartial  = fun \phi -> fun \r -> VPartial phi r
wiValue WiPartialP = fun \phi -> fun \r -> VPartialP phi r
wiValue WiSub      = fun \a -> fun \phi -> fun \u -> VSub a phi u
wiValue WiInS      = forallI \a -> forallI \phi -> fun \u -> VInc a phi u
wiValue WiOutS     = forallI \a -> forallI \phi -> forallI \u -> fun \x -> outS a phi u x
wiValue WiComp     = fun \a -> forallI \phi -> fun \u -> fun \x -> comp a phi u x

wiValue WiGlue     = fun \a -> forallI \phi -> fun \t -> fun \e -> glueType a phi t e
wiValue WiGlueElem = forallI \a -> forallI \phi -> forallI \ty -> forallI \eqv -> fun \x -> fun \y -> glueElem a phi ty eqv x y
wiValue WiUnglue   = forallI \a -> forallI \phi -> forallI \ty -> forallI \eqv -> fun \x -> unglue a phi ty eqv x

(~>) :: Value -> Value -> Value
a ~> b = VPi P.Ex a (Closure (Bound "_" 0) (const b))
infixr 7 ~>

fun, line :: (Value -> Value) -> Value
fun k = VLam P.Ex $ Closure (Bound "x" 0) (k . force)
line k = VLam P.Ex $ Closure (Bound "i" 0) (k . force)

forallI :: (Value -> Value) -> Value
forallI k = VLam P.Im $ Closure (Bound "x" 0) (k . force)

dprod' :: String -> Value -> (Value -> Value) -> Value
dprod' t a b = VPi P.Ex a (Closure (Bound (T.pack t) 0) b)

dprod :: Value -> (Value -> Value) -> Value
dprod = dprod' "x"

exists' :: String -> Value -> (Value -> Value) -> Value
exists' s a b = VSigma a (Closure (Bound (T.pack s) 0) b)

exists :: Value -> (Value -> Value) -> Value
exists = exists' "x"

forAll' :: String -> Value -> (Value -> Value) -> Value
forAll' n a b = VPi P.Im a (Closure (Bound (T.pack n) 0) b)

forAll :: Value -> (Value -> Value) -> Value
forAll = forAll' "x"


wiredInNames :: Map Text WiredIn
wiredInNames = Map.fromList
  [ ("Pretype", WiPretype)
  , ("Type", WiType)
  , ("Interval", WiInterval)
  , ("i0", WiI0)
  , ("i1", WiI1)
  , ("iand", WiIAnd)
  , ("ior",  WiIOr)
  , ("inot", WiINot)
  , ("PathP", WiPathP)

  , ("IsOne", WiIsOne)
  , ("itIs1", WiItIsOne)

  , ("Partial", WiPartial)
  , ("PartialP", WiPartialP)
  , ("Sub", WiSub)
  , ("inS", WiInS)
  , ("outS", WiOutS)

  , ("comp", WiComp)
  , ("Glue", WiGlue)
  , ("glue", WiGlueElem)
  , ("unglue", WiUnglue)
  ]

newtype NoSuchPrimitive = NoSuchPrimitive { getUnknownPrim :: Text }
  deriving (Show, Typeable)
  deriving anyclass (Exception)

-- Interval operations

iand, ior :: Value -> Value -> Value
iand x = case force x of
  VI1 -> id
  VI0 -> const VI0
  VIAnd x y -> \z -> case force z of
    VI0 -> VI0
    VI1 -> VI1
    z -> iand x (iand y z)
  x -> \y -> case force y of
    VI0 -> VI0
    VI1 -> x
    y -> VIAnd x y

ior x = case force x of
  VI0 -> id
  VI1 -> const VI1
  VIOr x y -> \z -> case force z of
    VI1 -> VI1
    VI0 -> VIOr x y
    _ -> ior x (ior y z)
  x -> \y -> case force y of
    VI1 -> VI1
    VI0 -> x
    y -> VIOr x y

inot :: Value -> Value
inot x = case force x of
  VI0 -> VI1
  VI1 -> VI0
  VIOr x y -> VIAnd (inot x) (inot y)
  VIAnd x y -> VIOr (inot x) (inot y)
  VINot x -> x
  x -> VINot x

ielim :: Value -> Value -> Value -> Value -> NFEndp -> Value
ielim line left right (GluedVl h sp vl) i =
  GluedVl h (sp Seq.:|> PIElim line left right i) (ielim line left right vl i)
ielim line left right fn i =
  case force fn of
    VLine _ _ _ fun -> fun @@ i
    x -> case force i of
      VI1 -> right
      VI0 -> left
      _ -> case x of
        VNe n sp -> VNe n (sp Seq.:|> PIElim line left right i)
        VSystem map -> VSystem (fmap (flip (ielim line left right) i) map)
        VInc (VPath _ _ _) _ u -> ielim line left right u i
        VCase x xs -> VCase x (fmap (fmap (flip (ielim line left right) i)) xs)
        _ -> error $ "can't ielim " ++ show (prettyTm (quote fn))

outS :: NFSort -> NFEndp -> Value -> Value -> Value
outS _ (force -> VI1) u _  = u @@ VItIsOne

outS _ _phi _ (VInc _ _ x) = x
outS _ VI0 _  x            = x

outS a phi u  (GluedVl x sp vl) = GluedVl x (sp Seq.:|> POuc a phi u) (outS a phi u vl)
outS a phi u  (VNe x sp) = VNe x (sp Seq.:|> POuc a phi u)
outS _ _ _ v             = error $ "can't outS " ++ show (prettyTm (quote v))

-- Composition
comp :: NFLine -> NFEndp -> Value -> Value -> Value
comp _ VI1 u _ = u @@ VI1 @@ VItIsOne
comp a psi@phi u (compOutS (a @@ VI1) phi (u @@ VI1 @@ VItIsOne) -> a0) =
  case force $ a @@ VVar (Bound (T.pack "i") 0) of
    VPi{} ->
      let
        plic i = let VPi p _ _ = force (a @@ i)              in p
        dom  i = let VPi _ d _ = force (a @@ i)              in d
        rng  i = let VPi _ _ (Closure _ r) =  force (a @@ i) in r

        y' i y = fill (fun (dom . inot)) VI0 (fun \_ -> fun \_ -> VSystem mempty) (VInc (dom VI0) phi y) i
        ybar i y = y' (inot i) y
       in VLam (plic VI1) . Closure (Bound "x" 0) $ \arg ->
              comp (line \i -> rng i (ybar i arg))
                   phi
                   (system \i isone -> vApp (plic i) (u @@ i @@ isone) (ybar i arg))
                   (VInc (rng VI0 (ybar VI0 arg)) phi (vApp (plic VI0) a0 (ybar VI0 arg)))
    VSigma{} ->
      let
        dom i = let VSigma d _             = force (a @@ i) in d
        rng i = let VSigma _ (Closure _ r) = force (a @@ i) in r

        w i = fill (fun dom) phi (system \i isone -> vProj1 (u @@ i @@ isone)) (VInc (dom VI0) phi (vProj1 a0)) i
        c1 = comp (fun dom) phi (system \i isone -> vProj1 (u @@ i @@ isone)) (VInc (dom VI0) phi (vProj1 a0))
        c2 = comp (fun \x -> rng x (w x)) phi (system \i isone -> vProj2 (u @@ i @@ isone)) (VInc (rng VI0 (w VI0)) phi (vProj2 a0))
       in
         VPair c1 c2

    VPath{} ->
      let
        a' i = let VPath thea _ _ = force (a @@ i) in thea
        u' i = let VPath _ theu _ = force (a @@ i) in theu
        v' i = let VPath _ _ thev = force (a @@ i) in thev
      in
        VLine (a' VI1 @@ VI1) (u' VI1) (v' VI1) $ fun \j ->
          comp (fun \x -> a' x @@ x)
            (phi `ior` j `ior` inot j)
            (system \i isone -> mkVSystem (Map.fromList [ (phi, ielim (a' VI0) (u' VI0) (v' VI0) (u @@ i @@ isone) j)
                                                        , (j, v' i)
                                                        , (inot j, u' i)]))
            (VInc (a' VI0 @@ VI0 @@ j) phi (ielim (a' VI0 @@ VI0) (u' VI0) (v' VI0) a0 j))

    VGlueTy{} ->
      let
        b = u
        b0 = a0
        fam = a
      in
      let
        base i    = let VGlueTy base _ _ _   = force (fam @@ i) in base
        phi i     = let VGlueTy _ phi  _ _   = force (fam @@ i) in phi
        types i   = let VGlueTy _ _ types _  = force (fam @@ i) in types
        equivs i  = let VGlueTy _ _ _ equivs = force (fam @@ i) in equivs

        a i = fun \u -> unglue (base i) (phi i) (types i @@ u) (equivs i @@ u) (b @@ i @@ u)
        a0 = unglue (base VI0) (phi VI0) (types VI0) (equivs VI0) b0

        del = faceForall phi
        a1' = comp (line base) psi (line a) (VInc undefined undefined a0)
        t1' = comp (line types) psi (line (b @@)) (VInc undefined undefined b0)

        (omega_st, omega_t, omega_rep) = pres types base equivs psi (b @@) b0
        omega = outS omega_t psi omega_rep omega_st

        (t1alpha_st, t1a_t, t1a_rep) = opEquiv (base VI1) (types VI1 @@ VItIsOne) (equivs VI1 @@ VItIsOne) (del `ior` psi) (fun ts) (fun ps) a1'
        t1alpha = outS t1a_t (del `ior` psi) t1a_rep t1alpha_st

        (t1, alpha) = (vProj1 t1alpha, vProj2 t1alpha)

        ts isone = mkVSystem . Map.fromList $ [(del, t1'), (psi, (b @@ VI1 @@ isone))]
        ps _isone = mkVSystem . Map.fromList $ [(del, omega), (psi, VLine (line (const (base VI1))) a1' a1' (fun (const a1')))]

        a1 = comp
          (fun (const (base VI1 @@ VItIsOne)))
          (phi VI1 `ior` psi)
          (system \j _u -> mkVSystem (Map.fromList [ (phi VI1, ielim (base VI1) a1' (vProj1 (equivs VI1 @@ VItIsOne)) alpha j)
                                                   , (psi, a VI1)]))
          a1'
        b1 = glueElem (base VI1) (phi VI1) (types VI1) (equivs VI1) (fun (const t1)) a1
       in b1

    VType -> VGlueTy a0 phi (system \_ _ -> mkVSystem (Map.fromList [(phi, u @@ VI1 @@ VItIsOne)]))
                            (system \_ _  -> mkVSystem (Map.fromList [(phi, makeEquiv (\j -> (u @@ inot j)))]))

    VNe HData{} args ->
      case force a0 of
        VNe (HCon con_type con_name) con_args ->
          VNe (HCon con_type con_name) $ compConArgs (length args) (a @@) con_type con_args phi u
        _ -> VComp a phi u (VInc (a @@ VI0) phi a0)

    _ -> VComp a phi u (VInc (a @@ VI0) phi a0)

compConArgs :: Int -> (NFEndp -> Value) -> Value -> Seq.Seq Projection -> NFEndp -> Value -> Seq.Seq Projection
compConArgs total_args fam = go total_args where
  go _ _ Seq.Empty _ _ = Seq.Empty
  go nargs (VPi p dom (Closure _ rng)) (PApp p' y Seq.:<| xs) phi u
    | nargs > 0 = assert (p == p') $ go (nargs - 1) (rng (smuggle (fun (\i -> nthArg (total_args - nargs) (fam i))))) xs phi u
    | otherwise = assert (p == p') $
      let fill = makeFiller nargs dom phi u y
       in PApp p' (fill @@ VI1) Seq.:<| go (nargs - 1) (rng fill) xs phi u
  go _ _ _ _ _ = error $ "invalid constructor"

  nthArg i (VNe hd s) =
    case s Seq.!? i of
      Just (PApp _ t) -> t
      _ -> error $ "invalid " ++ show i ++ "th argument to data type " ++ show hd
  nthArg i (VSystem vs) = VSystem (fmap (nthArg i) vs)
  nthArg i xs = error $ "can't get " ++ show i ++ "th argument of " ++ show (prettyTm (quote xs))

  makeFiller nth (VNe (HData (Bound _ (negate -> 10))) args) phi u a0 =
    fun $ fill (makeDomain args) phi (system \i is1 -> nthArg nth (u @@ i @@ is1) ) a0
  makeFiller _ _ _ _ a0 = fun (const a0)

  makeDomain (PApp _ x Seq.:<| xs) = fun \i -> foldl (\t (~(PApp _ x)) -> t @@ (x @@ i)) (x @@ i) xs
  makeDomain _ = error "somebody smuggled something that smells"

  smuggle x = VNe (HData (Bound "__comp_con_tyarg" (negate 10))) (Seq.singleton (PApp P.Ex x))

compOutS :: NFSort -> NFEndp -> Value -> Value -> Value
compOutS a b c d = compOutS a b c (force d) where
  compOutS _ _hi _0 vl@VComp{}    = vl
  compOutS _ _hi _0 (VInc _ _ x) = x
  compOutS _ _   _  v = v

system :: (Value -> Value -> Value) -> Value
system k = fun \i -> fun \isone -> k i isone

fill :: NFLine -> NFEndp -> Value -> Value -> NFEndp -> Value
fill a phi u a0 j =
  comp (line \i -> a @@ (i `iand` j))
       (phi `ior` inot j)
       (fun \i -> fun \isone -> mkVSystem (Map.fromList [ (phi, u @@ (i `iand` j) @@ isone)
                                                        , (inot j, a0)]))
       a0

glueType :: NFSort -> NFEndp -> NFPartial -> NFPartial -> Value
glueType a phi tys eqvs = VGlueTy a phi tys eqvs

glueElem :: NFSort -> NFEndp -> NFPartial -> NFPartial -> NFPartial -> Value -> Value
glueElem _a (force -> VI1) _tys _eqvs t _vl = t @@ VItIsOne
glueElem a phi tys eqvs t vl = VGlue a phi tys eqvs t vl

unglue   :: NFSort -> NFEndp -> NFPartial -> NFPartial -> Value -> Value
unglue _a (force -> VI1) _tys eqvs x = vProj1 (eqvs @@ VItIsOne) @@ x
unglue _a _phi _tys _eqvs (force -> VGlue _ _ _ _ _ vl) = vl
unglue a phi tys eqvs (force -> VSystem fs) = VSystem (fmap (unglue a phi tys eqvs) fs)
unglue a phi tys eqvs vl = VUnglue a phi tys eqvs vl
-- Definition of equivalences

faceForall :: (NFEndp -> NFEndp) -> Value
faceForall phi = T.length (getNameText name) `seq` go (phi (VVar name)) where
  {-# NOINLINE name #-}
  name = unsafePerformIO newName

  go x@(VVar n)
    | n == name = VI0
    | otherwise = x
  go x@(VINot (VVar n))
    | n == name = VI0
    | otherwise = x
  go (VIAnd x y) = iand (go x) (go y)
  go (VIOr x y) = ior (go x) (go y)
  go (VINot x) = inot (go x)
  go vl = vl

isContr :: Value -> Value
isContr a = exists' "x" a \x -> dprod' "y" a \y -> VPath (line (const a)) x y

fiber :: NFSort -> NFSort -> Value -> Value -> Value
fiber a b f y = exists' "x" a \x -> VPath (line (const b)) (f @@ x) y

isEquiv :: NFSort -> NFSort -> Value -> Value
isEquiv a b f = dprod' "y" b \y -> isContr (fiber a b f y)

equiv :: NFSort -> NFSort -> Value
equiv a b = exists' "f" (a ~> b) \f -> isEquiv a b f

pres :: (NFEndp -> NFSort) -> (NFEndp -> NFSort) -> (NFEndp -> Value) -> NFEndp -> (NFEndp -> Value) -> Value -> (Value, NFSort, Value)
pres tyT tyA f phi t t0 = (VInc pathT phi (VLine (tyA VI1) c1 c2 (line path)), pathT, fun $ \u -> VLine (fun (const (tyA VI1))) c1 c2 (fun (const (f VI1 @@ (t VI1 @@ u))))) where
  pathT = VPath (fun (const (tyA VI1))) c1 c2
  c1 = comp (line tyA) phi (system \i u -> f i @@ (t i @@ u)) (VInc (tyA VI0) phi (f VI0 @@ t0))
  c2 = f VI1 @@ comp (line tyT) phi (system \i u -> t i @@ u) t0

  a0 = f VI0 @@ t0
  v = fill (fun tyT) phi (system \i u -> t i @@ u) t0

  path j = comp (fun tyA) (phi `ior` j) (system \i _ -> f i @@ (v i)) a0

opEquiv :: HasCallStack => Value -> Value -> Value -> NFEndp -> Value -> Value -> Value -> (Value, NFSort, Value)
opEquiv aT tT f phi t p a = (VInc ty phi v, ty, fun \u -> VPair (t @@ u) (p @@ u)) where
  fn = vProj1 f
  ty  = exists' "f" tT \x -> VPath (line (const aT)) a (fn @@ x)
  v   = contr ty (vProj2 f @@ a) phi (\u -> VPair (t @@ u) (p @@ u))

contr :: HasCallStack => Value -> Value -> NFEndp -> (Value -> Value) -> Value
contr a aC phi u =
  comp (line (const a))
    phi
    (system \i is1 -> ielim (line (const a)) (vProj1 aC) (u is1) (vProj2 aC @@ u is1) i)
    (vProj1 aC)

makeEquiv :: (NFEndp -> Value) -> Value
makeEquiv line = comp (fun \i -> equiv a (line i)) VI0 (system \_ _ -> VSystem mempty) (VPair idfun idisequiv) where
  a = line VI0
  idfun = fun id
-- idEquiv y = ((y, \i -> y), \u i -> (u.2 (inot i), \j -> u.2 (ior (inot i) j)))
  u_ty = exists' "y" a \x -> VPath (fun (const a)) x x
  idisequiv = fun \y -> VPair (id_fiber y) $ fun \u ->
    VLine u_ty (id_fiber y) u $ fun \i -> VPair (ielim (fun (const a)) y y (vProj2 u) i) $
      VLine (fun (const a)) y (vProj1 u) $ fun \j ->
        ielim (fun (const a)) y y (vProj2 u) (iand i j)

  id_fiber y = VPair y (VLine a y y (fun (const y)))