{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveFunctor, DeriveTraversable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
module M where

import Bound
import Data.Deriving
import Data.Functor.Classes
import Data.Map (Map)
import Data.Set (Set)
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Void
import Control.Monad
import Data.String

type Id = String

freshen :: Foldable t => Id -> t Id -> Id
freshen v l | v `notElem` l = v
            | otherwise = freshen (v ++ "'") l

-- Hmmm... When we use this below, should we be taking care to make sure that
-- the new variable is also fresh for the context and other nearby objects,
-- rather than just for the thing being opened...?
openFresh :: (Monad f, Foldable f) => Id -> Scope b f Id -> (Id, f Id)
openFresh v scope = (new, instantiate1 (pure new) scope)
  where new = freshen v scope

-- Possibly-[o]pen [ty]pe term.
data OTy fr
  -- "VoidTy" to avoid confusion with Haskell-level Void type
  = TyVar fr | VoidTy | Unit
  | OTy fr :-> OTy fr | OTy fr :+ OTy fr | OTy fr :* OTy fr
  | Mu (Scope () OTy fr) | Nu (Scope () OTy fr)
  deriving (Functor, Foldable, Traversable)
infixr 5 :->
infixr 6 :+
infixr 7 :*
type Ty = OTy Void
type TyOp = Scope () OTy Void

makeBound ''OTy
deriveShow1 ''OTy
deriveEq1 ''OTy
instance Show fr => Show (OTy fr) where showsPrec = showsPrec1
instance Eq fr => Eq (OTy fr) where (==) = eq1
instance IsString (OTy Id) where fromString = TyVar . fromString

ty0 :: OTy (Var () a)
ty0 = TyVar (B ())

tyOp :: Id -> OTy Id -> Maybe TyOp
tyOp v = closed . abstract1 v

mu, nu :: Id -> OTy Id -> OTy Id
mu v = Mu . abstract1 v
nu v = Nu . abstract1 v

-- We include some extra type annotations on the map and rec terms. These
-- aren't present in the syntax presented in the book, but otherwise we'd have
-- non-unique types for terms, which is complexity that I don't want to deal
-- with. On top of that, the sum case for map's dynamics needs extra type
-- annotations on the terms in order to simplify without figuring out the type
-- of a subterm! This doesn't seem to be accounted for in the book, either,
-- because the type annotations on the sum constructors are elided in the
-- relevant rule (14.3c).
-- This ends up requiring us to also add an extra type annotation to the gen
-- term in order to account in the dynamics for the previously added
-- annotations.
data Exp fr
  = ExpVar fr | Absurd Ty (Exp fr) | Empty
  | Lam Ty (Scope () Exp fr) | Exp fr :@ Exp fr
  | Inl Ty Ty (Exp fr) | Inr Ty Ty (Exp fr)
  | Case (Exp fr) (Scope () Exp fr) (Scope () Exp fr)
  | Pair (Exp fr) (Exp fr) | Fst (Exp fr) | Snd (Exp fr)
  | MapPos TyOp Ty Ty (Scope () Exp fr) (Exp fr)
  | Fold TyOp (Exp fr)
  | Rec TyOp Ty (Scope () Exp fr) (Exp fr)
  | Unfold TyOp (Exp fr)
  | Gen TyOp Ty (Scope () Exp fr) (Exp fr)
  deriving (Functor, Foldable, Traversable)
infixl 8 :@
type CExp = Exp Void -- Closed Exp

exp0 :: Exp (Var () a)
exp0 = ExpVar (B ())

lam :: Ty -> Id -> Exp Id -> Exp Id
lam t v = Lam t . abstract1 v

case' :: Exp Id -> Id -> Exp Id -> Id -> Exp Id -> Exp Id
case' s v1 b1 v2 b2 = Case s (abstract1 v1 b1) (abstract1 v2 b2)

mapPos :: TyOp -> Ty -> Ty -> Id -> Exp Id -> Exp Id -> Exp Id
mapPos o ti1 ti2 v b e = MapPos o ti1 ti2 (abstract1 v b) e

rec, gen :: TyOp -> Ty -> Id -> Exp Id -> Exp Id -> Exp Id
rec o t v step e = Rec o t (abstract1 v step) e
gen o t v step e = Gen o t (abstract1 v step) e

makeBound ''Exp
deriveShow1 ''Exp
deriveEq1 ''Exp
instance Show fr => Show (Exp fr) where showsPrec = showsPrec1
instance Eq fr => Eq (Exp fr) where (==) = eq1
instance IsString (Exp Id) where fromString = ExpVar . fromString

-- janky but useful for debugging
showC :: Show fr => Exp (Var Int fr) -> String
showC e' = case e' of
  ExpVar v -> show v
  Absurd _ e -> "Absurd (" ++ showC e ++ ")"
  Empty -> "Empty"
  Lam _ body -> "Lam (" ++ showCS body ++ ")"
  ef :@ ex -> "(" ++ showC ef ++ ") :@ (" ++ showC ex ++ ")"
  Inl _ _ e -> "Inl (" ++ showC e ++ ")"
  Inr _ _ e -> "Inr (" ++ showC e ++ ")"
  Case e body1 body2 -> "Case (" ++ showC e ++ ") " ++
    "(" ++ showCS body1 ++ ") (" ++ showCS body2 ++ ")"
  Pair e1 e2 -> "Pair (" ++ showC e1 ++ ") (" ++ showC e2 ++ ")"
  Fst e -> "Fst (" ++ showC e ++ ")"
  Snd e -> "Snd (" ++ showC e ++ ")"
  MapPos o _ _ body e -> "MapPos (" ++ show o ++ ") " ++
    "(" ++ showCS body ++ ") (" ++ showC e ++ ")"
  Fold _ e -> "Fold (" ++ showC e ++ ")"
  Rec _ _ body e -> "Rec (" ++ showCS body ++ ") (" ++ showC e ++ ")"
  Unfold _ e -> "Unfold (" ++ showC e ++ ")"
  Gen _ _ body e -> "Gen (" ++ showCS body ++ ") (" ++ showC e ++ ")"
  where showCS = showC . fmap flattenVar . fromScope
        flattenVar (B ()) = B 0
        flattenVar (F (B n)) = B (n + 1)
        flattenVar (F (F fr)) = F fr


-- Statics

-- Given a set of in-scope type variables, is this type term a well-formed
-- type?
isTy :: Set Id -> OTy Id -> Bool
isTy d t' = case t' of
  TyVar v -> v `S.member` d
  VoidTy -> True
  Unit -> True
  s :-> t -> isTy d s && isTy d t
  s :+ t -> isTy d s && isTy d t
  s :* t -> isTy d s && isTy d t
  Mu o -> let (new, inst) = openFresh "t" o
          in isTy (S.insert new d) inst && isPos d o
  Nu o -> let (new, inst) = openFresh "t" o
          in isTy (S.insert new d) inst && isPos d o

-- Given a set of in-scope type variables, is this a strictly positive type
-- operator?
isPos :: Set Id -> Scope () OTy Id -> Bool
isPos d o@(Scope unwrapped) = case unwrapped of
  TyVar (B ()) -> True
  -- Bound uses this "generalized de Bruijn" so we have a Ty and not an Id.
  -- This technically doesn't follow the definition of positivity given in the
  -- book as written, but I'm pretty sure that allowing arbitrary
  -- "coefficients" is standard?
  TyVar (F t) -> isTy d t
  VoidTy -> True
  Unit -> True
  s :-> t -> let (_, s') = openFresh "t" (Scope s)
             in isTy d s' && isPos d (Scope t)
  s :+ t -> isPos d (Scope s) && isPos d (Scope t)
  s :* t -> isPos d (Scope s) && isPos d (Scope t)
  -- It seems safest to just require Mu and Nu to not have the bound type
  -- variable free.
  Mu _ -> let (_, inst) = openFresh "t" o
          in isTy d inst
  Nu _ -> let (_, inst) = openFresh "t" o
          in isTy d inst

isPosClosed :: TyOp -> Bool
isPosClosed = isPos S.empty . vacuous

-- Given a set of in-scope expression variables, is this expression well-typed,
-- and if so, what is its type?
tyOf :: Map Id Ty -> Exp Id -> Maybe Ty
tyOf g e' = case e' of
  ExpVar v -> M.lookup v g
  Absurd t e -> do VoidTy <- tyOf g e; return t
  Empty -> Just Unit
  Lam dom body ->
    let (new, inst) = openFresh "x" body
    in (dom :->) <$> tyOf (M.insert new dom g) inst
  ef :@ ex -> do
    dom :-> cod <- tyOf g ef
    dom' <- tyOf g ex
    guard (dom' == dom)
    return cod
  Inl _ tr e -> (:+ tr) <$> tyOf g e
  Inr tl _ e -> (tl :+) <$> tyOf g e
  Case e body1 body2 -> do
    tl :+ tr <- tyOf g e
    let (new1, inst1) = openFresh "l" body1
        (new2, inst2) = openFresh "r" body2
    tbody1 <- tyOf (M.insert new1 tl g) inst1
    tbody2 <- tyOf (M.insert new2 tr g) inst2
    guard (tbody1 == tbody2)
    return tbody1
  Pair e1 e2 -> (:*) <$> tyOf g e1 <*> tyOf g e2
  Fst e -> do t :* _ <- tyOf g e; return t
  Snd e -> do _ :* t <- tyOf g e; return t
  MapPos o ti1 ti2 body e -> do
    guard (isPosClosed o)
    let (new, inst) = openFresh "x" body
    ti2' <- tyOf (M.insert new ti1 g) inst
    guard (ti2' == ti2)
    tc <- tyOf g e
    guard (tc == instantiate1 ti1 o)
    return (instantiate1 ti2 o)
  Fold o e -> do
    guard (isPosClosed o)
    t <- tyOf g e
    guard (t == instantiate1 (Mu o) o)
    return (Mu o)
  Rec o t body e -> do
    guard (isPosClosed o)
    let (new, inst) = openFresh "x" body
    t' <- tyOf (M.insert new (instantiate1 t o) g) inst
    guard (t' == t)
    te <- tyOf g e
    guard (te == Mu o)
    return t
  Unfold o e -> do
    guard (isPosClosed o)
    t <- tyOf g e
    guard (t == Nu o)
    return (instantiate1 (Nu o) o)
  Gen o t body e -> do
    guard (isPosClosed o)
    t' <- tyOf g e
    guard (t' == t)
    let (new, inst) = openFresh "x" body
    tb <- tyOf (M.insert new t' g) inst
    guard (tb == instantiate1 t o)
    return (Nu o)


-- Dynamics
-- These will be strict, unlike in the book.

-- Is this expression a value?
val :: CExp -> Bool
val e = case e of
  Empty -> True
  Lam _ _ -> True
  Inl _ _ e' -> val e'
  Inr _ _ e' -> val e'
  Pair e1 e2 -> val e1 && val e2
  Fold _ e' -> val e'
  Gen _ _ _ e' -> val e' -- I *think* this is the correct definition!
  _ -> False

reduce :: CExp -> Maybe CExp
reduce r = case r of
  ef :@ ex
    | Just ef' <- reduce ef -> Just (ef' :@ ex)
    | val ef, Just ex' <- reduce ex -> Just (ef :@ ex')
    | Lam _ body <- ef, val ex -> Just (instantiate1 ex body)
  Inl tl tr e | Just e' <- reduce e -> Just (Inl tl tr e')
  Inr tl tr e | Just e' <- reduce e -> Just (Inr tl tr e')
  Case e body1 body2
    | Just e' <- reduce e -> Just (Case e' body1 body2)
    | Inl _ _ e1 <- e, val e1 -> Just (instantiate1 e1 body1)
    | Inr _ _ e2 <- e, val e2 -> Just (instantiate1 e2 body2)
  Pair e1 e2
    | Just e1' <- reduce e1 -> Just (Pair e1' e2)
    | val e1, Just e2' <- reduce e2 -> Just (Pair e1 e2')
  Fst e
    | Just e' <- reduce e -> Just (Fst e')
    | Pair e1 e2 <- e, val e1, val e2 -> Just e1
  Snd e
    | Just e' <- reduce e -> Just (Snd e')
    | Pair e1 e2 <- e, val e1, val e2 -> Just e2
  MapPos (Scope ouw) ti1 ti2 body e -> case ouw of
    TyVar (B ()) -> Just (instantiate1 e body)
    dom :-> cod
      | Just dom' <- closed dom ->
        let app = ExpVar (F e) :@ exp0
        in Just (Lam dom' (Scope
             (MapPos (Scope cod) ti1 ti2 (vacuous body) app)))
      | otherwise -> Nothing
    s :+ t ->
      let (os, ot) = (Scope s, Scope t)
          (s', t') = (instantiate1 ti2 os, instantiate1 ti2 ot)
          branch c o' = Scope (c s' t'
            (MapPos o' ti1 ti2 (vacuous body) exp0))
      in Just (Case e (branch Inl os) (branch Inr ot))
    s :* t -> Just (Pair
      (MapPos (Scope s) ti1 ti2 body (Fst e))
      (MapPos (Scope t) ti1 ti2 body (Snd e)))
    -- The other branches all eta-expand, but here we won't. The book
    -- includes eta expansion for void, for some reason, which we won't do
    -- (it doesn't do it for unit, though!)... Is this significant?
    _ -> Just e
  Fold o e | Just e' <- reduce e -> Just (Fold o e')
  Rec o t body e
    | Just e' <- reduce e -> Just (Rec o t body e')
    | Fold _ e' <- e, val e' ->
      let body' = Scope (Rec o t (vacuous body) exp0)
      in Just (instantiate1 (MapPos o (Mu o) t body' e') body)
  Unfold o e
    | Just e' <- reduce e -> Just (Unfold o e')
    | Gen _ t body e' <- e, val e' ->
      let body' = Scope (Gen o t (vacuous body) exp0)
      in Just (MapPos o t (Nu o) body' (instantiate1 e' body))
  Gen o t body e | Just e' <- reduce e -> Just (Gen o t body e')
  _ -> Nothing

