Pool.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Candidate-pool scoring and boolean expansion: penalized scoring, diverse
-- top-K selection, AND/OR saturation, and structural/truth-vector dedup. The
-- per-node scoring scans run in parallel chunks.
module DataFrame.DecisionTree.Pool (
    evalWithPenaltyVec,
    primaryColExpr,
    primaryColCV,
    takeDiverse,
    candidateParChunk,
    bestDiscreteCandidate,
    boolExprsVec,
    DedupMode (..),
    saturateCandidates,
    roundProducts,
    admitKeys,
    admitVecs,
    dedupCVByExpr,
    nubByExpr,
) where

import DataFrame.DecisionTree.CondVec (CondVec (..), combineAndVec, combineOrVec, countErrorsByVec)
import DataFrame.DecisionTree.Types (CarePoint, SynthConfig (..), TreeConfig (..))
import DataFrame.Internal.Expression (Expr, compareExpr, eSize, eqExpr, getColumns, normalize)

import Control.Parallel.Strategies (parListChunk, rdeepseq, using)
import Data.Function (on)
import Data.List (minimumBy, sortBy)
import qualified Data.Map.Strict as M
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Vector.Unboxed as VU

-- | Penalized score of a candidate: care-point errors plus a complexity
-- penalty, tie-broken by expression size.
evalWithPenaltyVec :: TreeConfig -> [CarePoint] -> CondVec -> (Int, Int)
evalWithPenaltyVec cfg carePoints cv = (countErrorsByVec (cvVec cv) carePoints + penalty, sz)
  where
    sz = eSize (cvExpr cv)
    penalty = floor (complexityPenalty (synthConfig cfg) * fromIntegral sz)

-- | First referenced column of a condition (a sentinel for literal-only ones),
-- used by 'takeDiverse' to enforce per-column diversity.
primaryColExpr :: Expr Bool -> T.Text
primaryColExpr e = case getColumns e of
    [] -> "<noncol>"
    (c : _) -> c

primaryColCV :: CondVec -> T.Text
primaryColCV = primaryColExpr . cvExpr

-- | Keep the first @k@ of an already-sorted list, admitting at most @quota@ per
-- primary column (@Nothing@ disables the per-column cap).
takeDiverse :: Int -> Maybe Int -> (a -> T.Text) -> [a] -> [a]
takeDiverse k Nothing _ = take k
takeDiverse k (Just quota) primary = go M.empty 0
  where
    go !_ !_ [] = []
    go !seen !n (x : xs)
        | n >= k = []
        | M.findWithDefault 0 col seen >= quota = go seen n xs
        | otherwise = x : go (M.insertWith (+) col 1 seen) (n + 1) xs
      where
        !col = primary x

-- | Chunk size for the parallel per-node candidate scans; tuned by an -N
-- sweep, not correctness-affecting.
candidateParChunk :: Int
candidateParChunk = 64

-- | Decorate candidates with their penalty in parallel chunks, forcing only
-- the @(Int, Int)@ key so the order (hence later sorts/minima) is preserved.
decorate :: (CondVec -> (Int, Int)) -> [CondVec] -> [((Int, Int), CondVec)]
decorate penaltyCV xs = zip (map penaltyCV xs `using` parListChunk candidateParChunk rdeepseq) xs

-- | The diverse top-@expressionPairs@ valid candidates by penalty.
sortedTopK :: TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> [CondVec]
sortedTopK cfg penaltyCV validCondVecs =
    map snd (takeDiverse (expressionPairs cfg) (perColumnQuota (synthConfig cfg)) (primaryColCV . snd) sorted)
  where
    sorted = sortBy (compare `on` fst) (decorate penaltyCV validCondVecs)

-- | Lowest-penalty candidate after boolean saturation of the diverse top-K.
bestDiscreteCandidate :: TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> Maybe CondVec
bestDiscreteCandidate _ _ [] = Nothing
bestDiscreteCandidate cfg penaltyCV validCondVecs =
    case saturateCandidates Structural (boolExpansion (synthConfig cfg)) (sortedTopK cfg penaltyCV validCondVecs) of
        [] -> Nothing
        xs -> Just (snd (minimumBy (compare `on` fst) (decorate penaltyCV xs)))

