{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
module Eval where

import Syntax
import qualified Data.Map.Strict as Map
import Data.Foldable
import Control.Exception
import Data.Typeable
import System.IO.Unsafe (unsafePerformIO)
import Data.IORef
import Systems
import Presyntax (Formula)
import qualified Presyntax as P
import Data.Maybe
import Debug.Trace
import GHC.Stack

iand :: Value -> Value -> Value
iand = \case
  VI1 -> id
  VI0 -> const VI0
  x -> \case
    VI0 -> VI0
    VI1 -> x
    y -> VIAnd x y

ior :: Value -> Value -> Value
ior = \case
  VI0 -> id
  VI1 -> const VI1
  x -> \case
    VI1 -> VI1
    VI0 -> x
    y -> VIOr x y

inot :: Value -> Value
inot VI1 = VI0
inot VI0 = VI1
inot (VIOr x y) = iand (inot x) (inot y)
inot (VIAnd x y) = ior (inot x) (inot y)
inot (VINot x) = x
inot x = VINot x

(@@) :: Value -> Value -> Value
VNe hd xs @@ vl = VNe hd (PApp vl:xs)
VLam _ _ k        @@ vl = k vl
VEqGlued a b      @@ vl = VEqGlued (a @@ vl) (b @@ vl)
VOfSub a phi u0 x @@ vl = x @@ vl
f @@ _  = error $ "can't apply argument to " ++ show f

proj1 :: Value -> Value
proj1 (VPair x _) = x
proj1 (VEqGlued x y) = VEqGlued (proj1 x) (proj1 y)
proj1 (VNe s xs) = VNe s (PProj1:xs)
proj1 (VOfSub (VSigma _ d _) phi u0 x) = VOfSub d phi (proj1 u0) (proj1 x)
proj1 x = error $ "can't proj1 " ++ show x

proj2 :: Value -> Value
proj2 (VPair _ y) = y
proj2 (VEqGlued x y) = VEqGlued (proj1 x) (proj1 y)
proj2 (VNe s xs) = VNe s (PProj2:xs)
proj2 (VOfSub (VSigma _ d r) phi u0 x) =
  VOfSub (r (proj1 x)) phi (proj2 u0) (proj2 x)
proj2 x = error $ "can't proj2 " ++ show x

pathp :: Env -> Value -> Value -> Value -> Value -> Value -> Value
pathp env p x y f@(VLine _a _x _y e) i =
  case reduceCube env i of
    Just P.Bot -> VEqGlued (e i) x
    Just P.Top -> VEqGlued (e i) y
    _          -> e i
pathp env p x y (VEqGlued e e') i = VEqGlued (pathp env p x y e i) (pathp env p x y e' i)
pathp env p x y (VNe hd sp) i =
  case reduceCube env i of
    Just P.Bot -> VEqGlued (VNe hd (PPathP p x y i:sp)) x
    Just P.Top -> VEqGlued (VNe hd (PPathP p x y i:sp)) y
    _ -> VNe hd (PPathP p x y i:sp)
pathp env p x y (VOfSub _ _ _ v) i = pathp env p x y v i

comp :: Env -> Value -> Formula -> Value -> Value -> Value
comp env a@(VLam ivar VI fam) phi u a0 = go (fam undefined) phi u a0 where
  i = VVar ivar
  stuck :: Value
  stuck = maybeAddEq $ VComp a (toValue phi) u a0

  maybeAddEq :: Value -> Value
  maybeAddEq =
    if phi == P.Top
      then flip VEqGlued (u @@ VI1)
      else id

  go :: HasCallStack => Value -> Formula -> Value -> Value -> Value
  go VPi{} phi u a0 =
    let
      dom x = let VPi _ d _ = fam x in d
      rng x = let VPi _ d _ = fam x in d

      ai1 = dom VI0
      y' i y   = fill env i (dom . inot . fam) P.Bot (VSystem emptySystem) y
      ybar i y = y' (inot i) y
     in VLam "x" ai1 \arg ->
       comp env
        (VLam ivar VI (\i -> rng (ybar i arg)))
        phi
        (VLam "i" VI \i -> mapVSystem (u @@ i) (@@ ybar i arg)) 
        (a0 @@ ybar VI0 arg)

  go VSigma{} phi u a0 =
    let
      dom x = let VSigma _ d _ = fam x in d
      rng x = let VSigma _ d _ = fam x in d

      a i = fill env i (dom . fam) phi (VLam "j" VI \v -> mapVSystem (u @@ v) proj1) (proj1 a0)
      c1 = comp env (VLam ivar VI (getd . fam))        phi (VLam "i" VI \v -> mapVSystem (u @@ v) proj1) (proj1 a0)
      c2 = comp env (VLam ivar VI (apr (a VI1) . fam)) phi (VLam "i" VI \v -> mapVSystem (u @@ v) proj2) (proj2 a0)

      getd (VSigma _ d _) = d
      apr x (VSigma _ _ r) = r x
     in VPair c1 c2

  go VPath{} phi p p0 =
    let
      ~(VPath ai1 u1 v1) = fam VI1
      ~(VPath ai0 u0 v0) = fam VI0
      getA (VPath a _ _) = a
      u' x = let ~(VPath _ u _) = fam x in u
      v' x = let ~(VPath _ _ v) = fam x in v
     in
       VLine (ai1 @@ VI1) u1 v1 \j ->
         let
           jc = reduceCube' env j
          in comp env (VLam ivar VI (getA . fam))
              (orFormula [phi, jc, notFormula jc])
              (VLam "j" VI \v ->
                let
                  VSystem (FMap sys) = p @@ v
                  sys' = fmap (flip (pathp env ai0 u0 v0) j) sys
               in mkVSystem $ Map.fromList [(phi, mapVSystem (p @@ v) (flip (pathp env ai0 u0 v0) j))
                                           , (notFormula jc, u' v), (jc, v' v)])
              (pathp env (ai0 @@ VI0) u0 v0 p0 j)

  go a P.Top u a0 = u @@ VI1
  go a phi u a0 = maybeAddEq stuck

comp env va phi u a0 =
  if phi == P.Top
    then VEqGlued (VComp va phi' u a0) (u @@ VI1)
    else VComp va phi' u a0
  where
    phi' = toValue phi

mkVSystem :: Map.Map Formula Value -> Value
mkVSystem mp
  | Just e <- Map.lookup P.Top mp = e
  | otherwise = VSystem $ FMap $ Map.filterWithKey f mp
  where
    f P.Bot _ = False
    f _ _ = True

reduceCube' :: Env -> Value -> Formula
reduceCube' env = fromJust . reduceCube env

mapVSystem :: Value -> (Value -> Value) -> Value
mapVSystem (VSystem ss) f = VSystem (mapSystem ss f)
mapVSystem x f = f x

evalSystem :: Env -> Map.Map Formula Term -> Value
evalSystem env face = mk . Map.mapMaybeWithKey go $ face where
  go :: Formula -> Term -> Maybe Value
  go face tm
    | VI0 <- toValue' env face = Nothing
    | otherwise = Just (eval env tm)

  differsFromEnv :: String -> Bool -> Bool
  differsFromEnv x True =
    case Map.lookup x (names env) of
      Just (VI, VI0) -> True
      _ -> False
  differsFromEnv x False =
    case Map.lookup x (names env) of
      Just (VI, VI1) -> True
      _ -> False

  mk x = case Map.toList x of
    [(_, x)] -> x
    _ -> mkVSystem x

eval :: Env -> Term -> Value
eval env = \case
  Var v ->
    case Map.lookup v (names env) of
      Just (_, vl) -> vl
      Nothing      -> error $ "variable not in scope: " ++ show v

  App f x -> eval env f @@ eval env x

  Lam s d b ->
    let d' = eval env d
     in VLam s d' \a -> eval env{ names = Map.insert s (d', a) (names env) } b

  Let s t b d ->
    let b' = eval env b
        t' = eval env t
     in eval env{ names = Map.insert s (t', b') (names env) } d

  Pi s d r  ->
    let d' = eval env d
     in VPi s d' \a -> eval env{ names = Map.insert s (d', a) (names env) } r

  Sigma s d r  ->
    let d' = eval env d
     in VSigma s d' \a -> eval env{ names = Map.insert s (d', a) (names env) } r

  Pair a b -> VPair (eval env a) (eval env b)
  Proj1 x  -> proj1 (eval env x)
  Proj2 y  -> proj2 (eval env y)

  Type     -> VType

  I   -> VI
  I0  -> VI0
  I1  -> VI1

  Path p x y      -> VPath (eval env p) (eval env x) (eval env y)
  Partial r a     -> VPartial (eval env r) (eval env a)

  PathI p x y s e -> VLine (eval env p) (eval env x) (eval env y) (\ a -> eval env{ names = Map.insert s (VI, a) (names env) } e)
  PathP p x y f i -> pathp env (eval env p) (eval env x) (eval env y) (eval env f) (eval env i)

  Sub p x y          -> VSub (eval env p) (eval env x) (eval env y)
  InclSub a phi u a0 -> VOfSub (eval env a) (eval env phi) (eval env u) (eval env a0)

  IAnd x y        -> iand (eval env x) (eval env y)
  IOr x y         -> ior (eval env x) (eval env y)
  INot x          -> inot (eval env x)

  Comp a phi u a0 ->
    case reduceCube env (eval env phi) of
      Just formula -> comp env (eval env a) formula (eval env u) (eval env a0)
      Nothing      -> VComp (eval env a) (eval env phi) (eval env u) (eval env a0)

  System fs       -> evalSystem env (getSystem fs)


