module MinHS.TyInfer where

import qualified MinHS.Env as E
import MinHS.Syntax
import MinHS.Subst
import MinHS.TCMonad

import Data.Monoid (Monoid (..), (<>))
import Data.Foldable (foldMap)
import Data.List (nub, union, (\\))
import Control.Monad (zipWithM)
import Debug.Trace

primOpType :: Op -> QType
primOpType Gt   = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Bool)
primOpType Ge   = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Bool)
primOpType Lt   = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Bool)
primOpType Le   = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Bool)
primOpType Eq   = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Bool)
primOpType Ne   = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Bool)
primOpType Neg  = Ty $ Base Int `Arrow` Base Int
primOpType Fst  = Forall "a" $ Forall "b" $ Ty $ (TypeVar "a" `Prod` TypeVar "b") `Arrow` TypeVar "a"
primOpType Snd  = Forall "a" $ Forall "b" $ Ty $ (TypeVar "a" `Prod` TypeVar "b") `Arrow` TypeVar "b"
primOpType _    = Ty $ Base Int `Arrow` (Base Int `Arrow` Base Int)

constType :: Id -> Maybe QType
constType "True"  = Just $ Ty $ Base Bool
constType "False" = Just $ Ty $ Base Bool
constType "()"    = Just $ Ty $ Base Unit
constType "Pair"  = Just
                  $ Forall "a"
                  $ Forall "b"
                  $ Ty
                  $ TypeVar "a" `Arrow` (TypeVar "b" `Arrow` (TypeVar "a" `Prod` TypeVar "b"))
constType "Inl"   = Just
                  $ Forall "a"
                  $ Forall "b"
                  $ Ty
                  $ TypeVar "a" `Arrow` (TypeVar "a" `Sum` TypeVar "b")
constType "Inr"   = Just
                  $ Forall "a"
                  $ Forall "b"
                  $ Ty
                  $ TypeVar "b" `Arrow` (TypeVar "a" `Sum` TypeVar "b")
constType _       = Nothing

type Gamma = E.Env QType

initialGamma :: Gamma
initialGamma = E.empty

tv :: Type -> [Id]
tv = tv'
 where
   tv' (TypeVar x) = [x]
   tv' (Prod  a b) = tv a `union` tv b
   tv' (Sum   a b) = tv a `union` tv b
   tv' (Arrow a b) = tv a `union` tv b
   tv' (Base c   ) = []

tvQ :: QType -> [Id]
tvQ (Forall x t) = filter (/= x) $ tvQ t
tvQ (Ty t) = tv t

tvGamma :: Gamma -> [Id]
tvGamma = nub . foldMap tvQ