normalize :: Int -> CExp -> CExp
normalize 0 e = e
normalize n e = case reduce e of
  Just e' -> normalize (if n > 0 then n - 1 else n) e'
  Nothing -> e


-- Some sample definitions in the language.

natOp :: TyOp
nat :: Ty
Just natOp = tyOp "nat" $ Unit :+ "nat"
nat = Mu natOp

zero, suc :: Exp Id
zero = Fold natOp (Inl Unit nat Empty)
suc = lam nat "n" (Fold natOp (Inr Unit nat "n"))

unfoldNat :: Exp Id
unfoldNat = lam nat "n"
  (rec natOp (Unit :+ nat) "d"
    (mapPos natOp (Unit :+ nat) nat "u" (Fold natOp "u") "d") "n")

embedNat :: Int -> Exp Id
embedNat n | n < 0 = error "not a nat"
embedNat 0 = zero
embedNat n = Fold natOp (Inr Unit nat (embedNat (n - 1)))

iterNat :: Ty -> Exp Id
iterNat t = lam nat "n" . lam (t :-> t) "f" . lam t "x" $
  rec natOp t "u" (case' "u" "_" "x" "p" ("f" :@ "p")) "n"

add :: Exp Id
add = lam nat "n" $ iterNat nat :@ "n" :@ suc

conat :: Ty
conat = Nu natOp

cozero, cosuc :: Exp Id
cozero = gen natOp Unit "_" (Inl Unit Unit Empty) Empty
cosuc = lam conat "n" (gen natOp stTy "st"
  (case' "st"
    "_" (Inr Unit stTy (Inr Unit conat "n"))
    "c" (mapPos natOp conat stTy
      "c'" (Inr Unit conat "c'") (Unfold natOp "c")))
  (Inl Unit conat Empty))
  where stTy = Unit :+ conat

omega :: Exp Id
omega = gen natOp Unit "_" (Inr Unit Unit Empty) Empty

natToConat :: Exp Id
natToConat = lam nat "n" (gen natOp nat "st"
  (case' (unfoldNat :@ "st")
    "_" (Inl Unit nat Empty)
    "p" (Inr Unit nat "p"))
  "n")

bool :: Ty
bool = Unit :+ Unit

tt, ff :: Exp Id
(tt, ff) = (Inr Unit Unit Empty, Inl Unit Unit Empty)

natLeConat :: Exp Id
natLeConat = lam nat "n" (iterNat (conat :-> bool) :@ "n" :@
  lam (conat :-> bool) "ple" (lam conat "c"
    (case' (Unfold natOp "c") "_" ff "c'" ("ple" :@ "c'"))) :@
  lam conat "_" tt)

listOp :: TyOp
list :: Ty
Just listOp = tyOp "list" $ Unit :+ vacuous nat :* "list"
list = Mu listOp

nil, cons :: Exp Id
nil = Fold listOp (Inl Unit (nat :* list) Empty)
cons = lam nat "h" . lam list "t" $
  Fold listOp (Inr Unit (nat :* list) (Pair "h" "t"))

embedList :: [Int] -> Exp Id
embedList [] = nil
embedList (n:ns) = Fold listOp (Inr Unit (nat :* list)
  (Pair (embedNat n) (embedList ns)))

streamOp :: TyOp
stream :: Ty
Just streamOp = tyOp "stream" $ vacuous nat :* "stream"
stream = Nu streamOp

strgen :: Ty -> Id -> Exp Id -> Exp Id -> Exp Id -> Exp Id
strgen stTy v initial hd tl = gen streamOp stTy v (Pair hd tl) initial

seqn :: Ty
seqn = nat :-> nat

streamOfSeq :: Exp Id
streamOfSeq = lam seqn "a" (strgen nat "st" zero ("a" :@ "st") (suc :@ "st"))

streamPrefix :: Exp Id
streamPrefix = lam stream "s" . lam nat "n" $ sub "n" :@ "s"
  where sub = rec natOp (stream :-> list) "u"
          (case' "u" "_" (lam stream "_" nil)
                     "p" (lam stream "s'" $ cons :@
                           Fst (Unfold streamOp "s'") :@
                           ("p" :@ Snd (Unfold streamOp "s'"))))


-- demo
main :: IO ()
main = print (tt' == normalize (-1) cmp1, tt' == normalize (-1) cmp2)
  where Just [tt', cmp1, cmp2] = traverse closed [tt,
          natLeConat :@ embedNat 30 :@ (natToConat :@ embedNat 29),
          natLeConat :@ embedNat 30 :@ omega]