module Ntha.Type.Type where

import           Ntha.State
import           Ntha.Z3.Assertion
import           Ntha.Z3.Class
import           Ntha.Z3.Logic

import           Z3.Monad
import           Control.Monad     (foldM, liftM)
import           System.IO.Unsafe  (unsafePerformIO)
import           Data.IORef
import           Data.List         (intercalate)
import           Data.Maybe        (fromMaybe)
import qualified Data.Map          as M
import qualified Data.Set          as S
import qualified Text.PrettyPrint  as PP

type Id = Int
type TName = String
type TField = String
type Types = [Type]
type TInstance = Maybe Type
type Z3Pred = Pred Term RType Assertion

data Type = TVar Id (IORef TInstance) TName -- type variable
          | TOper TName Types -- type operator
          | TRecord (M.Map TField Type)
          | TCon TName Types Type
          | TSig Type
          | TRefined String Type Term

-- extract normal type from refined type for type inference
extractType :: Type -> Type
extractType t = case t of
  -- just support arrow type for now
  TOper "→" args -> TOper "→" (map extractType args)
  TRefined _ t' _ -> t'
  _ -> t

extractTerm :: Type -> [Term]
extractTerm t = case t of
  TOper "→" args -> args >>= extractTerm
  TRefined _ _ tm -> [tm]
  _ -> []

getPredNames :: Type -> [String]
getPredNames t = case t of
  TOper "→" args -> args >>= getPredNames
  TRefined n _ _ -> [n]
  _ -> []

intT :: Type
intT = TOper "Number" []

boolT :: Type
boolT = TOper "Boolean" []

charT :: Type
charT = TOper "Char" []

listT :: Type -> Type -- list type is not polymorphism
listT t = TOper "List" [t]

productT :: Types -> Type -- tuple type, product type is a name from Algebraic Data type
productT ts = TOper "*" ts

arrowT :: Type -> Type -> Type -- function type with single param
arrowT fromType toType = TOper "→" $ [fromType, toType]

functionT :: Types -> Type -> Type
functionT paramsT rtnT = foldr (\paramT resT -> arrowT paramT resT) rtnT paramsT

strT :: Type
strT = listT charT

unitT :: Type
unitT = TOper "()" []

prune :: Type -> Infer Type
prune t = case t of
  TVar _ inst _ -> do
    instV <- readIORef inst
    case instV of
      Just inst' -> do
        newInstance <- prune inst'
        writeIORef inst $ Just newInstance
        return newInstance
      Nothing -> return t
  _ -> return t

stringOfType :: M.Map TName TName -> Type -> Infer String
stringOfType subrule (TVar _ inst name) = do
  instV <- readIORef inst
  case instV of
    Just inst' -> stringOfType subrule inst'
    Nothing -> return $ fromMaybe "α" $ M.lookup name subrule
stringOfType subrule (TOper name args) = case name of
  "*" -> do
    argsStr <- (intercalate " * ") <$> mapM (stringOfType subrule) args
    return $ "(" ++ argsStr ++ ")"
  "List" -> do
    argStr <- stringOfType subrule $ args!!0
    return $ "[" ++ argStr ++ "]"
  "→" -> do
    argT <- prune $ args!!0
    rtnT <- prune $ args!!1
    argStr <- stringOfType subrule argT
    rtnStr <- stringOfType subrule rtnT
    let adjust t s = case t of
                      TOper "→" _ -> "(" ++ s ++ ")"
                      _ -> s
    let argStr' = adjust argT argStr
    let rtnStr' = adjust rtnT rtnStr
    return $ argStr' ++ " → " ++ rtnStr'
  _ -> if (length args) == 0
      then return name
      else do
        argsStr <- unwords <$> mapM (stringOfType subrule) args
        return $ "(" ++ name ++ " " ++ argsStr ++ ")"
stringOfType subrule (TRecord pairs) = do
  pairsStr <- (intercalate ", ") <$> (mapM (\(k, v) -> ((k ++ ": ") ++) <$> stringOfType subrule v) $ M.toList pairs)
  return $ "{" ++ pairsStr ++ "}"
stringOfType subrule (TCon name types dataType) = do
  dataTypeStr <- stringOfType subrule dataType
  case types of
    [] -> return dataTypeStr
    _ -> do
      typesStr <- (intercalate ", ") <$> mapM (stringOfType subrule) types
      return $ "(" ++ name ++ " " ++ typesStr ++ " ⇒ " ++ dataTypeStr ++ ")"
stringOfType subrule (TSig t) = liftM ("typesig: " ++) $ stringOfType subrule t
stringOfType subrule (TRefined _ t _) = liftM ("refined: " ++) $ stringOfType subrule t

getFreeVars :: Type -> Infer (S.Set TName)
getFreeVars (TVar _ inst name) = do
  instV <- readIORef inst
  case instV of
    Just inst' -> getFreeVars inst'
    Nothing -> return $ S.singleton name
getFreeVars (TOper _ args) =
  foldM (\acc arg -> do
    freeVars <- getFreeVars arg
    return $ S.union freeVars acc)
  S.empty args
getFreeVars (TRecord pairs) =
  foldM (\acc (_, v) -> do
    freeVars <- getFreeVars v
    return $ S.union freeVars acc)
  S.empty $ M.toList pairs
