{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# HLINT ignore "Use const" #-}
{-# HLINT ignore "Unused LANGUAGE pragma" #-}
{-# HLINT ignore "Avoid lambda" #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

module Test.Datalog.NaiveDatabaseSpec where

import Data.Map qualified as Map
import Data.Set qualified as Set
import Datalog.DatalogParser
import Datalog.NaiveDatabase
import Datalog.NaiveDatabase qualified as NaiveDatabase
import Test.Hspec
import Datalog.DatalogDB

spec :: Spec
spec = do
  describe "NaiveDatabase operations" do
    it "can ingest facts into relations & a universe" $ do
      let db =
            NaiveDatabase.withFacts
              [ "parent(\"alice\", \"bob\")."
              , "parent(\"bob\", \"carol\")."
              ]
      constants db
        `shouldBe` Set.fromList (Sym <$> ["alice", "bob", "carol"])
      relations db
        `shouldBe` Map.fromList
          [
            ( "parent"
            , Relation "parent" 2 (Set.fromList $ Sym <<$>> [["alice", "bob"], ["bob", "carol"]]) []
            )
          ]
    it "can ingest facts and rules" do
      let db =
            NaiveDatabase.withFactsAndRules
              [ "parent(\"alice\", \"bob\")."
              , "parent(\"bob\", \"carol\")."
              ]
              [ "ancestor(X,Y) :- parent(X,Z), ancestor(Z,Y)."
              , "ancestor(X,Y) :- parent(X,Y)."
              ]
          parentRelation =
            Relation
              { _name = "parent"
              , _arity = 2
              , _tuples =
                  Set.fromList $
                    map (Sym <$>) [["alice", "bob"], ["bob", "carol"]]
              , _rules = []
              }
          ancestorRule = RelationRule
              { headVariables = ["X", "Y", "Z"]
              , bodyElements =
                  [ RuleBodyElement
                      { _subRelationId = "parent"
                      , _ruleElements =
                          [ RuleElementVariable 0
                          , RuleElementVariable 2
                          ]
                      }
                  , RuleBodyElement
                      { _subRelationId = "ancestor"
                      , _ruleElements =
                          [ RuleElementVariable 2
                          , RuleElementVariable 1
                          ]
                      }
                  ]
              }
          ancestorRule2 = RelationRule
              { headVariables = ["X", "Y"]
              , bodyElements =
                  [ RuleBodyElement
                      { _subRelationId = "parent"
                      , _ruleElements = [RuleElementVariable 0, RuleElementVariable 1]
                      }
                  ]
              }
          ancestorRelation =
            Relation
              { _arity = 2
              , _name = "ancestor"
              , _tuples = Set.empty
              , _rules = [ancestorRule, ancestorRule2]
              }

      constants db
        `shouldBe` Set.fromList (Sym <$> ["alice", "bob", "carol"])

      relations db
        `shouldBe` Map.fromList
          [ ("ancestor", ancestorRelation)
          , ("parent", parentRelation)
          ]

    it "can ingest facts and rules with constants" do
      let db =
            NaiveDatabase.withFactsAndRules
              []
              ["ancestor(X,\"patriarch\") :- ."]
          ancestorRule =
            RelationRule
              { headVariables = ["X"]
              , bodyElements = []
              }
          ancestorRelation =
            Relation
              { _arity = 2
              , _name = "ancestor"
              , _tuples = Set.empty
              , _rules = [ancestorRule]
              }
      relations db
        `shouldBe` Map.singleton "ancestor" ancestorRelation

      constants db
        `shouldBe` Set.fromList (Sym <$> ["patriarch"])

    it "can ingest facts and rules with duplicate head entries" do
      let db =
            NaiveDatabase.withFactsAndRules
              []
              ["equivalent(Q,Q) :- ."]
          equivalentRule =
            RelationRule
              { headVariables = ["Q"]
              , bodyElements = []
              }
          equivalentRelation =
            Relation
              { _arity = 2
              , _name = "equivalent"
              , _tuples = Set.empty
              , _rules = [equivalentRule]
              }
      relations db
        `shouldBe` Map.singleton "equivalent" equivalentRelation

      constants db
        `shouldBe` Set.empty

    it "can ingest facts and rules with duplicate head entries" do
      let db =
            NaiveDatabase.withFactsAndRules
              []
              ["equivalent(Q,Q) :- ."]
          rule1 =
            RelationRule
              { headVariables = ["Q"]
              , bodyElements = []
              }

          equivalentRelation =
            Relation
              { _arity = 2
              , _name = "equivalent"
              , _tuples = Set.empty
              , _rules = [rule1]
              }
      relations db
        `shouldBe` Map.singleton "equivalent" equivalentRelation

      constants db
        `shouldBe` Set.empty

    it "can ingest a theory of equivalence relations" do
      let db =
            NaiveDatabase.withFactsAndRules
              []
              [ "equivalent(Q,Q) :- ."
              , "equivalent(R,Q) :- equivalent(Q,R)."
              , "equivalent(Q,S) :- equivalent(Q,R), equivalent(R,S)."
              ]
          rule1 =
            RelationRule
              { headVariables = ["Q"]
              , bodyElements = []
              }
          rule2 =
            RelationRule
              { headVariables = ["R", "Q"]
              , bodyElements =
                  [ RuleBodyElement
                      { _subRelationId = "equivalent"
                      , _ruleElements =
                          [ RuleElementVariable 1
                          , RuleElementVariable 0
                          ]
                      }
                  ]
              }
          rule3 =
            RelationRule
              { headVariables = ["Q", "S", "R"]
              , bodyElements =
                  [ RuleBodyElement
                      { _subRelationId = "equivalent"
                      , _ruleElements =
                          [ RuleElementVariable 0
                          , RuleElementVariable 2
                          ]
                      }
                  , RuleBodyElement
                      { _subRelationId = "equivalent"
                      , _ruleElements =
                          [ RuleElementVariable 2
                          , RuleElementVariable 1
                          ]
                      }
                  ]
              }
          equivalentRelation =
            Relation
              { _arity = 2
              , _name = "equivalent"
              , _tuples = Set.empty
              , _rules = [rule1, rule2, rule3]
              }
      relations db
        `shouldBe` Map.singleton "equivalent" equivalentRelation

      constants db
        `shouldBe` Set.empty

    it "can do basic queries" do
      let db :: NaiveDatabase =
            NaiveDatabase.withFacts
              [ "parent(\"alice\", \"bob\")."
              , "parent(\"bob\", \"carol\")."
              ]
      query db "?- parent(alice,X)." `shouldBe` "#NYI" -- ideally, 'bob'

(<<$>>) :: (Functor f1, Functor f2) => (a -> b) -> f1 (f2 a) -> f1 (f2 b)
(<<$>>) = fmap fmap fmap

-- (<<<$>>>) :: (Functor f1, Functor f2, Functor f3) => f1 (f2 (a -> b)) -> f1 (f2 (f3 a -> f3 b))
-- (<<<$>>>) :: Functor f => (a1 -> b) -> (a2 -> a1) -> f a2 -> f b
-- (<<<$>>>) :: (Functor f1, Functor f2) => (a1 -> a2 -> b) -> f1 a1 -> f1 (f2 a2 -> f2 b)
-- (<<<$>>>) :: (Functor f1, Functor f2, Functor f3) => f1 (a -> b) -> f1 (f2 (f3 a) -> f2 (f3 b))
(<<<$>>>) :: (Functor f1, Functor f2, Functor f3) => (a -> b) -> f1 (f2 (f3 a)) -> f1 (f2 (f3 b))
(<<<$>>>) = fmap fmap fmap fmap fmap fmap fmap fmap