{-# LANGUAGE NoImplicitPrelude #-}

module Felix.Test.Unit.Core (unitTests) where

import Base hiding (Empty)
import Felix.Checking.Core
import Felix.Checking.Foundation qualified as Foundation
import Felix.Checking.SetConstruction

import Control.DeepSeq (NFData(..), force)
import Hedgehog
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Data.Set qualified as Set
import Test.Tasty
import Test.Tasty.HUnit hiding (assert)
import Test.Tasty.Hedgehog (testPropertyNamed)


data TestGlobal
    = TestGlobal
    | TestPairGlobal
    deriving (Show, Eq, Ord)

instance NFData TestGlobal where
    rnf _global =
        ()

testGlobalType :: TestGlobal -> Maybe CoreType
testGlobalType TestGlobal =
    Just TySet
testGlobalType TestPairGlobal =
    Just (TySet `TyArrow` (TySet `TyArrow` TySet))

unitTests :: TestTree
unitTests =
    testGroup "Checked core"
        [ testCase
            "freezes explicit binder positions"
            freezesExplicitBinderPositions
        , testCase
            "freezes generalized substitution without capture"
            freezesGeneralizedSubstitution
        , testCase
            "rejects ill-typed and open terms"
            rejectsInvalidTerms
        , testCase
            "rechecks canonical terms without unchecked construction"
            rechecksCanonicalTerms
        , testCase
            "checks scoped canonical weakening and substitution"
            checksScopedCanonicalOperations
        , testCase
            "builds set-induction hypotheses from the complete property"
            buildsSetInductionHypotheses
        , testCase
            "specializes the checked separation characteristic"
            specializesCheckedSeparationCharacteristic
        , testCase
            "specializes the checked replacement characteristic"
            specializesCheckedReplacementCharacteristic
        , testCase
            "derives named construction views from checked components"
            derivesNamedConstructionViews
        , testCase
            "thaws checked closed terms without changing them"
            thawsCheckedClosedTerms
        , testPropertyNamed
            "optimized freeze agrees with bounded reference"
            "prop_freezeAgreesWithReference"
            prop_freezeAgreesWithReference
        ]

freezesExplicitBinderPositions :: Assertion
freezesExplicitBinderPositions = do
    let x = 0 :: Int
        y = 1 :: Int
        freeze body = do
            checked <-
                either
                    (assertFailure . show)
                    pure
                    (checkClosedCore
                        testGlobalType
                        (coreLambda TySet x
                            (coreLambda TyProp y body)))
            either
                (assertFailure . show)
                (pure . frozenCoreTerm)
                (freezeClosed checked)
    nearest <-
        freeze (coreLocal y)
    outer <-
        freeze (coreLocal x)
    global <-
        freeze (coreGlobal TestGlobal)
    assertEqual
        "nearest binder"
        (CLam TySet (CLam TyProp (CBound 0)))
        nearest
    assertEqual
        "outer binder"
        (CLam TySet (CLam TyProp (CBound 1)))
        outer
    assertEqual
        "global with both vacuous binders"
        (CLam TySet (CLam TyProp (CGlobal TestGlobal)))
        global

freezesGeneralizedSubstitution :: Assertion
freezesGeneralizedSubstitution = do
    let outer = 0 :: Int
        inner = 1 :: Int
        placeholder = 2 :: Int
        innerTerm =
            coreLambda TySet inner (coreLocal placeholder)
        substituted =
            innerTerm >>= \local ->
                if local == placeholder
                    then coreLocal outer
                    else coreLocal local
        term =
            coreLambda TyProp outer substituted
    checked <-
        either
            (assertFailure . show)
            pure
            (checkClosedCore testGlobalType term)
    optimized <-
        either
            (assertFailure . show)
            pure
            (freezeClosed checked)
    reference <-
        either
            (assertFailure . show)
            pure
            (referenceFreezeClosed checked)
    assertEqual "reference result" reference optimized
    assertEqual
        "outer variable remains outside the inner binder"
        (CLam TyProp (CLam TySet (CBound 1)))
        (frozenCoreTerm optimized)

rejectsInvalidTerms :: Assertion
rejectsInvalidTerms = do
    assertEqual
        "free local"
        (Left UnboundCoreLocal)
        (checkedCoreType
            <$> checkClosedCore
                testGlobalType
                (coreLocal (0 :: Int)))
    assertEqual
        "application argument"
        (Left
            (ApplicationArgumentTypeMismatch
                TySet
                TyProp))
        (checkedCoreType
            <$> checkClosedCore
                testGlobalType
                (coreApply
                    (coreIntrinsic FamilyUnion)
                    coreFalsum))
    assertEqual
        "equality operand"
        (Left
            (EqualityOperandTypeMismatch
                TySet
                TyProp))
        (checkedCoreType
            <$> checkClosedCore
                testGlobalType
                (coreEquality
                    TySet
                    coreFalsum
                    (coreGlobal TestGlobal)))

rechecksCanonicalTerms :: Assertion
rechecksCanonicalTerms = do
    let term =
            CForall TySet
                (CEq TySet
                    (CBound 0)
                    (CBound 0))
    assertEqual
        "well-typed canonical proposition"
        (Right (TyProp, term))
        ( (\checked ->
                ( frozenCoreType checked
                , frozenCoreTerm checked
                ))
            <$> checkCanonicalCore testGlobalType term
        )
    assertEqual
        "out-of-scope de Bruijn index"
        (Left (UnboundCoreIndex 1))
        (frozenCoreType
            <$> checkCanonicalCore
                testGlobalType
                (CForall TySet (CBound 1)))
    assertEqual
        "canonical application still checks argument types"
        (Left
            (ApplicationArgumentTypeMismatch
                TySet
                TyProp))
        (frozenCoreType
            <$> checkCanonicalCore
                testGlobalType
                (CApp
                    (CIntrinsic FamilyUnion)
                    CFalsum))

checksScopedCanonicalOperations :: Assertion
checksScopedCanonicalOperations = do
    scoped <-
        either
            (assertFailure . show)
            pure
            (checkScopedCanonicalCore
                testGlobalType
                [TySet]
                (CBound 0))
    weakened <-
        either
            (assertFailure . show)
            pure
            (weakenScopedCore
                testGlobalType
                TyProp
                scoped)
    assertEqual
        "nearest binder insertion shifts the prior local"
        ( [TyProp, TySet]
        , TySet
        , CBound 1
        )
        ( scopedCoreContext weakened
        , scopedCoreType weakened
        , scopedCoreTerm weakened
        )
    assertEqual
        "top-level binder substitution removes its index"
        (CEq TySet
            (CIntrinsic Empty)
            (CIntrinsic Empty))
        (instantiateCanonical
            (CIntrinsic Empty
                :: CanonicalTerm TestGlobal)
            (CEq TySet
                (CBound 0)
                (CBound 0)))
    root <-
        either
            (assertFailure . show)
            pure
            (checkScopedCanonicalCore
                testGlobalType
                []
                (CGlobal TestGlobal))
    assertEqual
        "empty scoped context closes"
        (Just
            (TySet, CGlobal TestGlobal))
        ( (\closed ->
                ( frozenCoreType closed
                , frozenCoreTerm closed
                ))
            <$> closeScopedCore root
        )

buildsSetInductionHypotheses :: Assertion
buildsSetInductionHypotheses = do
    let propertyTerm =
            CImp
                (CEq TySet
                    (CBound 0)
                    (CIntrinsic Empty))
                (CEq TySet
                    (CBound 0)
                    (CBound 0))
    claimProperty <-
        either
            (assertFailure . show)
            pure
            (checkScopedCanonicalCore
                testGlobalType
                [TySet]
                propertyTerm)
    (predicate, hypothesis, step, result) <-
        maybe
            (assertFailure "set-induction instance was not constructed")
            pure
            (scopedSetInductionInstance 0 claimProperty)
    let abstractedProperty =
            CImp
                (CEq TySet
                    (CBound 0)
                    (CIntrinsic Empty))
                (CEq TySet
                    (CBound 0)
                    (CBound 0))
        memberHypothesis =
            CForall TySet
                (CImp
                    (CApp
                        (CApp
                            (CIntrinsic Member)
                            (CBound 0))
                        (CBound 1))
                    abstractedProperty)
    assertEqual
        "set induction abstracts the selected property once"
        (CLam TySet abstractedProperty)
        (scopedCoreTerm predicate)
    assertEqual
        "antecedent and goal are both generalized over the member"
        memberHypothesis
        (scopedCoreTerm hypothesis)
    assertEqual
        "set-induction step owns its member-wise hypothesis"
        (CForall TySet
            (CImp memberHypothesis abstractedProperty))
        (scopedCoreTerm step)
    assertEqual
        "set-induction result closes the complete property"
        (CForall TySet abstractedProperty)
        (scopedCoreTerm result)

specializesCheckedSeparationCharacteristic :: Assertion
specializesCheckedSeparationCharacteristic = do
    foundation <-
        either
            (assertFailure . show)
            pure
            Foundation.checkedFoundation
    let bound = CIntrinsic Empty
        predicate =
            CLam TySet
                (CEq TySet (CBound 0) (CBound 0))
        separation =
            CApp
                (CApp (CIntrinsic Sep) bound)
                predicate
    body <-
        either
            (assertFailure . show)
            pure
            (checkScopedCanonicalCore
                testGlobalType
                []
                separation)
    definition <-
        maybe
            (assertFailure "separation did not form a set definition")
            pure
            (scopedSetDefinition
                (Foundation.foundationAxiomFrozen
                    foundation
                    Foundation.SeparationCharacteristic)
                body)
    let generated =
            instantiateCanonical
                separation
                (scopedCoreTerm definition)
        expected =
            betaNormalize
                (specializeForall predicate
                (specializeForall bound
                    (mapCanonicalGlobals absurd
                        (frozenCoreTerm
                            (Foundation.foundationAxiomFrozen
                                foundation
                                Foundation.SeparationCharacteristic)))))
    assertEqual
        "local characteristic is the checked rule specialization"
        expected
        generated

specializesCheckedReplacementCharacteristic :: Assertion
specializesCheckedReplacementCharacteristic = do
    foundation <-
        either
            (assertFailure . show)
            pure
            Foundation.checkedFoundation
    pair <- checked [] (CGlobal TestPairGlobal)
    domain <- checked [] (CIntrinsic Empty)
    value <- checked [TySet] (CBound 0)
    (graph, checkedDomain, function) <-
        maybe
            (assertFailure "checked values did not form a replacement graph")
            pure
            (scopedReplacementGraph pair domain value)
    definition <-
        maybe
            (assertFailure "replacement graph did not form a definition")
            pure
            (scopedCharacteristicDefinition
                (Foundation.foundationAxiomFrozen
                    foundation
                    Foundation.ReplacementCharacteristic)
                graph
                (checkedDomain :| [function]))
    let generated =
            instantiateCanonical
                (scopedCoreTerm graph)
                (scopedCoreTerm definition)
        expected =
            betaNormalize
                (specializeForall (scopedCoreTerm function)
                (specializeForall (scopedCoreTerm checkedDomain)
                    (mapCanonicalGlobals absurd
                        (frozenCoreTerm
                            (Foundation.foundationAxiomFrozen
                                foundation
                                Foundation.ReplacementCharacteristic)))))
    assertEqual
        "local graph characteristic is the checked rule specialization"
        expected
        generated
  where
    checked context term =
        either
            (assertFailure . show)
            pure
            (checkScopedCanonicalCore testGlobalType context term)

derivesNamedConstructionViews :: Assertion
derivesNamedConstructionViews = do
    foundation <-
        either (assertFailure . show) pure Foundation.checkedFoundation
    bound <- checked [] (CIntrinsic Empty)
    predicate <- checked [TySet] (CEq TySet (CBound 0) (CBound 0))
    separation <-
        maybe
            (assertFailure "checked separation descriptor failed")
            pure
            (checkedSeparationConstruction testGlobalType bound predicate)
    (separationView, separationEquation) <-
        maybe
            (assertFailure "checked separation views failed")
            pure
            (namedSetConstructionLocalViews
                (checkedFoundationSetConstruction foundation)
                separation)
    expectedSeparationView <-
        maybe
            (assertFailure "checked separation characteristic failed")
            pure
            (scopedSetDefinition
                (Foundation.foundationAxiomFrozen foundation
                    Foundation.SeparationCharacteristic)
                (namedSetConstructionTerm separation))
    assertEqual "separation view is the checked specialization"
        expectedSeparationView separationView
    assertEqual "separation view is first-order"
        (Set.singleton Foundation.EmptyCharacteristic)
        (Foundation.foundationAxiomDependencies
            (scopedCoreTerm separationView))
    assertEqual "separation equation retains exact construction"
        (Set.fromList
            [ Foundation.EmptyCharacteristic
            , Foundation.SeparationCharacteristic
            ])
        (Foundation.foundationAxiomDependencies
            (scopedCoreTerm separationEquation))

    firstDomain <- checked [] (CIntrinsic Empty)
    singleValue <- checked [TySet] (CBound 0)
    singleCondition <- checked [TySet]
        (CEq TySet (CBound 0) (CBound 0))
    singleReplacement <-
        maybe
            (assertFailure "checked one-domain replacement failed")
            pure
            (checkedFunctionalReplacementConstruction
                testGlobalType
                (firstDomain :| [])
                singleValue
                (Just singleCondition))
    assertEqual "one-domain replacement has one canonical term"
        (CApp
            (CApp
                (CIntrinsic Repl)
                (CApp
                    (CApp (CIntrinsic Sep) (CIntrinsic Empty))
                    (CLam TySet
                        (CEq TySet (CBound 0) (CBound 0)))))
            (CLam TySet (CBound 0)))
        (scopedCoreTerm
            (namedSetConstructionTerm singleReplacement))

    secondDomain <- checked [TySet] (CBound 0)
    value <- checked [TySet, TySet] (CBound 0)
    condition <- checked [TySet, TySet]
        (CEq TySet (CBound 0) (CBound 0))
    replacement <-
        maybe
            (assertFailure "checked replacement descriptor failed")
            pure
            (checkedFunctionalReplacementConstruction
                testGlobalType
                (firstDomain :| [secondDomain])
                value
                (Just condition))
    (replacementView, replacementEquation) <-
        maybe
            (assertFailure "checked replacement views failed")
            pure
            (namedSetConstructionLocalViews
                (checkedFoundationSetConstruction foundation)
                replacement)
    let terminal =
            andP
                (CEq TySet (CBound 0) (CBound 0))
                (CEq TySet (CBound 2) (CBound 0))
        secondWitness =
            existsP
                (andP
                    (memberP (CBound 0) (CBound 1))
                    terminal)
        firstWitness =
            existsP
                (andP
                    (memberP (CBound 0) (CIntrinsic Empty))
                    secondWitness)
        expectedReplacementTerm =
            CForall TySet
                (CEq TyProp
                    (memberP (CBound 0) (CBound 1))
                    firstWitness)
    expectedReplacementView <- checked [TySet] expectedReplacementTerm
    assertEqual "replacement view preserves bounds and condition"
        expectedReplacementView replacementView
    assertEqual "flattened replacement view is first-order"
        (Set.singleton Foundation.EmptyCharacteristic)
        (Foundation.foundationAxiomDependencies
            (scopedCoreTerm replacementView))
    assertEqual "replacement equation retains every helper"
        (Set.fromList
            [ Foundation.FamilyUnionCharacteristic
            , Foundation.EmptyCharacteristic
            , Foundation.SeparationCharacteristic
            , Foundation.ReplacementCharacteristic
            ])
        (Foundation.foundationAxiomDependencies
            (scopedCoreTerm replacementEquation))
  where
    checked context term =
        either
            (assertFailure . show)
            pure
            (checkScopedCanonicalCore testGlobalType context term)

    memberP element set =
        CApp (CApp (CIntrinsic Member) element) set

    notP proposition = CImp proposition CFalsum

    andP left right =
        notP (CImp left (notP right))

    existsP proposition =
        notP (CForall TySet (notP proposition))

specializeForall
    :: CanonicalTerm global
    -> CanonicalTerm global
    -> CanonicalTerm global
specializeForall argument = \case
    CForall _binderType body ->
        instantiateCanonical argument body
    _ ->
        error "the checked characteristic lost a binder"

betaNormalize :: CanonicalTerm global -> CanonicalTerm global
betaNormalize = \case
    CApp function argument ->
        case betaNormalize function of
            CLam _binderType body ->
                betaNormalize
                    (instantiateCanonical
                        (betaNormalize argument)
                        body)
            normalizedFunction ->
                CApp normalizedFunction (betaNormalize argument)
    CLam binderType body ->
        CLam binderType (betaNormalize body)
    CImp premise conclusion ->
        CImp
            (betaNormalize premise)
            (betaNormalize conclusion)
    CEq operandType left right ->
        CEq operandType
            (betaNormalize left)
            (betaNormalize right)
    CForall binderType body ->
        CForall binderType (betaNormalize body)
    term -> term

thawsCheckedClosedTerms :: Assertion
thawsCheckedClosedTerms = do
    let source =
            coreLambda TySet (0 :: Int)
                (coreForall TySet 1
                    (coreEquality
                        TySet
                        (coreLocal 0)
                        (coreLocal 1)))
        freeze syntax = do
            checked <-
                either
                    (assertFailure . show)
                    pure
                    (checkClosedCore testGlobalType syntax)
            either
                (assertFailure . show)
                pure
                (freezeClosed checked)
    original <-
        freeze source
    roundTrip <-
        freeze (thawFrozenCore original)
    assertEqual "frozen term" original roundTrip
    assertEqual
        "global inventory"
        mempty
        (frozenCoreGlobals original)

prop_freezeAgreesWithReference :: Property
prop_freezeAgreesWithReference = property do
    depth <-
        forAll (Gen.int (Range.linear 1 10))
    binderTypes <-
        forAll
            (Gen.list
                (Range.singleton depth)
                (Gen.element
                    [ TyProp
                    , TySet
                    , TySet `TyArrow` TySet
                    ]))
    selected <-
        forAll
            (Gen.maybe
                (Gen.int (Range.linear 0 (depth - 1))))
    let binders =
            zip [0 :: Int ..] binderTypes
        body =
            maybe
                (coreGlobal TestGlobal)
                coreLocal
                selected
        term =
            foldr
                (\(local, binderType) ->
                    coreLambda binderType local)
                body
                binders
    checked <-
        evalEither
            (checkClosedCore testGlobalType term)
    optimized <-
        evalEither (force <$> freezeClosed checked)
    reference <-
        evalEither (force <$> referenceFreezeClosed checked)
    optimized === reference