{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE NamedFieldPuns         #-}
{-# LANGUAGE OverlappingInstances   #-}
{-# LANGUAGE QuasiQuotes            #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeSynonymInstances   #-}
{-# LANGUAGE UndecidableInstances   #-}
{-# LANGUAGE ViewPatterns           #-}

module Forml.TypeCheck where

import           Data.List                     (intersect, partition,
                                                union, (\\), reverse)

import qualified Data.List                     as L
import qualified Data.Map                      as M
import qualified Data.Set                      as S

import           Text.ParserCombinators.Parsec

import           Forml.Types.Axiom
import           Forml.Types.Definition
import           Forml.Types.Expression
import           Forml.Types.Namespace        hiding (Module)
import           Forml.Types.Pattern
import           Forml.Types.Statement        hiding (Test, find, modules,
                                                namespace)
import           Forml.Types.Symbol

import           Forml.Types.Type
import           Forml.Types.TypeDefinition

import           Forml.Parser
import           Forml.Parser.Utils

import           Forml.TypeCheck.Types

import           Forml.Deps

import Data.Graph (graphFromEdges, SCC(..))
import Data.Graph.SCC (sccList)


-- Type Inference
-- --------------------------------------------------------------------------------


instance Infer (Expression Definition) Type where
    infer (ApplyExpression _ []) = fail "This should not be"
    infer (ApplyExpression e (x:[])) =

        do te <- infer e
           tx <- infer x
           t  <- newTVar Star
           unify (tx `fn` t) te
           return t

    infer (ApplyExpression e (x:xs)) = infer (ApplyExpression (ApplyExpression e (x:[])) xs)

    infer (IfExpression a b (Just c)) =

        do ta <- infer a
           tb <- infer b
           tc <- infer c
           t  <- newTVar Star
           unify ta bool_type
           unify t tb
           unify t tc
           return t

    infer (IfExpression a b Nothing) =

        do ta <- infer a
           tb <- infer b
           t  <- newTVar Star
           unify ta bool_type
           unify t tb
           is_js <- t `can_unify` TypeRecord (TRecord M.empty TComplete Star)
           if is_js
               then t `unify` TypeRecord (TRecord M.empty TComplete Star)
               else t `unify` TypeApplication (Type (TypeConst "JS" (KindFunction Star Star))) (TypeRecord (TRecord M.empty TComplete Star))
           return t


    infer (LiteralExpression s) = infer s

    infer (SymbolExpression i) =

        do sc <- find (show i)
           t <- freshInst sc
           return t

    infer (JSExpression _) =

        do t <- newTVar Star
           return (TypeApplication (Type (TypeConst "JS" (KindFunction Star Star))) t)

    infer (LazyExpression x _) =

        do t <- newTVar Star
           t' <- infer x
           unify t t'
           return (TypeApplication (Type (TypeConst "JS" (KindFunction Star Star))) t)


    -- TODO this may be removeable at no perf cost?
    infer (FunctionExpression rs) =

        do t <- newTVar Star
           as <- get_assumptions
           [_ :>: q] <- with_scope$ ims as
           t' <- freshInst q
           unify t t'
           return t

        where ims as =

                  do infer [Definition Public False (Symbol "") rs]
                     as'' <- get_assumptions
                     return$ as'' \\ as

    infer (AccessorExpression (Addr s f x) y) = infer (acc $ reverse y)

        where acc :: [Symbol] -> Expression Definition
              acc [] = x
              acc (y:ys) = --Addr undefined undefined $
                  ApplyExpression
                  (FunctionExpression
                   [ EqualityAxiom
                     (Match [RecordPattern (M.fromList [(y, VarPattern "__x__")]) Partial] Nothing)
                     (Addr s f (SymbolExpression (Symbol "__x__"))) ] )
                  [acc ys]

    infer (RecordExpression (unzip . M.toList -> (names, xs))) =

        do ts <- mapM infer xs
           let r = TypeRecord (TRecord (M.fromList (zip (map f names) ts)) TComplete Star)
           t' <- newTVar Star
           sc <- find $ quantify (tv r) r
           case sc of
             Nothing ->
                 do unify t' r
                    return t'
             Just (Forall _ scr, sct) ->
                 do t''  <- freshInst sct
                    t''' <- return$ inst (map TypeVar$ tv t'') (scr)
                    t    <- freshInst (quantify (tv t''' \\ tv t'') t''')
                    unify t r
                    unify t' t''
                    return t'

        where f (Symbol x) = x
              f (Operator x) = x

    infer (LetExpression xs x) =

        with_scope$ do mapM_ ((>>= assume) . infer) defs
                       infer x

        where defs = to_group (map DefinitionStatement xs)

    infer (ListExpression x) =

        do t <- newTVar Star
           ts <- mapM infer x
           mapM (unify t) ts
           t' <- newTVar Star
           unify t' (TypeApplication (Type (TypeConst "Array" (KindFunction Star Star))) t)
           return t'
           
    infer x = error $ "Unimplemented: " ++ show x

-- Axioms

instance (Infer a t) => Infer (Addr a) t where

    infer (Addr s _ x) = do m <- get_msg
                            set_msg new_msg
                            z <- infer x
                            set_msg m
                            return z

        where new_msg = "  at line " ++ show (sourceLine s) ++ ", column " ++ show (sourceColumn s) ++ "\n"

instance Infer (Axiom (Expression Definition)) Type where

    infer (EqualityAxiom (Match y z) x) =

        do ts <- mapM infer y
           case z of
             (Just q) -> infer q >>= (flip unify) bool_type
             _ -> return ()
           t  <- infer x
           return (foldr fn t ts)

    infer _ = newTVar Star





instance Infer [Definition] () where
    infer bs =

        do def_types <- mapM (\_ -> newTVar Star) bs

           let is    = map get_name bs
               scs   = map toScheme def_types
               altss = map get_axioms bs

           axiom_types <- with_scope$
                do assume $ zipWith (:>:) is scs
                   mapM (mapM (with_scope . infer)) altss

           let f _ []     = return ()
               f g (x:xs) = do s <- get_substitution
                               g x
                               g (apply s x)
                               f g xs

           mapM (\(t, as) -> f (unify t) as) (zip def_types axiom_types)

           as  <- get_assumptions
           ss  <- get_substitution
           fs' <- substitute as

           let ts' = apply ss def_types
               fs  = tv fs'
               vss = map tv ts'
               gs  = foldr1 union vss \\ fs

           if restricted then
               let scs' = map (quantify gs) ts'
               in do assume (zipWith (:>:) is scs')
                     return ()
             else
               let scs' = map (quantify gs) ts'
               in do assume (zipWith (:>:) is scs')
                     return ()

        where get_name (Definition _ _ (Symbol x) _) = x
              get_name (Definition _ _ (Operator x) _) = x
              get_axioms (Definition _ _ _ x) = x

              restricted = any simple bs
              simple (Definition _ _ _ axs) = any (null . f) axs

              f (EqualityAxiom (Match p _) _) = p
              f _ = error "Fatal error occurred while reticulating splines"


instance Infer Definition () where

    infer (Definition _ _ name axs) =

        do sc <- find$ f name
           t  <- freshInst sc
           axiom_types <- with_scope$ mapM (with_scope . infer) axs

           s <- get_substitution
           mapM (flip unify t) axiom_types  -- TODO apply sub to axiom_types?

           as <- get_assumptions

           let t'  = apply s t
               fs  = tv (apply s as)
               gs  = tv t' \\ fs
               sc' = quantify gs t'


           if sc /= sc' then
               add_error$ "Signature too general\n\n    Expected: " ++ show sc ++ "\n    Actual: " ++ show sc'
             else
               assume (f name :>: sc)

           return ()

        where f (Symbol x) = x
              f (Operator x) = x

instance Infer Test () where

    infer (Test ex) = do t <- newTVar Star
                         x <- infer ex
                         unify t x
                         unify t bool_type

newtype Test = Test (Addr (Expression Definition))

instance Infer BindGroup [Assumption] where
    infer (Scope imps tts es iss ts) =

        do --as <- get_assumptions
           mapM import' imps
           as' <- get_assumptions
           infer tts
           mapM assume$ sigs es
           mapM infer iss
           with_scope$ mapM infer es
           mapM infer (map Test ts)
           as'' <- get_assumptions
           set_assumptions as'
           return (as'' \\ as')

        where f (TypeAxiom t) = True
              f _ = False

              g name (TypeAxiom t) = [ name :>: to_scheme' t' | t' <- enumerate_types t ]

              to_scheme' :: Type -> Scheme
              to_scheme' t = quantify (tv t) t

              sigs :: [Definition] -> [Assumption]
              sigs [] = []
              sigs (Definition _ _ name as:xs) =
                  case L.find f as of
                    Nothing -> sigs xs
                    Just x -> g (h name) x ++ sigs xs

              import' (Namespace ns, Nothing) =

                  do z <- get_modules
                     a <- get_assumptions
                     (Namespace ns') <- get_namespace
                     case Namespace ns `lookup` z of
                       Just z' -> assume$ a ++ z'
                       Nothing -> if length ns' > 0 && head ns' /= head ns
                                  then import' (Namespace (head ns' : ns), Nothing)
                                  else add_error$ "Unknown namespace " ++ show (Namespace ns)

              import' (Namespace ns, Just alias) =

                  do z <- get_modules
                     a <- get_assumptions
                     (Namespace ns') <- get_namespace
                     case Namespace ns `lookup` z of
                       Just z' -> 
                           do record <- to_record z'
                              assume $ alias :>: record
                       Nothing -> if length ns' > 0 && head ns' /= head ns
                                  then import' (Namespace (head ns' : ns), Just alias)
                                  else add_error$ "Unknown namespace " ++ show (Namespace ns)
                                  
              to_record assumptions =
              
                  let f (_ :>: scheme) =

                          [ do t <- freshInst scheme
                               return t ]

                      f _ = []
                      
                      g (i :>: _) = [i]
                      g _ = []
                      
                  in  do schemes <- sequence $ concat $ map f assumptions
                         let symbols = concat $ map g assumptions 
                         let rec = TypeRecord (TRecord (M.fromList (zip symbols schemes)) TComplete Star)
                         return $ quantify (tv rec) rec


              h (Symbol x) = x
              h (Operator x) = x

    infer (Module name bgs) =

        do as <- get_assumptions
           with_module name$ infer' bgs
           set_assumptions as
           return []

         where infer' [] = return []
               infer' (x:xs) =

                  do a <- infer x
                     assume a
                     as' <- infer' xs
                     return$ a ++ as'



to_scheme :: TypeDefinition -> UnionType -> [Assumption]
to_scheme (TypeDefinition n vs) t = [ quantify (vars y) y :>>: def_type y
                                          | y <- enumerate_types t ]

    where vars y = map (\x -> TVar x (infer_kind x y)) vs

          def_type y = quantify (vars y) (foldl app poly_type (map TypeVar (vars y)))

          poly_type = Type (TypeConst n (to_kind (length vs)))

          to_kind 0 = Star
          to_kind n = KindFunction Star (to_kind$ n - 1)

          app :: Type -> Type -> Type
          app y x = TypeApplication y x

          -- TODO this is still wrong - have to check for all enumerated types

          infer_kind x y = let ks = infer_kinds x y
                           in if ks == []
                              then Star
                              else if all (\x -> x == head ks) ks
                                   then head ks
                                   else error "Kind mismatch in scheme"

          infer_kinds x (TypeApplication a b) = infer_kinds x a ++ infer_kinds x b
          infer_kinds x (TypeVar (TVar y k)) | x == y = [k]
          infer_kinds x (TypeRecord (TRecord m _ _)) = concat$ map (infer_kinds x) (M.elems m)
          infer_kinds _ _ = []

-- | Computes all possible types from a type signature AST.

enumerate_types :: UnionType -> [Type]
enumerate_types (UnionType types) = to_unit . concat . map enumerate_type . S.toList $ types

    where term_type (VariableType x)      = [ TypeVar (TVar x Star) ]
          term_type (SymbolType x)        = [ Type (TypeConst (show x) Star) ]
          term_type (PolymorphicType a b) = [ foldl TypeApplication a' b'
                                                  | b' <- map enumerate_types b
                                                  , a' <- to_kind' (length b')$ term_type a ]

          to_kind 0 = Star
          to_kind n = KindFunction Star (to_kind$ n - 1)

          to_unit [] = [TypeRecord (TRecord M.empty TComplete Star)]
          to_unit x = x

          to_kind' _ [] = []
          to_kind' n (TypeVar (TVar x _) : xs) = TypeVar (TVar x (to_kind n)) : to_kind' n xs
          to_kind' n (Type (TypeConst x _) : xs) = Type (TypeConst x (to_kind n)) : to_kind' n xs

          enumerate_type (SimpleType x) = term_type x

          enumerate_type (FunctionType a b) =
              [ a' `fn` b' | a' <- enumerate_types a, b' <- enumerate_types b ]

          enumerate_type (RecordType (unzip . M.toList -> (names, types'))) =

              map f permutations

              where f = TypeRecord . (\x -> TRecord x TComplete Star) . M.fromList . zip (map show names)
                    permutations = permutations' . map enumerate_types $ types'

                        where permutations' [] = []
                              permutations' (x:[]) = [ x ]
                              permutations' (x:xs) = [ x' : xs' | x' <- x, xs' <- permutations' xs ]





instance Infer [Statement] () where

    infer [] = return ()
    infer (TypeStatement t c : xs) =

        do assume     $ to_scheme t c
           infer xs

    infer (_ : xs) = infer xs


js_type :: Type
js_type = Type (TypeConst "JS" (KindFunction Star Star))

tiProgram :: Program -> [(Namespace, [Assumption])] -> ([(Namespace, [Assumption])], [String])
tiProgram (Program bgs) env =

    runTI $ do TI (\x -> (x { modules = env }, ()))
               assume$ "true"  :>: (Forall [] (Type (TypeConst "Bool" Star)))
               assume$ "false" :>: (Forall [] (Type (TypeConst "Bool" Star)))
               assume$ "error" :>: (Forall [Star] (TypeGen 0))
               assume$ "run"   :>: (Forall [Star] (TypeApplication js_type (TypeGen 0) -:> TypeGen 0))
               infer'$ to_group bgs
               s  <- get_substitution
               ms <- TI (\y -> (y, modules y))
               e  <- get_errors
               return ((apply s ms), S.toList . S.fromList $ e)

    where infer' [] = return ()
          infer' (x:xs) =

              do a <- infer x
                 assume a
                 infer' xs