infer :: Program -> Either TypeError Program
infer program = do (p',tau, s) <- runTC $ inferProgram initialGamma program
                   return p'

unquantify :: QType -> TC Type
{-
Normally this implementation would be possible:

unquantify (Ty t) = return t
unquantify (Forall x t) = do x' <- fresh
                             unquantify (substQType (x =:x') t)

However as our "fresh" names are not checked for collisions with names bound in the type
we avoid capture entirely by first replacing each bound
variable with a guaranteed non-colliding variable with a numeric name,
and then substituting those numeric names for our normal fresh variables
-}

unquantify = unquantify' 0 emptySubst
unquantify' :: Int -> Subst -> QType -> TC Type
unquantify' i s (Ty t) = return $ substitute s t
unquantify' i s (Forall x t) = do x' <- fresh
                                  unquantify' (i + 1)
                                              ((show i =: x') <> s)
                                              (substQType (x =:TypeVar (show i)) t)

-- substitute, then merge.
submerge :: Type -> Type -> Type -> Type -> TC Subst
submerge a1 a2 b1 b2 = do
  s <- unify a1 a2
  s' <- unify (substitute s b1) (substitute s b2)
  return (s <> s')  

unify :: Type -> Type -> TC Subst
unify (TypeVar a) (TypeVar b) 
  | a == b = return mempty
  | otherwise = return (a =: TypeVar b)
unify (Base a) (Base b)
  | a == b = return mempty
unify (Prod a1 b1) (Prod a2 b2) = submerge a1 a2 b1 b2
unify (Sum a1 b1) (Sum a2 b2) = submerge a1 a2 b1 b2
unify (Arrow a1 b1) (Arrow a2 b2) = submerge a1 a2 b1 b2
unify (TypeVar a) t
  | a `elem` tv t = typeError $ OccursCheckFailed a t
  | otherwise = return (a =: t)
unify t (TypeVar a) = unify (TypeVar a) t
unify a b = typeError $ TypeMismatch a b

{- 
  returns a subst S such that
  if match t t' = s
  then st = t'.
  we use this to match our types against user ones
  since this doesn't shift around with existing fresh
  variables and whatnot, this almost certainly has a bug
  where the user overrides a fresh variable in the environment.
  But, this is due in three hours and I don't feel like implementing
  a safe environment rewrite-around-user-types rule
-}
match :: Type -> Type -> TC Subst
match (Prod a1 b1) (Prod a2 b2) = matchmerge a1 a2 b1 b2
match (Sum a1 b1) (Sum a2 b2) = matchmerge a1 a2 b1 b2
match (Arrow a1 b1) (Arrow a2 b2) = matchmerge a1 a2 b1 b2
match (TypeVar a) b = return $ a =: b
match (Base a) (Base b)
  | a == b = return mempty
match a b = typeError $ TypeMismatch a b

matchmerge :: Type -> Type -> Type -> Type -> TC Subst
matchmerge a1 a2 b1 b2 = do
  s1 <- match a1 a2
  s2 <- match b1 b2
  return $ s1 <> s2

combineUserType :: QType -> Maybe QType -> TC Subst
combineUserType t Nothing = return emptySubst
combineUserType t1 (Just t2) = do
  t1' <- unquantify t1
  t2' <- unquantify t2
  match t1' t2'

letfunCombineUserType :: QType -> Maybe QType -> TC Subst
letfunCombineUserType t Nothing = combineUserType t Nothing
letfunCombineUserType t t'@(Just (Ty _)) = combineUserType t t'
letfunCombineUserType _ _ = typeError ForallInLetfun  

generalise :: Gamma -> Type -> QType
generalise g t = 
  let free = tv t \\ tvGamma g
  in foldr Forall (Ty t) free 

inferProgram :: Gamma -> Program -> TC (Program, Type, Subst)
inferProgram env [Bind "main" userType [] body] = do
  (body', t, s) <- inferExp env body
  s' <- combineUserType (Ty t) userType
  let t' = s' `substitute` t
  let qsubst = substQType (s <> s')
  let body'' = allTypes qsubst body'
  let t'' = generalise env t'
  let main = Bind "main" (Just t'') [] body''
  return ([main], t', s)
inferProgram env bs = 
  error $ "ugh gotta do program for\n" ++
    show env ++ "\nwith body\n" ++ show bs
 --"implement me! don't forget to run the result substitution on the\n" ++
 --"entire expression using allTypes from Syntax.hs"

freshFor :: [Id] -> TC [(Id, QType)]
freshFor = mapM (\i -> (\t -> (i, Ty t)) `fmap` fresh)

inferExp :: Gamma -> Exp -> TC (Exp, Type, Subst)
inferExp _ (Num n) = return (Num n, Base Int, emptySubst)
inferExp g (Var x) = case g `E.lookup` x of
  Just t -> do
    t' <- unquantify t
    return (Var x, t', mempty)
  Nothing -> typeError $ NoSuchVariable x
inferExp g (Prim op) = do
  t <- unquantify $ primOpType op
  return (Prim op, t, mempty)
inferExp g (Con c) = case constType c of
    Just t -> do
      t' <- unquantify t
      return (Con c, t', mempty)
    Nothing -> typeError $ NoSuchConstructor c
inferExp g (App f e) = do
  (f', t1, ts) <- inferExp g f
  let g' = ts `substGamma` g
  (e', t2, ts') <- inferExp g' e
  a <- fresh
  u <- unify (ts' `substitute` t1) (Arrow t2 a)
  return (
    App f' e', 
    u `substitute` a, 
    u <> ts' <> ts)
inferExp g (If e e1 e2) = do
  (e', t, ts) <- inferExp g e
  u <- unify t (Base Bool)
  let g' = (u <> ts) `substGamma` g
  (e1', t1, ts1) <- inferExp g' e1
  let g'' = ts1 `substGamma` g'
  (e2', t2, ts2) <- inferExp g'' e2
  u' <- unify (ts2 `substitute` t1) t2
  return (
    If e' e1' e2',
    u' `substitute` t2,
    u' <> ts2 <> ts1 <> u <> ts)
inferExp g (Let binds body) = let 
  magic g' [] = do 
    (body', bt, bts) <- inferExp g' body
    return (
      (body', bt),
      [],
      bts)
  magic g' (Bind name userType [] def : bs) = do
    (def', t, ts) <- inferExp g' def

    s <- combineUserType (Ty t) userType
    --trace ("deft:\n" ++ (show t)) $ return ()
    --trace ("s deft:\n" ++ (show $ s `substitute` t)) $ return ()

    let t' = s `substitute` t
    let tsg = ts `substGamma` g'
    let g'' = tsg `E.add` (name, generalise tsg t')
    (bodyStuff, bst, bsts) <- magic g'' bs
    let tgen = generalise g'' t'
    return (
      bodyStuff,
      Bind name (Just tgen) [] def' : bst,
      s <> bsts <> ts)
  in do
    ((body', bodyt), binds', ts) <- magic g binds
    return (
      Let binds' body',
      bodyt,
      ts)
inferExp g (Letfun (Bind name userType args body)) = do
  argts <- freshFor args
  a2 <- fresh
  let g' = g `E.addAll` ((name, Ty a2) : argts)
  (body', t, ts) <- inferExp g' body
  let a2' = ts `substitute` a2
  argtypes <- mapM (unquantify . snd) argts 
  let a1' = foldr (Arrow . substitute ts) t argtypes
  u <- unify a2' a1'
  let tfinal = u `substitute` a1'
  s <- letfunCombineUserType (Ty tfinal) userType
  --trace ("tfinal:\n" ++ (show tfinal)) $ return ()
  --trace ("s tfinal:\n" ++ (show $ s `substitute` tfinal)) $ return ()
  let tfinal' = s `substitute` tfinal
  return (
    Letfun $ Bind name (Just $ Ty tfinal') args body',
    tfinal',
    s <> u <> ts)
-- Note: this is the only case you need to handle for case expressions
-- inferExp g (Case e [Alt "Inl" [x] e1, Alt "Inr" [y] e2])
-- inferExp g (Case e _) = typeError MalformedAlternatives
inferExp g (Case e [Alt "Inl" [x] e1, Alt "Inr" [y] e2]) = do
  (e', t, ts) <- inferExp g e
  al <- fresh
  let g' = ts `substGamma` (g `E.add` (x, Ty al))
  (e1', tl, ts1) <- inferExp g' e1
  ar <- fresh
  let g'' = ts1 `substGamma` (g' `E.add` (y, Ty ar))
  (e2', tr, ts2) <- inferExp g'' e2
  let alr = (ts2 <> ts1 <> ts) `substitute` Sum al ar
  u <- unify alr ((ts2 <> ts1) `substitute` t) 
  u' <- unify ((u <> ts2) `substitute` tl) (u `substitute` tr)
  return (
    Case e' [Alt "Inl" [x] e1', Alt "Inr" [y] e2'],
    (u' <> u) `substitute` tr,
    u <> u' <> ts2 <> ts1 <> ts)
inferExp g (Case e _) = typeError MalformedAlternatives
inferExp g (Letrec binds body) = do
  let binds' = (\(Bind name Nothing [] def) -> (name, def)) `map` binds
  let (names, defs) = unzip binds'
  namedTypes <- freshFor names
  let g' = g `E.addAll` namedTypes
  tmp <- mapM (inferExp g') defs
  let (defs', defts, defsubs) = unzip3 tmp
  let freshQTypes = (snd . unzip) namedTypes
  let freshTypes = (\(Ty t) -> t) `map` freshQTypes
  let bigDefSub = foldl (<>) emptySubst defsubs
  unifiers <- zipWithM unify (substitute bigDefSub `map` freshTypes) defts
  let bigUnifier = foldl (<>) emptySubst unifiers
  let bigSub = bigUnifier <> bigDefSub
  let g'' = bigSub `substGamma` g'
  let newTypes = (generalise g' . substitute bigSub) `map` defts
  let newNamedTypes = names `zip` newTypes
  (body', bodyt, bodysub) <- inferExp (g'' `E.addAll` newNamedTypes) body
  --trace ("g'':\n" ++ (show g'') ++ "\n") $ return ()
  --trace ("\nnewTypes:\n" ++ (show (names `zip` newTypes)) ++ "\n") $ return ()
  let toBind name typ = Bind name (Just typ) []
  let newBinds = zipWith3 toBind names newTypes defs'
  return (
    Letrec newBinds body',
    bodyt,
    bodysub <> bigSub)
