module STLC where 

import Iso
import Games 
import BasicGames

import Data.Maybe
import FilterGames

-- Simple types
-- /TyExp/
data Ty = TyNat | TyArr Ty Ty deriving (Eq, Show)
data Exp = Var Nat | Lam Ty Exp | App Exp Exp 
-- /End/
  deriving (Eq,Show)

-- Game for types 
-- /tyG/
tyG :: Game Ty 
tyG = Split (Iso ask bld) unitGame (prodGame tyG tyG)
 where ask TyNat = Left ()
       ask (TyArr t1 t2) = Right (t1,t2) 
       bld (Left ()) = TyNat
       bld (Right (t1,t2)) = TyArr t1 t2
-- /End/


-- Environment is just a list of types
-- Precondition: expression well typed in environment
-- /typeOf/
type Env = [Ty] 
typeOf :: Env -> Exp -> Ty
typeOf env (Var i) = env !! i
typeOf env (App e _) = let TyArr _ t = typeOf env e in t
typeOf env (Lam t e) = TyArr t (typeOf (t:env) e)
-- /End/

-- Matching 
-- /Pat/
data Pat = Any | PArr Ty Pat
matches :: Pat -> Ty -> Bool
matches Any _ = True 
matches (PArr t p) (TyArr t1 t2) = t1==t && matches p t2
matches _ _ = False 
-- /End/

{- Let the Games begin!   
   ~~~~~~~~~~~~~~~~~~~~ -} 



-- Game for matching variables 
-- /mkVarGame/
varGame :: (Ty -> Bool) -> Env -> Maybe (Game Nat)
varGame f [] = Nothing 
varGame f (t:env) = case varGame f env of 
 Nothing -> if f t then Just (constGame 0) else Nothing
 Just g  -> if f t then Just (Split succIso unitGame g)
            else Just (g +> Iso pred succ)
-- /End/

-- Returns an expression with a type that that matches match 
-- Satisfies the "all bits count" property
-- /expGame/
expGame :: Env -> Pat -> Game Exp
-- forall (env:Env) (p:Pat), 
--   Game { e | exists t, env |- e : t && matches p t = true }
expGame env p
  = case varGame (matches p) env of 
      Nothing -> appLamG 
      Just varG -> Split varI varG appLamG
  where appLamG = Split appLamI appG (lamG p)
        appG = depGame (expGame env Any) $ \e -> 
               expGame env (PArr (typeOf env e) p) 
        lamG (PArr t p) = prodGame (constGame t) $
                          expGame (t:env) p
        lamG Any = depGame tyG $ \t -> 
                   expGame (t:env) Any 

varI = Iso ask bld where ask (Var x)   = Left x
                         ask e         = Right e
                         bld (Left x)  = Var x 
                         bld (Right e) = e 
appLamI = Iso ask bld 
  where ask (App e1 e2)    = Left (e2,e1) 
        ask (Lam t e)      = Right (t,e) 
        bld (Left (e2,e1)) = App e1 e2 
        bld (Right (t,e))  = Lam t e 

-- /End/

progGame :: Game Exp
progGame = expGame [] Any

-- Returns a game for terms in a *given* environment and *given* type.
-- /expGameCheck/
-- forall (env:Env) (t:Ty), Game { e | env |- e : t }
expGameCheck :: Env -> Ty -> Game Exp 
expGameCheck env t 
  = case varGame (== t) env of
      Nothing -> appLamG t
      Just varG -> Split varI varG (appLamG t)
  where appLamG TyNat 
          = appG +> Iso (\(App e1 e2)->(e2,e1))
                        (\(e2,e1)->App e1 e2)
        appLamG (TyArr t1 t2) 
          = let ask (App e1 e2)    = Left (e2,e1) 
                ask (Lam t e)      = Right e 
                bld (Left (e2,e1)) = App e1 e2
                bld (Right e)      = Lam t1 e 
            in Split (Iso ask bld) appG (lamG t1 t2)  
        appG = depGame (expGame env Any) $ \e -> 
               expGameCheck env (TyArr (typeOf env e) t)
        lamG t1 t2 = expGameCheck (t1:env) t2

-- /End/  

-- -- A strong model for terms (will be strong only if there are infinite inhabitants) 
-- -- [Satisfies the all bits count property] 
-- /expGameCheckProper/
expGameCheckProper env t 
  = filterInfGame (\_ -> True) (expGameCheck env t) 
-- /End/ 


-- /allTerms/ 
all01 = [O] : map (O:) all01
-- Games for the empty environment and type Nat -> Nat
allNat2Nat = map (fst . dec game) all01
  where game = expGameCheckProper [] (TyArr TyNat TyNat)
-- /End/

-- decRandomTm i = run decClosedTm (mkRandom i)

listsOfLength :: Int -> [[Bit]]
listsOfLength 0 = [[]]
listsOfLength (n+1) = map (O:) (listsOfLength n) ++ map (I:) (listsOfLength n)

allLists n = listsOfLength n ++ allLists (n+1)

enumerateTms (x:l) =
  case decOpt (expGame [] Any) x of
    Just (e,[]) -> e : enumerateTms l
    _ -> enumerateTms l

allTms = enumerateTms (allLists 0)


ex = Lam TyNat (Lam TyNat (Var 1))