{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
module Syntax where

import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map.Strict (Map)
import Systems
import qualified Data.Map.Strict as Map
import Text.Show (showListWith)
import Presyntax (Formula)

type Space = Term

data Term
  = Var String
  | App Term   Term
  | Lam String Term Term
  | Let String Term Term Term

  | Pi  String Term Term
  | Type

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

  | Path  Space Term Term
  | PathI Space Term Term String Term
  | PathP Space Term Term Term Term

  | Sigma String Term Term
  | Pair Term Term
  | Proj1 Term
  | Proj2 Term

  | System (System Term)
  | Partial Term Term
  | Comp Term Term Term Term
  | Sub Term Term Term
  | InclSub Term Term Term Term
  deriving (Eq, Ord)

instance Show Term where
  showsPrec p =
    \case
      Var s -> showString s
      System fs -> showListWith showPE (Map.toList (getSystem fs))
      -- ew
      App (App (App (Lam _ _ (Lam _ _ (Lam _ _ Path{}))) a) x) y -> showsPrec p (Path a x y)
      App (App (App (Lam _ _ (Lam _ _ (Lam _ _ Sub{}))) a) x) y -> showsPrec p (Sub a x y)
      App (App (Lam _ _ (Lam _ _ Partial{})) phi) r -> showsPrec p (Partial phi r)

      App (App (App (App (Lam _ _ (Lam _ _ (Lam _ _ (Lam _ _ Comp{})))) a) phi) u) a0 ->
        showsPrec p (Comp a phi u a0)

      App f x -> showParen (p >= app_prec) $
        showsPrec fun_prec f
        . showChar ' '
        . showsPrec app_prec x

      Lam s t b ->
        let
          getLams (Lam s _ b) =
            let (as, b') = getLams b
             in (s:as, b')
          getLams (PathI _a _x _y s b) =
            let (as, b') = getLams b
             in (("(" ++ s ++ " : I)"):as, b')
          getLams t = ([], t)
          (args, bd) = getLams (Lam s t b)
        in showParen (p >= fun_prec) $
            showString ("λ " ++ unwords args ++ " -> ")
            . shows bd

      Let s t d b -> showParen (p > fun_prec) $
        showString "let\n  "
        . showString s
        . showString " : "
        . shows t
        . showString " = "
        . shows d
        . showString " in "
        . shows b

      Pi "_" d r ->
        showParen (p >= domain_prec) $
          showsPrec domain_prec d
          . showString " -> "
          . shows r

      Pi v d r -> showParen (p >= domain_prec) $
        let
          showBinder (Pi "_" d r) =
            showsPrec domain_prec d
            . showString " -> "
            . shows r
          showBinder (Pi n d r) =
            let 
              arr = case r of
                Pi n _ _ | n /= "_" -> " "
                _ -> " -> "
             in
               showParen True (showString n . showString " : " . shows d)
               . showString arr
               . showBinder r
          showBinder t = shows t
         in showBinder (Pi v d r)

      Type -> showString "Type"
      I    -> showChar 'I'
      I0   -> showString "i0"
      I1   -> showString "i1"

      IAnd i j -> showParen (p >= and_prec) $
        showsPrec or_prec i
        . showString " && "
        . showsPrec or_prec j

      IOr i j -> showParen (p >= or_prec) $
        showsPrec app_prec i
        . showString " || "
        . showsPrec app_prec j

      INot s -> showChar '~' . showsPrec p s

      Path a x y -> showsPrec p (App (App (App (Var "Path") a) x) y)
      Sub a x y -> showsPrec p (App (App (App (Var "Sub") a) x) y)
      Partial r a -> showsPrec p (App (App (Var "Partial") r) a)
      Comp a phi u a0 -> showsPrec p (foldl App (Var "comp") [a, phi, u, a0])
      InclSub _a _phi _u0 a0 -> showsPrec p a0

      PathI a x y s b -> showParen (p >= fun_prec) $
        showString ("λ " ++ s ++ " -> ")
        . shows b
      PathP _a _x _y f i -> showsPrec p (App f i)

      Pair a b -> showParen True $
        shows a
        . showString ", "
        . shows b

      Proj1 b -> showsPrec p b . showString ".1"
      Proj2 b -> showsPrec p b . showString ".1"
      Sigma v d r ->
        showParen (p >= app_prec) $
          showParen True (showString v . showString " : " . shows d)
          . showString " × "
          . shows r
    where
      app_prec    = 6
      domain_prec = 5
      and_prec    = 4
      or_prec     = 3
      fun_prec    = 1

showPE :: (Formula, Term) -> String -> String
showPE (f, t) = shows f . showString " -> " . shows t

data Value
  = VNe   String [Proj]
  | VLam  String Value (Value -> Value)
  | VPi   String Value (Value -> Value)
  | VType
  | VI | VI0 | VI1
  | VEqGlued Value Value -- e which is def. eq. to e'
  | VPair Value Value
  | VSigma String Value (Value -> Value)
  | VLine Value Value Value (Value -> Value)
    -- (λ i → ...) : Path A x y
    -- order: A x y k
  | VSystem (System Value)
  | VOfSub Value Value Value Value

  | VIAnd Value Value
  | VIOr Value Value
  | VINot Value

  | VPath Value Value Value
  | VSub Value Value Value
  | VPartial Value Value
  | VComp Value Value Value Value

data Proj
  = PApp Value
  | PPathP Value Value Value Value
    -- a x y i
  | PProj1
  | PProj2

pattern VVar :: String -> Value
pattern VVar x = VNe x []

quote :: Value -> Term
quote = go mempty where
  go :: Set String -> Value -> Term
  go scope (VNe hd spine) = foldl (goSpine scope) (Var hd) (reverse spine)
  go scope (VLam s a k) =
    let n = rename s scope
     in Lam n (go scope a) (go (Set.insert n scope) (k (VVar n)))
  go scope (VPi s d r) =
    let n = rename s scope
     in Pi n (go scope d) (go (Set.insert n scope) (r (VVar n)))
  go scope (VSigma s d r) =
    let n = rename s scope
     in Sigma n (go scope d) (go (Set.insert n scope) (r (VVar n)))

  go scope VType = Type
  go scope VI0   = I0
  go scope VI1   = I1
  go scope VI    = I

  go scope (VIAnd x y) = IAnd (go scope x) (go scope y)
  go scope (VIOr x y) = IOr (go scope x) (go scope y)
  go scope (VINot x) = INot (go scope x)

  go scope (VPath a x y) = Path (go scope a) (go scope x) (go scope y)
  go scope (VSub a x y) = Sub (go scope a) (go scope x) (go scope y)
  go scope (VPartial r a) = Partial (go scope r) (go scope a)
  go scope (VComp a b c d) = Comp (go scope a) (go scope b) (go scope c) (go scope d)

  go scope (VEqGlued e _) = go scope e
  go scope (VPair a b) = Pair (go scope a) (go scope b)
  go scope (VLine a x y k) =
    let n = rename "i" scope
     in PathI (go scope a) (go scope x) (go scope y) n (go (Set.insert n scope) (k (VVar n)))
  go scope (VSystem (FMap fs)) = System (FMap (fmap (go scope) fs))
  go scope (VOfSub _ _ _ x) = go scope x

  goSpine :: Set String -> Term -> Proj -> Term
  goSpine scope t (PApp x) = App t (go scope x)
  goSpine scope t (PPathP a x y i) = PathP (go scope a) (go scope x) (go scope y) t (go scope i)
  goSpine scope t PProj1 = Proj1 t
  goSpine scope t PProj2 = Proj2 t

rename :: String -> Set String -> String
rename x s
  | x == "_"         = x
  | x `Set.member` s = rename (x ++ "'") s
  | otherwise        = x

instance Show Value where
  showsPrec p = showsPrec p . quote

data Env =
  Env { names :: Map String (Value, Value)
      }
  deriving (Show)