{-# LANGUAGE OverloadedLists #-}

module Clustering.FrequentItemSet.BruteForceSpec where

import Control.Arrow (second)
import Control.Monad (forM_)
import Data.Function (($), (.))
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.List qualified as List
import Data.Ord (Ord)
import Data.Ratio (Rational, (%))
import Data.Semigroup (Semigroup (..))
import GHC.IsList (toList)
import Logic
import Numeric.Probability (assertProbability)
import System.FilePath ((<.>))
import Test.Syd
import Text.Show (Show (..))
import Prelude (Num)

import Clustering.FrequentItemSet.BruteForce
import Utils

-- | From https://research.nii.ac.jp/~uno/code/lcm.html#IntroductionstoFrequentItemsetMining
databaseTakeakiUno :: Ord item => Num item => Transactions item
databaseTakeakiUno =
  [ [1, 2, 5, 6, 7]
  , [2, 3, 4, 5]
  , [1, 2, 7, 8, 9]
  , [1, 7, 9]
  , [2, 7, 9]
  , [2, 7, 9] -- Copy-paste typo on the original example
  , [1, 9] -- Add this to increase the support of [1,9] because the original example is wrong…
  , [2]
  ]

-- | From https://hal.science/hal-03500847
databaseHAL03500847T2 :: Ord item => Num item => Transactions item
databaseHAL03500847T2 =
  [ [1, 3, 7, 6]
  , [1, 2, 7]
  , [2, 8, 9, 10]
  , [5, 11]
  ]

spec :: Spec
spec = do
  describe "allFrequentItemSets" do
    forM_ ([2 .. 3] :: [Int]) \minSupp ->
      golden ("db=TakeakiUno" <.> "minSupp=" <> show minSupp) $
        -- Alas, this is not a zero-cost `coerce`
        unName <$> allFrequentItemSets @Int (unitName databaseTakeakiUno) (assertStrictlyPositive minSupp)
    forM_ ([2 .. 3] :: [Int]) \minSupp ->
      golden ("db=HAL03500847T2" <.> "minSupp=" <> show minSupp) $
        unName <$> allFrequentItemSets @Int (unitName databaseHAL03500847T2) (assertStrictlyPositive minSupp)

  describe "associationRules" do
    forM_ ([2 .. 3] :: [Int]) \minSupp ->
      forM_ ([(75 % 100)] :: [Rational]) \minConf ->
        golden
          ("db=TakeakiUno" <.> "minSupp=" <> show minSupp <.> "minConf=75%")
          [ associationRules fis (unitName databaseTakeakiUno) (unitName (assertProbability minConf))
          | fis <- allFrequentItemSets @Int (unitName databaseTakeakiUno) (assertStrictlyPositive minSupp)
          ]

  describe "allClosedFrequentItemSets" do
    forM_ ([2 .. 3] :: [Int]) \minSupp ->
      forM_ ([2 .. 3] :: [Int]) \minSize ->
        golden ("db=TakeakiUno" <.> "minSupp=" <> show minSupp <.> "minSize=" <> show minSize) $
          second (List.sort . toList) . unName
            <$> allClosedFrequentItemSets @Int (unitName databaseTakeakiUno) (assertStrictlyPositive minSupp) (assertStrictlyPositive minSize)
    forM_ ([1 .. 3] :: [Int]) \minSupp ->
      forM_ ([2 .. 4] :: [Int]) \minSize ->
        golden ("db=HAL03500847T2" <.> "minSupp=" <> show minSupp <.> "minSize=" <> show minSize) $
          second (List.sort . toList) . unName
            <$> allClosedFrequentItemSets @Int (unitName databaseHAL03500847T2) (assertStrictlyPositive minSupp) (assertStrictlyPositive minSize)