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

import Syntax
import Data.Map.Strict (Map)
import Data.Text (Text)
import qualified Data.Map.Strict as Map
import Control.Exception
import Data.Typeable
import qualified Presyntax.Presyntax as P
import Elab.Eval
import qualified Data.Sequence as Seq
import qualified Data.Text as T

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 WiIsOne1   = forAll VI \i -> forAll VI \j -> VIsOne i ~> VIsOne (ior i j)
wiType WiIsOne2   = forAll VI \i -> forAll VI \j -> VIsOne j ~> VIsOne (ior i j)

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 (VLam P.Ex (Closure "x" (const u)))
wiType WiOutS     = forAll VType \a -> forAll VI \phi -> forAll (VPartial phi a) \u -> VSub a phi u ~> a

wiType WiComp     = dprod (VI ~> VType) \a -> forAll VI \phi -> dprod (dprod VI \i -> VPartial phi (a @@ i)) \u -> VSub (a @@ VI0) phi (u @@ VI0) ~> a @@ VI1

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 WiIsOne1  = forallI \_ -> forallI \_ -> fun VIsOne1
wiValue WiIsOne2  = forallI \_ -> forallI \_ -> fun VIsOne2

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     = forAll (VI ~> VType) \a -> forAll VI \phi -> dprod (dprod VI \i -> VPartial phi (a @@ i)) \u -> VSub (a @@ VI0) phi (u @@ VI0) ~> a @@ VI1
wiValue WiComp     = fun \a -> forallI \phi -> fun \u -> fun \x -> comp a phi u x

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

fun :: (Value -> Value) -> Value
fun k = VLam P.Ex $ Closure "x" (k . force)

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

dprod :: Value -> (Value -> Value) -> Value
dprod a b = VPi P.Ex a (Closure "x" b)

forAll :: Value -> (Value -> Value) -> Value
forAll a b = VPi P.Im a (Closure "x" b)

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)
  , ("isOneL", WiIsOne1)
  , ("isOneR", WiIsOne2)

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

  , ("comp", WiComp)
  ]

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

-- Interval operations

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

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

inot :: Value -> Value
inot = \case
  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 fn i =
  case force fn of
    VLine _ _ _ fun -> fun @@ i
    VNe n sp -> VNe n (sp Seq.:|> PIElim _line _left _right i)
    _ -> error $ "can't ielim " ++ show fn

outS :: NFSort -> NFEndp -> Value -> Value -> Value
outS _ (force -> VI1) u _  = u @@ VItIsOne
outS _ _phi _ (VInc _ _ x) = x 
outS a phi u  (VNe x sp)   = VNe x (sp Seq.:|> POuc a phi u)
outS _ _ _ v               = error $ "can't outS " ++ show v

-- Composition
comp :: NFLine -> NFEndp -> Value -> Value -> Value
comp _ VI1 u _ = u @@ VI1 @@ VItIsOne
comp a phi u (VInc _ _ a0) =
  case a @@ VNe (HVar (Bound (T.pack "x"))) Seq.empty of
    VPi{} -> 
      let
        plic i = let VPi p _ _ = a @@ i in p
        dom  i = let VPi _ d _ = a @@ i in d
        rng  i = let VPi _ _ (Closure _ r) = 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 "x" $ \arg ->
              comp (fun \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 _ = a @@ i in d
        rng i = let VSigma _ (Closure _ r) = 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 (($ w VI1) . rng)) phi (system \i isone -> vProj1 (u @@ i @@ isone)) (VInc (rng VI0 (dom VI0)) phi (vProj2 a0))
       in
         VPair c1 c2
    VPath{} ->
      let
        a' i = let VPath a _ _ = a @@ i in a
        u' i = let VPath _ u _ = a @@ i in u
        v' i = let VPath _ _ v = a @@ i in v
      in
        VLine (a' VI1 @@ VI1) (u' VI1) (v' VI1) $ fun \j ->
          comp (fun a')
            (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))

    _ -> VComp a phi u a0
comp a phi u a0 = VComp a phi u a0

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 (fun \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