data UnifyError
  = Mismatch       Value Value
  | NotPiType      Value
  | NotPartialType Formula Value
  | NotSigmaType   Value
  deriving (Show, Typeable, Exception)

unify :: Env -> Value -> Value -> IO ()
unify env (VEqGlued a b) c =
  unify env a c `catch` \e -> const (unify env b c) (e :: UnifyError)
unify env c (VEqGlued a b) =
  unify env c a `catch` \e -> const (unify env c b) (e :: UnifyError)

unify env (VLine a x y f) e = unify env (f (VVar "i")) (pathp env a x y e (VVar "i"))
unify env e (VLine a x y f) = unify env (f (VVar "i")) (pathp env a x y e (VVar "i"))

unify env (VPartial r b) (VPartial r' b') = do
  unify env b b'
  case sameCube env r r' of
    Just True -> pure ()
    _ -> unify env r r'

unify env (VPartial r b) x = do
  case sameCube env r VI1 of
    Just True -> pure ()
    _ -> unify env r VI1
  unify env b x

unify env x (VPartial r b) = do
  case sameCube env r VI1 of
    Just True -> pure ()
    _ -> unify env r VI1
  unify env x b

unify env (VSub a phi _u0) vl = unify env a vl

unify env u1 (VOfSub _a phi u0 a) = do
  case sameCube env phi VI1 of
    Just True -> unify env u1 u0
    _ -> unify env u1 a

