{-# OPTIONS_GHC -fglasgow-exts #-} 
module PolyLet where 

import Iso
import Games 
import BasicGames

import Data.Maybe
import FilterGames

-- Simple types
-- /TyExp/
data Ty = TyVar Nat | TyArr Ty Ty | TyProd Ty Ty deriving (Eq, Show)
data Exp = Var Nat [Ty] | Lam Ty Exp | App Exp Exp | Let Int Exp Exp
-- /End/
  deriving (Eq,Show)

-- Type schemes, with number of bound vars first
data TySch = TySch Int Ty deriving Eq

instance Show TySch where
  show (TySch 0 ty) = show ty
  show (TySch n ty) = "(forall " ++ show n ++ ")" ++ show ty

instantiate :: [Ty] -> Ty -> Ty
instantiate tys (TyVar i) = if i < length tys then tys !! i else TyVar (i - length tys)
instantiate tys (TyArr t1 t2) = TyArr (instantiate tys t1) (instantiate tys t2)
instantiate tys (TyProd t1 t2) = TyProd (instantiate tys t1) (instantiate tys t2)

instantiateSch :: TySch -> [Ty] -> Ty
instantiateSch (TySch _ ty) tys = instantiate tys ty

type Subst = [Maybe Ty]
subst :: Subst -> Ty -> Ty
subst [] ty = ty
subst (Just ty:s) (TyVar 0) = TyVar 0
subst (Nothing:s) (TyVar i) = TyVar i
subst (_:s) (TyVar (i+1)) = subst s (TyVar i)
subst s (TyArr t1 t2) = TyArr (subst s t1) (subst s t2)
subst s (TyProd t1 t2) = TyProd (subst s t1) (subst s t2)

merge [] [] = []
merge (Nothing:s1) (Nothing:s2) = Nothing:merge s1 s2
merge (_:s1) (Just ty:s2) = Just ty:merge s1 s2
merge (Just ty:s1) (_:s2) = Just ty:merge s1 s2

singleton ntyvars i ty = replicate i Nothing ++ [Just ty] ++ replicate (ntyvars-i-1) Nothing

-- Attempt to match the first n type variables in the second type against the first
matchTy :: Ty -> Subst -> Ty -> Maybe Subst
matchTy ty s (TyVar i) = 
  if i<length s
  then case s!!i of Nothing -> Just (merge (singleton (length s) i ty) s) ; Just ty' -> if ty==ty' then Just s else Nothing
  else if ty == TyVar (i-length s) then Just s else Nothing
matchTy (TyArr ty1a ty1b) s (TyArr ty2a ty2b) = 
  case matchTy ty1a s ty2a of
    Nothing -> Nothing
    Just s' -> matchTy ty1b s' ty2b
matchTy (TyProd ty1a ty1b) s (TyProd ty2a ty2b) =  
  case matchTy ty1a s ty2a of
    Nothing -> Nothing
    Just s' -> matchTy ty1b s' ty2b
matchTy _ _ _ = Nothing

matchSch :: Ty -> TySch -> Maybe Subst
matchSch ty (TySch n ty') = matchTy ty (replicate n Nothing) ty'

intTy = TyVar 0
boolTy = TyVar 1

showNiceTy :: [String] -> Ty -> String
showNiceTy names (TyVar i) = names !! i
showNiceTy names (TyArr ty1 ty2) = "(" ++ showNiceTy names ty1 ++ "->" ++ showNiceTy names ty2 ++ ")"
showNiceTy names (TyProd ty1 ty2) = "(" ++ showNiceTy names ty1 ++ "*" ++ showNiceTy names ty2 ++ ")"

var n = Var n []

iAtBool = Lam boolTy (var 0)
iAtBoolToBool = Lam (TyArr boolTy boolTy) (var 0)
iAtInt = Lam intTy (var 0)
kAtBool = Lam boolTy (Lam boolTy (var 1))
kAtInt = Lam intTy (Lam intTy  (var 1))
ii = App iAtBoolToBool iAtBool
twiceTm = Lam (TyArr intTy intTy) (Lam intTy (App (var 1) (App (var 1) (var 0))))

type Env = (Int, [TySch])

-- Types for fst, snd, pair, zero, succ
exEnv :: Env
exEnv = (2, [TySch 2 (TyArr (TyProd (TyVar 0) (TyVar 1)) (TyVar 0)),
            TySch 2 (TyArr (TyProd (TyVar 0) (TyVar 1)) (TyVar 1)),
            TySch 2 (TyArr (TyVar 0) (TyArr (TyVar 1) (TyProd (TyVar 0) (TyVar 1)))),
            TySch 0 intTy,
            TySch 0 (TyArr intTy intTy),
            TySch 1 (TyArr (TyArr (TyVar 0) (TyVar 0)) (TyArr (TyVar 0) (TyVar 0)))
            ])

typeOf :: Env -> Exp -> Ty
typeOf (_,env) (Var i tys) = instantiateSch (env !! i) tys
typeOf env (App e1 e2) = case typeOf env e1 of TyArr t1 t2 -> t2 
typeOf (n,env) (Lam t e) = TyArr t (typeOf (n, TySch 0 t:env) e)
typeOf (n,env) (Let m e1 e2) = typeOf (m, TySch m (typeOf (m+n,env) e1) : env) e2

showTys names [] = ""
showTys names [ty] = showNiceTy names ty
showTys names (ty:tys) = showNiceTy names ty ++ "," ++ showTys names tys

niceName names = let name = [toEnum (length names + fromEnum 'a')] in (name, name:names)

niceNames 0 names = names
niceNames (n+1) names = let (_,names') = niceName names in niceNames n names'

showNice :: Env -> [String] -> [String] -> Exp -> String
showNice (env @ (ntyvars,tyenv)) names tynames t =
  case t of 
    Var i [] -> names !! i
    Var i tys -> names !! i ++ "{" ++ showTys tynames tys ++ "}"
    App t1 t2 -> showNice env names tynames t1 ++ " " ++ showNice env names tynames t2
    Lam ty t -> let (name,names') = niceName names in "(\\" ++ name ++ ":" ++ showNiceTy tynames ty ++ "." ++ showNice (ntyvars, TySch 0 ty : tyenv) names' tynames t ++ ")" 
    Let n t1 t2 -> 
      let tynames' = niceNames n tynames in 
      let (name,names') = niceName names in 
      "let(" ++ show n ++ ")" ++ name ++ " = " ++ showNice (n+ntyvars,tyenv) names tynames' t1 ++ " in " ++ showNice (ntyvars,TySch n (typeOf (n+ntyvars,tyenv) t1) : tyenv) names' tynames t2

showClosed t = showNice exEnv ["fst", "snd", "pair", "zero", "succ", "twice"] ["Int", "Bool"] t

ex1 = 
  Let 1 
    (Lam (TyArr (TyVar 0) (TyVar 0)) 
      (Lam (TyVar 0) 
        (App (Var 1 []) (App (Var 1 []) (Var 0 [])))))
      (App (Var 0 [intTy]) (Var 5 []))

-- Match a type scheme against a pattern
data Pat = Any | PArr Ty Pat
matchMatch :: Pat -> Subst -> Ty -> Maybe Subst
matchMatch m s ty =
  case (m, ty) of
    (Any, _) -> Just s
    (PArr ty1 m', TyArr ty2 ty2') ->
      case matchTy ty1 s ty2 of
        Nothing -> Nothing
        Just s' -> matchMatch m' s' ty2'
    _ -> Nothing


matches :: Pat -> TySch -> Maybe Subst
matches p (TySch n t) = matchMatch p (replicate n Nothing) t

-- Game for types 
-- /tyG/
tyGame :: Nat -> Game Ty 
tyGame 0 = (prodGame (tyGame 0) (tyGame 0)) +> Iso (\(TyArr t1 t2) -> (t1,t2)) (\(t1,t2) -> TyArr t1 t2) 
tyGame ntyvars = Split (Iso ask bld) 
                   (rangeGame 0 (ntyvars-1)) (prodGame (tyGame ntyvars) (tyGame ntyvars))
 where ask (TyVar i) = Left i
       ask (TyArr t1 t2) = Right (t1,t2) 
       bld (Left i) = TyVar i
       bld (Right (t1,t2)) = TyArr t1 t2
-- /End/


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

-- Given a template for a list of a's that fills in some of the elements, create
-- a game that fills out the missing elements
partialVecGame :: [Maybe a] -> Game a -> Game [a]
partialVecGame [] g = constGame []
partialVecGame (x:xs) g = prodGame (maybe g constGame x) (partialVecGame xs g) +> nonemptyIso

instGame :: Int -> Subst -> Game [Ty]       
instGame ntyvars s = partialVecGame s (tyGame ntyvars)
  
-- Game for matching variables 
-- /mkVarGame/
varGame :: (TySch -> Maybe Subst) -> Env -> Maybe (Game (Nat,[Ty]))
varGame f (_,[]) = Nothing 
varGame f (ntyvars,t:env) =   
  case varGame f (ntyvars,env) of 
    Nothing -> 
      case f t of 
        Nothing -> Nothing
        Just s -> Just (prodGame (constGame 0) (instGame ntyvars s))
    Just g  -> 
      case f t of 
        Nothing -> Just (g +> Iso (\(n,i) -> (pred n,i)) (\(n,i) -> (succ n,i))) 
        Just s -> Just (Split (Iso ask bld) (instGame ntyvars s) g)
    where ask (0,i) = Left i
          ask (n+1,i) = Right (n,i) 
          bld (Left i) = (0,i) 
          bld (Right (n,i)) = (n+1,i)
-- /End/

progGame :: Game Exp
progGame = expGame exEnv Any

posGame :: Game Nat
posGame = unaryNatGame +> Iso pred succ

-- Returns an expression with a type that that matches match 
-- Satisfies the "all bits count" property
-- /expGame/
-- (env : Env) -> (p : Pat) -> 
--   Game {e | exists t, env |- e : t && matches p t} 
expGame :: Env -> Pat -> Game Exp
expGame (env@(ntyvars,tyschs)) p = 
  case varGame (matches p) env of 
    Nothing -> nonVarG
    Just varG -> Split varI varG nonVarG
  where nonVarG = Split nonVarI letG appLamG
        appLamG = Split appLamI appG (lamG p)
        letG = 
          depGame posGame $ \nbound ->
          depGame (expGame (nbound+ntyvars,tyschs) Any) $ \e1 ->
          expGame (ntyvars, TySch nbound (typeOf (nbound+ntyvars,tyschs) e1) : tyschs) p
          
        appG = depGame (expGame env Any) $ \e -> 
               expGame env (PArr (typeOf env e) p) 
        lamG (PArr t p) = prodGame (constGame t) $ 
                          expGame (ntyvars,TySch 0 t:tyschs) p
        lamG Any = depGame (tyGame ntyvars) $ \t -> 
                   expGame (ntyvars,TySch 0 t:tyschs) Any 

varI = Iso ask bld where ask (Var x inst)    = Left (x,inst)
                         ask e               = Right e
                         bld (Left (x,inst)) = Var x inst
                         bld (Right e)       = e 
                         
nonVarI = Iso ask bld                         
  where ask (Let n e1 e2) = Left (n, (e1,e2))
        ask e = Right e
        bld (Left (n, (e1,e2))) = Let n e1 e2
        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 

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

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

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

allTms = enumerateTms (allLists 0)