module CheckTests (tests) where

import Check
import qualified Data.Map as M
import Expr (Expr (..))
import Test.HUnit

sort :: Test
sort = TestCase $ assertEqual "*" (Right Square) (checkType M.empty [] Star)

stlc :: Test
stlc =
    TestCase $
        assertEqual
            "fun (x : A) (y : B) . x"
            (Right $ Pi "" (Var 0 "A") (Pi "" (Var 2 "B") (Var 2 "A")))
            (checkType M.empty [Star, Star] $ Abs "x" (Var 0 "A") (Abs "y" (Var 2 "B") (Var 1 "x")))

freeVar :: Test
freeVar =
    TestCase $
        assertEqual "{x = *} , [] |- x : □" (Right Square) (checkType (M.singleton "x" Star) [] (Free "x"))

polyIdent :: Test
polyIdent =
    TestCase $
        assertEqual
            "fun (A : *) (x : A) . x"
            (Right $ Pi "A" Star (Pi "" (Var 0 "A") (Var 1 "A")))
            (checkType M.empty [] (Abs "A" Star (Abs "x" (Var 0 "A") (Var 0 "x"))))

typeCons :: Test
typeCons =
    TestCase $
        assertEqual
            "fun (A B : *) . A -> B"
            (Right $ Pi "" Star (Pi "" Star Star))
            (checkType M.empty [] (Abs "A" Star (Abs "B" Star (Pi "" (Var 1 "A") (Var 1 "B")))))

useTypeCons :: Test
useTypeCons =
    TestCase $
        assertEqual
            "fun (C : * -> *) (A : *) (x : C A) . x"
            (Right $ Pi "C" (Pi "" Star Star) (Pi "A" Star (Pi "" (App (Var 1 "C") (Var 0 "A")) (App (Var 2 "C") (Var 1 "A")))))
            (checkType M.empty [] $ Abs "C" (Pi "" Star Star) (Abs "A" Star (Abs "x" (App (Var 1 "C") (Var 0 "A")) (Var 0 "x"))))

dependent :: Test
dependent =
    TestCase $
        assertEqual
            "fun (S : *) (x : S) . S -> S"
            (Right $ Pi "S" Star (Pi "" (Var 0 "S") Star))
            (checkType M.empty [] $ Abs "S" Star (Abs "x" (Var 0 "S") (Pi "" (Var 1 "S") (Var 2 "S"))))

useDependent :: Test
useDependent =
    TestCase $
        assertEqual
            "fun (S : *) (P : S -> *) (x : S) . P x"
            (Right $ Pi "S" Star (Pi "" (Pi "" (Var 0 "S") Star) (Pi "" (Var 1 "S") Star)))
            (checkType M.empty [] $ Abs "S" Star (Abs "P" (Pi "" (Var 0 "S") Star) (Abs "x" (Var 1 "S") (App (Var 1 "P") (Var 0 "x")))))

big :: Test
big =
    TestCase $
        assertEqual
            "fun (S : *) (P Q : S -> *) (H : forall (x : S), P x -> Q x) (G : forall (x : S), P x) (x : S) . H x (G x)"
            (Right $ Pi "S" Star (Pi "P" (Pi "" (Var 0 "S") Star) (Pi "Q" (Pi "" (Var 1 "S") Star) (Pi "" (Pi "x" (Var 2 "S") (Pi "" (App (Var 2 "P") (Var 0 "x")) (App (Var 2 "Q") (Var 1 "x")))) (Pi "" (Pi "x" (Var 3 "S") (App (Var 3 "P") (Var 0 "x"))) (Pi "x" (Var 4 "S") (App (Var 3 "Q") (Var 0 "x"))))))))
            (checkType M.empty [] $ Abs "S" Star (Abs "P" (Pi "" (Var 0 "S") Star) (Abs "Q" (Pi "" (Var 1 "S") Star) (Abs "H" (Pi "x" (Var 2 "S") (Pi "" (App (Var 2 "P") (Var 0 "x")) (App (Var 2 "Q") (Var 1 "x")))) (Abs "G" (Pi "x" (Var 3 "S") (App (Var 3 "P") (Var 0 "x"))) (Abs "x" (Var 4 "S") (App (App (Var 2 "H") (Var 0 "x")) (App (Var 1 "G") (Var 0 "x")))))))))

tests :: Test
tests =
    TestList
        [ TestLabel "sort" sort
        , TestLabel "λ→" $ TestList [stlc, freeVar]
        , TestLabel "λ2" polyIdent
        , TestLabel "λω" $ TestList [typeCons, useTypeCons]
        , TestLabel "λP2" $ TestList [dependent, useDependent]
        , TestLabel "λC" big
        ]