-- | AND/OR expansion of cached conditions to depth @maxDepth@ (each
-- combination is a single vector op, not an interpret).
boolExprsVec :: [CondVec] -> [CondVec] -> Int -> Int -> [CondVec]
boolExprsVec baseExprs prevExprs depth maxDepth
    | depth == 0 = baseExprs ++ boolExprsVec baseExprs prevExprs (depth + 1) maxDepth
    | depth >= maxDepth = []
    | otherwise = combined ++ boolExprsVec baseExprs combined (depth + 1) maxDepth
  where
    combined = roundProducts prevExprs baseExprs

data DedupMode = Structural | TruthVector
    deriving (Eq, Show)

-- | Saturate the pool with AND/OR combinations, deduplicating structurally
-- (byte-identical, first occurrence kept) or by truth vector (opt-in).
saturateCandidates :: DedupMode -> Int -> [CondVec] -> [CondVec]
saturateCandidates Structural maxDepth base = base' ++ go 1 base' seen0
  where
    (base', seen0) = admitKeys Set.empty base
    go !depth frontier seen
        | depth >= maxDepth || null frontier = []
        | otherwise = let (admitted, seen') = admitKeys seen (roundProducts frontier base) in admitted ++ go (depth + 1) admitted seen'
saturateCandidates TruthVector maxDepth base = M.elems (go 1 frontier0 reps0)
  where
    (reps0, frontier0) = admitVecs M.empty base
    go !depth frontier reps
        | depth >= maxDepth || null frontier = reps
        | otherwise = let (reps', admitted) = admitVecs reps (roundProducts frontier base) in go (depth + 1) admitted reps'

-- | One combination round: @frontier × base@ via AND then OR, skipping
-- self-pairs (mirrors 'boolExprsVec' for byte-identical structural output).
roundProducts :: [CondVec] -> [CondVec] -> [CondVec]
roundProducts frontier base =
    [c | e1 <- frontier, e2 <- base, not (eqExpr (cvExpr e1) (cvExpr e2)), c <- [combineAndVec e1 e2, combineOrVec e1 e2]]

-- | Admit candidates with a not-yet-seen normalized form, preserving order.
admitKeys :: Set.Set String -> [CondVec] -> ([CondVec], Set.Set String)
admitKeys = go []
  where
    go acc seen [] = (reverse acc, seen)
    go acc !seen (c : cs)
        | structuralKey c `Set.member` seen = go acc seen cs
        | otherwise = go (c : acc) (Set.insert (structuralKey c) seen) cs

structuralKey :: CondVec -> String
structuralKey = show . normalize . cvExpr

-- | Admit candidates by distinct truth vector, keeping the smallest-expression
-- representative per vector.
admitVecs :: M.Map (VU.Vector Bool) CondVec -> [CondVec] -> (M.Map (VU.Vector Bool) CondVec, [CondVec])
admitVecs = go []
  where
    go acc reps [] = (reps, reverse acc)
    go acc !reps (c : cs) = case M.lookup (cvVec c) reps of
        Nothing -> go (c : acc) (M.insert (cvVec c) c reps) cs
        Just r -> go acc (M.insert (cvVec c) (smaller r c) reps) cs

smaller :: CondVec -> CondVec -> CondVec
smaller a b = case compare (eSize (cvExpr a)) (eSize (cvExpr b)) of
    LT -> a
    GT -> b
    EQ -> if compareExpr (cvExpr a) (cvExpr b) /= GT then a else b

-- | Deduplicate 'CondVec's by normalized 'cvExpr', keeping the first.
dedupCVByExpr :: [CondVec] -> [CondVec]
dedupCVByExpr = go Set.empty
  where
    go _ [] = []
    go seen (cv : cvs)
        | structuralKey cv `Set.member` seen = go seen cvs
        | otherwise = cv : go (Set.insert (structuralKey cv) seen) cvs

-- | Deduplicate expressions by normalized form, keeping the first.
nubByExpr :: [Expr Bool] -> [Expr Bool]
nubByExpr = go Set.empty
  where
    go _ [] = []
    go seen (e : es)
        | k `Set.member` seen = go seen es
        | otherwise = e : go (Set.insert k seen) es
      where
        k = show (normalize e)