unify env (VOfSub _a phi u0 a) u1 = do
  case sameCube env phi VI1 of
    Just True -> unify env u1 u0
    _ -> unify env u1 a

unify env vl1@(VNe x sp) vl2@(VNe y sp')
  | x == y    = traverse_ (uncurry unifySp) (zip sp sp')
  | otherwise = throwIO $ Mismatch vl1 vl2
  where
    unifySp (PApp x) (PApp y) = unify env x y
    unifySp (PPathP _a _x _y i) (PPathP _a' _x' _y' i') = unify env i i'
    unifySp PProj1 PProj1 = pure ()
    unifySp PProj2 PProj2 = pure ()

unify env (VLam x _ k) e = unify env (k (VVar x)) (e @@ VVar x)
unify env e (VLam x _ k) = unify env (e @@ VVar x) (k (VVar x))

unify env (VPi x d r) (VPi _ d' r') = do
  unify env d d'
  unify env (r (VVar x)) (r' (VVar x))

unify env (VSigma x d r) (VSigma _ d' r') = do
  unify env d d'
  unify env (r (VVar x)) (r' (VVar x))

unify env VType VType = pure ()

unify env VI  VI = pure ()

unify env (VPair a b) (VPair c d) = unify env a c *> unify env b d
unify env (VPath a x y) (VPath a' x' y') = unify env a a' *> unify env x x' *> unify env y y'

unify env (VSystem fs) vl
  | ((_, vl'):_) <- Map.toList (Map.filterWithKey (\f _ -> isTrue (toValue' env f)) (getSystem fs))
  = unify env vl' vl

unify env vl (VSystem fs)
  | ((_, vl'):_) <- Map.toList (Map.filterWithKey (\f _ -> isTrue (toValue' env f)) (getSystem fs))
  = unify env vl' vl

unify env x y =
  case sameCube env x y of
    Just True -> pure ()
    _ -> throwIO $ Mismatch x y

reduceCube :: Env -> Value -> Maybe Formula
reduceCube env x = fmap (toDNF . simplify) (toFormula x) where
  simplify :: Formula -> Formula
  simplify (P.Is0 x) =
    case Map.lookup x (names env) of
      Just (VI, VI0) -> P.Top
      Just (VI, VI1) -> P.Bot
      _ -> P.Is0 x
  simplify (P.Is1 x) =
    case Map.lookup x (names env) of
      Just (VI, VI1) -> P.Top
      Just (VI, VI0) -> P.Bot
      _ -> P.Is0 x
  simplify (P.And x y) = P.And (simplify x) (simplify y)
  simplify (P.Or x y) = P.Or (simplify x) (simplify y)
  simplify x = x

sameCube :: Env -> Value -> Value -> Maybe Bool
sameCube env x y =
  case (reduceCube env x, reduceCube env y) of
    (Just x, Just y) -> Just (x == y)
    _ -> Nothing

toFormula :: Value -> Maybe Formula
toFormula VI0                = Just P.Bot
toFormula VI1                = Just P.Top
toFormula (VNe x [])  = Just (P.Is1 x)
toFormula (VINot f)    = notFormula <$> toFormula f
toFormula (VIAnd x y) = do
  s <- toFormula y
  t <- toFormula x
  pure $ andFormula [s, t]
toFormula (VIOr x y) = do
  s <- toFormula y
  t <- toFormula x
  pure $ orFormula [s, t]
toFormula _ = Nothing

faceInEnv :: Env -> Face -> Bool
faceInEnv e f = Map.isSubmapOf (getFace f) (faceOfEnv (names e)) where
  faceOfEnv = Map.map (\(_, v) -> case v of { VI1 -> True; VEqGlued _ VI1 -> True; _ -> False }) . Map.filter (\(_, v) -> isI v)

  isI VI1 = True
  isI VI0 = True
  isI (VEqGlued _ x) = isI x
  isI _ = False

isPiType :: Value -> IO (String, Value, Value -> Value)
isPiType (VPi x d r) = pure (x, d, r)
isPiType x           = throwIO $ NotPiType x

isSigmaType :: Value -> IO (String, Value, Value -> Value)
isSigmaType (VSigma x d r) = pure (x, d, r)
isSigmaType x              = throwIO $ NotSigmaType x

isPiOrPathType :: Value -> IO (Either (String, Value, Value -> Value) (Value, Value, Value))
isPiOrPathType (VPi x d r)   = pure (Left (x, d, r))
isPiOrPathType (VPath x d r) = pure (Right (x, d, r))
isPiOrPathType x           = throwIO $ NotPiType x

isPartialType :: Formula -> Value -> IO (Formula, Value)
isPartialType f p@(VPartial x y) =
  case toFormula x of
    Just x -> pure (x, y)
    Nothing -> throwIO $ NotPartialType f p
isPartialType f x = throwIO $ NotPartialType f x

getVar :: IO Value
getVar =
  do
    n <- atomicModifyIORef ref \x -> (x + 1, x)
    pure (VVar (show n))
  where
    ref :: IORef Int
    ref = unsafePerformIO (newIORef 0)
    {-# NOINLINE ref #-}

fill :: Env
     -> Value
     -> (Value -> Value) -- (Γ i : I, A : Type)
     -> Formula          -- (phi : I)
     -> Value            -- (u : (i : I) -> Partial phi (A i))
     -> Value            -- (Sub (A i0) phi (u i0))
     -> Value            -- -> A i
fill env i a phi u a0 =
  comp env
    (VLam "j" VI \j -> a (i `iand` j))
    (phi `P.Or` ifc)
    (VLam "j" VI \j ->
      mkVSystem (Map.fromList [ (phi, uiand j)
                              , (notFormula ifc, a0) ]))
    a0
  where
    uiand j = u @@ (i `iand` j)
    ifc = fromJust (reduceCube env i)

toValue :: Formula -> Value
toValue P.Top = VI1
toValue P.Bot = VI0
toValue (P.And x y) = toValue x `iand` toValue y
toValue (P.Or x y) = toValue x `ior` toValue y
toValue (P.Is0 x) = inot (VVar x)
toValue (P.Is1 x) = VVar x

toValue' :: Env -> Formula -> Value
toValue' env P.Top = VI1
toValue' env P.Bot = VI0
toValue' env (P.And x y) = toValue x `iand` toValue y
toValue' env (P.Or x y) = toValue x `ior` toValue y
toValue' env (P.Is0 x) =
  case Map.lookup x (names env) of
    Just (VI, VI0) -> VI1
    Just (VI, VI1) -> VI0
    Just (VI, x) -> inot x
    _ -> error $ "type error in toValue'"
toValue' env (P.Is1 x) =
  case Map.lookup x (names env) of
    Just (VI, x) -> x
    _ -> error $ "type error in toValue'"

isTrue :: Value -> Bool
isTrue VI1 = True
isTrue _ = False