{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Solver.Test where

import Control.Arrow (second)
import Control.Monad (Monad(..))
import Data.Bool
import Data.Either (Either(..))
import Data.Function (($), (.))
import Data.Functor ((<$>))
-- import Data.Functor.Identity (Identity(..))
import Data.Int (Int)
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import Data.Maybe ({-Maybe(..),-} fromMaybe)
import Data.Monoid ((<>))
import Data.Sequence (Seq)
import Data.String (IsString(..))
-- import Data.Text (Text)
import qualified Data.Text as Text
import Data.Text.Buildable (Buildable(..))
-- import qualified Data.Text.IO as Text
import qualified Data.Text.Lazy.Builder as Build
import qualified Data.Text.Lazy.IO as TL
-- import Data.Tuple (fst)
import Prelude (Num(..), error)
import System.IO (IO)
-- import Test.HUnit hiding (test)
import Test.Tasty
import Test.Tasty.HUnit
import Text.Show (Show(..))

import Language.LOL.Typing.Type
import Language.LOL.Typing.Expr
import Language.LOL.Typing.Solver
import Language.LOL.Typing.Collect.Constraint
-- import Language.LOL.Typing.Collect.Grammar
import qualified Language.LOL.Typing.Collect.Grammar as Collect
import Language.LOL.Typing.Lib.Data.Default (Default(..))
import qualified Language.LOL.Typing.Lib.Control.Monad.Classes.Instance as MC
import qualified Language.LOL.Typing.Lib.Data.Text.Buildable as Build

-- Convenient instances
instance IsString Polytype where
	fromString = polytype . Monotype_Const . fromString
instance Num Monotype where
	(+) = Monotype_App
	(*) = Monotype_App
	abs x = x
	signum _ = 0
	negate x = x
	fromInteger = Monotype_Var . fromInteger

type Test_Solver
 =   Solver_Greedy_Finite Collect_Infos IO

{-
polytype_of_expr
 :: Expr
 -> Either (Collect_Error (Info  Test_Solver)
                          (Error Test_Solver))
           Text
polytype_of_expr ex = do
	-- Text.putStrLn (Text.pack $ show ex)
	let Syn_Expr
		 { constraints_Syn_Expr = cs
		 , freshvar_Syn_Expr    = freshvar
		 , monoty_Syn_Expr      = monotype
		 , errors_Syn_Expr      = errors
		 } = wrap_Expr (sem_Expr ex) Inh_Expr
		 { polytys_Inh_Expr  = -- Map.fromList
			 [ ("eq", polytyref
				(Quantification [1] [(1, "a")] $
				[Class_Qualifier "Eq" 1]
				 .=>. (1.->.1.->."Bool")
				 :: Polytype))
			 ]
		 , freshvar_Inh_Expr = 0
		 }
	case errors of
	 (_:_) -> Left $ Collect_Error_Grammar errors
	 [] ->
		let cs' = {-hussel-} cs in
		let cfg = solver_config
			 { solver_config_freshvar = freshvar
			 , solver_config_check    = True
			 } in
		let (result :: Solver_Result (Info Test_Solver) (Error Test_Solver)) =
			runIdentity $
			solve (\_ -> return ()) $
			solver_greedy_finite cfg cs' in
		case solver_result_errors result of
		 [] -> Right $ Build.text $
			(def::Quantification_Build_Options,) $
			for_all $ solver_result_monotypes result `substitute` monotype
		 errs -> Left $ Collect_Error_Solver errs

-- | Like 'polytype_of_expr', but more verbose;
-- mainly for debugging.
debug_expr :: Expr -> IO Text
debug_expr ex = do
	Text.putStrLn (Text.pack $ Build.string ex)
	let Syn_Expr
		 { constraints_Syn_Expr = cs
		 , freshvar_Syn_Expr    = freshvar
		 , monoty_Syn_Expr      = monotype
		 , errors_Syn_Expr      = errors
		 } = wrap_Expr (sem_Expr ex) Inh_Expr
		 { polytys_Inh_Expr  = mempty
		 , freshvar_Inh_Expr = 0
		 }
	case errors of
	 (_:_) -> return $ Build.text $
		show (Collect_Error_Grammar errors::Collect_Error Collect_Info (Error Test_Solver))
	 [] -> do
		let cs' = {-hussel-} cs
		let cfg = solver_config
			 { solver_config_freshvar = freshvar
			 , solver_config_check    = True
			 }
		(result :: Solver_Result (Info  (Solver_Greedy_Finite Collect_Infos IO))
		                         (Error (Solver_Greedy_Finite Collect_Infos IO))) <-
			solve write_log $
			solver_greedy_finite cfg cs'
			 -- (\_ -> return ())
		return $
			case solver_result_errors result of
			 [] -> Build.text $
				(def::Quantification_Build_Options,) $
				for_all $ solver_result_monotypes result `substitute` monotype
			 errs -> Build.text $ show $ Collect_Error_Solver errs
-}

write_log :: Solver_Log -> IO ()
write_log l =
	case l of
	 Solver_Log_Constraint x ->
		TL.putStrLn $ Build.toLazyText $ "- " <> build x
	 Solver_Log_Class x ->
		TL.putStrLn $ Build.toLazyText $ "- " <> build x
	 Solver_Log_Polytype x ->
		TL.putStrLn $ Build.toLazyText $ "- " <> build x
	 Solver_Log_States states ->
		TL.putStrLn $ Build.toLazyText $
		Build.unlines $
			[ hline ] <>
			List.zipWith
			 (\i (MC.Instance s) ->
				build i <> ". " <> build (state_name s) <> "\n" <>
				Build.indent "    " (Build.unlines $
				 (build <$> state_options s) <>
				 [build $ state_show s]))
			 [1::Int ..]
			 states <>
			[ hline ]
	where
		hline :: Build.Builder = build $ Text.replicate 80 "-"

polytys_env :: [(Name, Polytype)]
polytys_env =
	[ ("equal", polytype
		(Quantification [1] [(1, "a")] $
		[Class_Qualifier "Eq" 1]
		 .=>. (1.->.1.->."Bool")
		 :: Polytype))
	, ("compare", polytype
		(Quantification [1] [(1, "a")] $
		[Class_Qualifier "Ord" 1]
		 .=>. (1.->.1.->."Ordering")
		 :: Polytype))
	, ("True", polytype
		(Quantification [] [] $
		[] .=>. "Bool"
		 :: Polytype))
	, ("one", polytype
		(Quantification [] [] $
		[] .=>. "Int"
		 :: Polytype))
	, ("pair", polytype
		(Quantification [1,2] [] $
		[] .=>. (1.->.2.->.type_Tuple [1,2])
		 :: Polytype))
	]


class Collect_Constraints a where
	collect_constraints
	 :: a -> Either (Collect_Error (Info Test_Solver) (Error Test_Solver))
	                (Freshvar, Seq Collect_Constraint)
instance Collect_Constraints Expr where
	collect_constraints expr =
		let Collect.Syn_Expr
			 { Collect.freshvar_Syn_Expr    = freshvar
			 , Collect.constraints_Syn_Expr = constrs
			 , Collect.errors_Syn_Expr      = errors
			 } = Collect.wrap_Expr (Collect.sem_Expr expr)
			 Collect.Inh_Expr
			 { Collect.polytys_Inh_Expr  = (polytyref `second`) <$> polytys_env
			 , Collect.freshvar_Inh_Expr = 0
			 } in
		case errors of
		 errs@(_:_) -> Left $ Collect_Error_Grammar errs
		 [] -> Right (freshvar, constrs)
instance Collect_Constraints Decl where
	collect_constraints decl =
		let Collect.Syn_Decl
			 { Collect.freshvar_Syn_Decl    = freshvar
			 , Collect.constraints_Syn_Decl = constrs
			 , Collect.errors_Syn_Decl      = errors
			 } = Collect.wrap_Decl (Collect.sem_Decl decl)
			 Collect.Inh_Decl
			 { Collect.polytys_Inh_Decl  = (polytyref `second`) <$> polytys_env
			 , Collect.freshvar_Inh_Decl = 0
			 } in
		case errors of
		 errs@(_:_) -> Left $ Collect_Error_Grammar errs
		 [] -> Right (freshvar, constrs)

infer :: Collect_Constraints a
 => Bool -> a
 -> IO (Either (Collect_Error (Info Test_Solver) (Error Test_Solver))
               (Solver_Result (Info Test_Solver) (Error Test_Solver)))
infer logging x =
	case collect_constraints x of
	 Left err -> return $ Left err
	 Right (solver_config_freshvar, constrs) -> do
		let cfg = solver_config
			 { solver_config_freshvar
			 , solver_config_check = True
			 }
		let wlog = if logging then write_log else (\_ -> return ())
		res <- solve wlog $ solver_greedy_finite cfg constrs
		return $
			case solver_result_errors res of
			 [] -> Right res
			 errs -> Left $ Collect_Error_Solver errs

class Collect_Constraints a => Infer_Polytype a where
	infer_polytype
	 :: Bool -> a
	 -> IO (Either (Collect_Error (Info Test_Solver) (Error Test_Solver))
	               Polytype)
instance Infer_Polytype Expr where
	infer_polytype logging ex = do
		res <- infer logging ex
		case res of
		 Left err -> return $ Left err
		 Right Solver_Result
			 { solver_result_monotypes  = monosub
			 , solver_result_qualifiers = quals
			 } -> return $ Right $
				for_all $ (quals .=>.) $
				(monosub `substitute`) $
				Monotype_Var 0
instance Infer_Polytype Decl where
	infer_polytype logging ex = do
		res <- infer logging ex
		case res of
		 Left err -> return $ Left err
		 Right Solver_Result
			 { solver_result_monotypes = monosub
			 , solver_result_polytypes = polysub
			 } -> return $ Right $
				(monosub `substitute`) $
				fromMaybe (error $ "Oops, Polytype missing" <> show polysub) $
				Map.lookup 0 polysub


tests :: TestTree
tests = testGroup "Solver"
 [ testGroup "Expr" $
		let test logging input expected = do
			p <- infer_polytype logging input
			let got = (Build.text . (def::Quantification_Build_Options,)) <$> p
			got @?= expected in
		let (==>) = test False in
		-- let (==>>) = test True in
	 [ testGroup "SKI" $
			-- DOC: https://en.wikipedia.org/wiki/SKI_combinator_calculus
			-- DOC: https://en.wikipedia.org/wiki/B,_C,_K,_W_system
			let i = "x".-> "x" in
			let k = "x".-> "y".-> "x" in
			let s = "x".-> "y".-> "z".-> ("x"!"z") ! ("y"!"z") in
			let c = "f".-> "x".-> "y".-> "f"!"y"!"x" in
			let b = "f".-> "g".-> "x".-> "f"!("g"!"x") in
			let w = "x".-> "y".-> "x"!"y"!"y" in
			let t = k in -- True
			let f = s ! k in -- False
			let no = f ! t in -- not
		 [ testCase "I" $ i
			 ==> Right "forall a. a -> a"
		 , testCase "K" $ k
			 ==> Right "forall a b. a -> b -> a"
		 , testCase "S" $ s
			 ==> Right "forall a b c. (a -> b -> c) -> (a -> b) -> a -> c"
		 , testCase "B" $ b
			 ==> Right "forall a b c. (a -> b) -> (c -> a) -> c -> b"
		 , testCase "B = S (K S) K" $
			 (s ! (k!s) ! k)
			 ==> Right "forall a b c. (a -> b) -> (c -> a) -> c -> b"
		 , testCase "C" $ c
			 ==> Right "forall a b c. (a -> b -> c) -> b -> a -> c"
		 , testCase "C = S (B B S) (K K)" $
			 (s ! (b!b!s) ! (k!k))
			 ==> Right "forall a b c. (a -> b -> c) -> b -> a -> c"
		 , testCase "W" $ w
			 ==> Right "forall a b. (a -> a -> b) -> a -> b"
		 , testCase "W = S S (S K)" $
			 (s ! s ! (s!k))
			 ==> Right "forall a b. (a -> a -> b) -> a -> b"
		 , testCase "W /= S I I" $
			 (s!i!i)
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos
					 [ Collect_Info_Solver (Solver_Info_Monotype $ Info_Monotype_Unification 1 (18 .->. 0))
					 , Collect_Info_Grammar (Collect_Info_Grammar_Expr "App") ]
				 , Solver_Error_Monotype $
					Error_Monotype_Unification $
					Unification_Error_Infinite_type 12
				 )])
		 , testCase "W W /= W W" $
			 (w ! w)
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos
					 [ Collect_Info_Solver (Solver_Info_Monotype $ Info_Monotype_Unification 1 (10 .->. 0))
					 , Collect_Info_Grammar (Collect_Info_Grammar_Expr "App") ]
				 , Solver_Error_Monotype $
					Error_Monotype_Unification $
					Unification_Error_Infinite_type 13
				 )])
		 , testCase "S = B (B W) (B B C)" $
			 (b ! (b!w) ! (b!b!c))
			 ==> Right "forall a b c. (a -> b -> c) -> (a -> b) -> a -> c"
		 , testCase "I = S K K" $
			 (s!k!k)
			 ==> Right "forall a. a -> a"
		 , testCase "I = W K" $
			 (w ! k)
			 ==> Right "forall a. a -> a"
		 , testCase "not" $ no
			 ==> Right "forall a. a -> a"
		 ]
	 , testGroup "Abst" $
			let o = "f".-> "g".-> "x".-> "f"!("g"!"x") in
		 [ testCase "(x:Int) -> x" $
			 (("x", "Int"::Monotype)..-> "x")
			 ==> Right "Int -> Int"
		 , testCase "(x:Int) -> (x:Bool)" $
			 (("x", "Int"::Monotype)..-> ("x".:("Bool"::Monotype)))
			 ==> Left (Collect_Error_Solver
				 [ ( Collect_Infos
					 [ Collect_Info_Solver (Solver_Info_Monotype (Info_Monotype_Unification 2 "Bool"))
					 , Collect_Info_Solver (Solver_Info_Polytype (Info_Polytype_Rigidified [1] "Bool"))
					 , Collect_Info_Grammar (Collect_Info_Grammar_Expr "Annot Bool") ]
				 , Solver_Error_Monotype (Error_Monotype_Unification (Unification_Error_Constant_clash "Int" "Bool")))
				 ])
		 , testCase "(x:Int -> Int) -> x" $
			 (("x", "Int".->."Int")..-> "x")
			 ==> Right "(Int -> Int) -> Int -> Int"
		 , testCase "x -> y" $
			 ("x".-> "y")
			 ==> Left (Collect_Error_Grammar [Collect_Error_Grammar_Variable_not_in_scope "y"])
		 , testCase "compose" $ o
			 ==> Right "forall a b c. (a -> b) -> (c -> a) -> c -> b"
		 , testCase "twice" $
			 ("f".-> "x".-> "f"!("f"!"x"))
			 ==> Right "forall a. (a -> a) -> a -> a"
		 , testCase "twice = \\f -> f . f" $
			 ("f".-> "c" .= o $ "c"!"f"!"f")
			 ==> Right "forall a. (a -> a) -> a -> a"
		 , testCase "(x y) z" $
			 ("x".-> "y".-> "z".-> "x"!"y"!"z")
			 ==> Right "forall a b c. (a -> b -> c) -> a -> b -> c"
		 , testCase "x (y z)" $
			 ("x".-> "y".-> "z".-> "x"!("y"!"z"))
			 ==> Right "forall a b c. (a -> b) -> (c -> a) -> c -> b"
		 , testCase "x u v t" $
			 ("x".-> "y".-> "z".-> "t".->
				"y" ! ("u".-> ("t" ! (("z"!"u") ! ("v".-> "x"!"u"!"v"!"t")))))
			 ==> Right "forall a b c d e f. (a -> b -> (c -> d) -> e) -> ((a -> d) -> f) -> (a -> (b -> e) -> c) -> (c -> d) -> f"
		 , testCase "x -> x x" $
			 ("x".-> "x"!"x")
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos
					 [ Collect_Info_Solver (Solver_Info_Monotype $ Info_Monotype_Unification 3 (4 .->. 2))
					 , Collect_Info_Grammar (Collect_Info_Grammar_Expr "App") ]
				 , Solver_Error_Monotype $
					Error_Monotype_Unification $
					Unification_Error_Infinite_type 1
				 )])
		 , testCase "y -> (x -> x) y y" $
			 ("y".-> ("x".-> "x")!"y"!"y")
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos
					 [ Collect_Info_Solver (Solver_Info_Monotype $ Info_Monotype_Unification 3 (8 .->. 2))
					 , Collect_Info_Grammar (Collect_Info_Grammar_Expr "App") ]
				 , Solver_Error_Monotype $
					Error_Monotype_Unification $
					Unification_Error_Infinite_type 1
				 )])
		 , testCase "x -> (a b c -> a c) (x -> x) (f -> f x x)" $
			 ("x".-> ("a".-> "b".-> "c".-> "a"!"c") ! ("x".-> "x") ! ("f".-> "f"!"x"!"x"))
			 ==> Right "forall a b. a -> b -> b"
		 ]
	 , testGroup "Let"
		 [ testCase "top_example_4_7" $
			 ("x".->
					"f".= ("g".= "x" $ "a".-> "g") $
						"h".= "f" $ "h"!"x")
			 ==> Right "forall a. a -> a"
		 , testCase "(i:Int) -> let (f:Int) = i in f" $
			 (("i", "Int"::Monotype)..->
				(("f", "Int"::Monotype)..= "i") "f")
			 ==> Right "Int -> Int"
		 , testCase "i -> let (f:forall a. a -> a) = i in f" $
			 ("i".-> (("f", for_all ([] .=>. (1 .->. 1)))..= "i") "f")
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos [ Collect_Info_Grammar (Collect_Info_Grammar_Expr "Let(sig) forall a. a -> a") ]
				 , Solver_Error_Polytype (Error_Polytype_Rigid_escaping [6])
				 )])
		 , testCase "(i:Int) -> let (f:Bool) = i in f" $
			 (("i", "Int"::Monotype)..->
				(("f", "Bool"::Monotype)..= "i") "f")
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos
					 [ Collect_Info_Solver (Solver_Info_Monotype (Info_Monotype_Unification 4 (Monotype_Const "Bool")))
					 , Collect_Info_Solver (Solver_Info_Polytype (Info_Polytype_Rigidified [1] "Bool"))
					 , Collect_Info_Grammar (Collect_Info_Grammar_Expr "Let(sig) Bool") ]
				 , Solver_Error_Monotype $
					Error_Monotype_Unification $
					Unification_Error_Constant_clash "Int" "Bool"
				 )])
		 , testCase "a -> b -> a" $
			 ("x".-> "f".= ("y".->"x") $ "f")
			 ==> Right "forall a b. a -> b -> a"
		 ]
	 , testGroup "Annot" $
			let id = "x".-> "x" in
		 [ testCase "id : a -> a" $
			 (id .: for_all ([] .=>. (1 .->. 1)))
			 ==> Right "forall a. a -> a"
		 , testCase "id : Int -> Int" $
			 (id .: for_all ([] .=>. ("Int" .->. "Int")))
			 ==> Right "Int -> Int"
		 , testCase "id /: a -> Int" $
			 (id .: for_all ([] .=>. (1 .->. "Int")))
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos [ Collect_Info_Grammar (Collect_Info_Grammar_Expr "Annot forall a. a -> Int") ]
				 , Solver_Error_Polytype $
					Error_Polytype_Rigid_type_mismatch
					[(3, Monotype_Const "Int")] )
				 ])
		 , testCase "id /: a -> b" $
			 (id .: for_all ([] .=>. (1 .->. 2)))
			 ==> Left (Collect_Error_Solver
				 [( Collect_Infos [ Collect_Info_Grammar (Collect_Info_Grammar_Expr "Annot forall a b. a -> b") ]
				 , Solver_Error_Polytype $
					Error_Polytype_Rigid_injectivity_lost $
					Map.fromList [(3, [3, 4])] )
				 ])
		 ]
	 {-
	 , testGroup "Class" $
		let (==>) input (expected::Text) = do
			got <- ((Build.text . (def::Quantification_Build_Options,)) <$>)
				 <$> polytype_of_expr input
			-- Text.putStrLn ("exp: " <> Build.text term)
			got @?= Right expected in
		 [ testCase "Eq a => a -> Bool" $
			("x".-> "=="!"x"!"x")
			-- ("f".= ("x".-> "=="!"x"!"x") "f")
			 ==> "forall a. Eq a => a -> Bool"
		 ]
	 , testGroup "Class" $
		let test logging input (expected::Text) = do
			p <- polytype_of_decl logging input
			let got = (Build.text . (def::Quantification_Build_Options,)) <$> p
			got @?= Right expected in
		let (==>) = test False in
		let (==>>) = test True in
		 [ testCase "Eq a => a -> a -> Bool" $
			Decl_Let Nothing "f" "=="
			 ==> "forall a. Eq a => a -> a -> Bool"
		 , testCase "Eq a => a -> Bool" $
			Decl_Let Nothing "f"
			 ("x".-> "=="!"x"!"x")
			 ==> "forall a. Eq a => a -> Bool"
		 , testCase "Eq a => a -> Bool" $
			Decl_Let Nothing "f"
			 ("x".-> "=="!("=="!"x"!"x")!("=="!"x"!"x"))
			 ==> "forall a. Eq a => a -> Bool"
		 , testCase "Ord a => a -> Bool" $
			Decl_Let Nothing "f" ("x".->
				"=="
				 !("=="!"x"!"x")
				 !("=="
					 !("compare"!"x"!"x")
					 !("compare"!"x"!"x")))
			 ==> "forall a. Ord a => a -> Bool"
		 , testCase "a -> b -> a" $
			 Decl_Let Nothing "f"
			 ("x".-> "g".= ("y".->"x") $ "g")
			 ==>> "forall a b. a -> b -> a"
		 ]
	-} ]
 ]