getFreeVars (TCon _ types dataType) =
  foldM (\acc t -> do
    freeVars <- getFreeVars t
    return $ S.union freeVars acc)
  S.empty $ types ++ [dataType]
getFreeVars (TSig t) = getFreeVars t
getFreeVars (TRefined _ t _) = getFreeVars t

{-# NOINLINE normalize #-}
normalize :: Type -> String
normalize t = unsafePerformIO $ do
  freeVars <- getFreeVars t
  let subrule = M.map (\c -> [c]) $ M.fromList $ zip (S.toList freeVars) ['α'..'ω']
  stringOfType subrule t

instance Show Type where
  showsPrec _ x = shows $ PP.text $ normalize x

instance Eq Type where
  TVar id1 inst1 vname1 == TVar id2 inst2 vname2 = id1 == id2 && instV1 == instV2 && vname1 == vname2 where
    instV1 = readState inst1
    instV2 = readState inst2
  TOper name1 args1 == TOper name2 args2 = name1 == name2 && args1 == args2
  TRecord pairs1 == TRecord pairs2 = pairs1 == pairs2
  TCon name1 types1 dataType1 == TCon name2 types2 dataType2 = name1 == name2 && types1 == types2 && dataType1 == dataType2
  TSig t1 == TSig t2 = t1 == t2
  TRefined x1 t1 tm1 == TRefined x2 t2 tm2 = x1 == x2 && t1 == t2 && tm1 == tm2
  _ == _ = False

instance Ord Type where
    TVar id1 inst1 vname1 <= TVar id2 inst2 vname2 = id1 <= id2 && instV1 <= instV2 && vname1 <= vname2 where
      instV1 = readState inst1
      instV2 = readState inst2
    TOper name1 args1 <= TOper name2 args2 = name1 <= name2 && args1 <= args2
    TRecord pairs1 <= TRecord pairs2 = pairs1 <= pairs2
    TCon name1 types1 dataType1 <= TCon name2 types2 dataType2 = name1 <= name2 && types1 <= types2 && dataType1 <= dataType2
    TSig t1 <= TSig t2 = t1 <= t2
    TRefined x1 t1 tm1 <= TRefined x2 t2 tm2 = x1 <= x2 && t1 <= t2 && tm1 <= tm2
    _ <= _ = False

makeVariable :: Infer Type
makeVariable = do
    i <- nextId
    name <- nextUniqueName
    instRef <- newIORef Nothing
    return $ TVar i instRef name

-- for refined type

data Term = TmVar   String
          | TmNum   Int
          | TmLT    Term Term
          | TmGT    Term Term
          | TmLE    Term Term
          | TmGE    Term Term
          | TmSub   Term Term
          | TmAdd   Term Term
          | TmMul   Term Term
          | TmDiv   Term Term
          | TmEqual Term Term
          | TmAnd   Term Term
          | TmOr    Term Term
          | TmNot   Term
          | TmIf    Term Term Term

deriving instance Eq Term
deriving instance Ord Term
deriving instance Show Term

-- currently just support integer
data RType = RTInt

deriving instance Eq RType
deriving instance Ord RType

instance Z3Encoded Term where
    encode (TmVar x) = do
        ctx <- getQualifierCtx
        case M.lookup x ctx of
            Just (idx, _) -> return idx
            Nothing -> smtError $ "Can't find variable " ++ x
    encode (TmNum n) = mkIntSort >>= mkInt n
    encode (TmLT t1 t2) = encode (Less t1 t2)
    encode (TmGT t1 t2) = encode (Greater t1 t2)
    encode (TmLE t1 t2) = encode (LessE t1 t2)
    encode (TmGE t1 t2) = encode (GreaterE t1 t2)
    encode (TmAdd t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkAdd [a1, a2]
    encode (TmSub t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkSub [a1, a2]
    encode (TmMul t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkMul [a1, a2]
    encode (TmDiv t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkDiv a1 a2
    encode (TmEqual t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkEq a1 a2
    encode (TmAnd t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkAnd [a1, a2]
    encode (TmOr t1 t2) = do
        a1 <- encode t1
        a2 <- encode t2
        mkOr [a1, a2]
    encode (TmNot t) = encode t >>= mkNot
    encode (TmIf p c a) = do
        a1 <- encode p
        a2 <- encode c
        a3 <- encode a
        mkIte a1 a2 a3

instance Z3Sorted Term where
    sort (TmVar x) = do
        ctx <- getQualifierCtx
        case M.lookup x ctx of
            Just (_, s) -> return s
            Nothing -> smtError $ "Can't find variable " ++ x
    sort (TmNum _) = mkIntSort
    sort (TmLT _ _) = mkBoolSort
    sort (TmGT _ _) = mkBoolSort
    sort (TmLE _ _) = mkBoolSort
    sort (TmGE _ _) = mkBoolSort
    sort (TmAdd _ _) = mkIntSort
    sort (TmSub _ _) = mkIntSort
    sort (TmMul _ _) = mkIntSort
    sort (TmDiv _ _) = mkIntSort
    sort (TmEqual _ _) = mkBoolSort
    sort (TmAnd _ _) = mkBoolSort
    sort (TmOr _ _) = mkBoolSort
    sort (TmNot _) = mkBoolSort
    sort (TmIf _ c _) = sort c

instance Z3Sorted RType where
    sort RTInt  = mkIntSort
