Segmented.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

{- | Fit a separate base model per categorical value-combination, routing each
row to its segment at predict time; unseen or too-small segments fall back to a
global fit. Optional partial pooling (linear base) shrinks small segments.
-}
module DataFrame.Segmented (
    module DataFrame.Model,
    Segmented (..),
    segmented,
    segmentOn,
    pooled,
    Segment (..),
    SegmentedModel (..),
    SegmentFit (..),
) where

import Data.List (foldl', (\\))
import qualified Data.Map.Strict as M
import Data.Maybe (isJust)
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import DataFrame.Featurize.Internal (featureNames, numericMatrix, targetDoubles)
import DataFrame.Internal.Column (
    Column (..),
    Columnable,
    bitmapTestBit,
    columnBitmap,
    hasElemType,
 )
import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn)
import DataFrame.Internal.Expression (Expr (..))
import DataFrame.Internal.Types (SBool (..), sIntegral)
import DataFrame.LinearAlgebra (Matrix, gram, matVec, tMatVec)
import DataFrame.LinearAlgebra.Solve (choleskySolve, qrLeastSquares)
import DataFrame.LinearModel.Logistic (LogisticConfig)
import DataFrame.LinearModel.Regression (
    LinearConfig (..),
    LinearRegressor (..),
 )
import DataFrame.Model
import DataFrame.Operations.Core (nRows)
import DataFrame.Operations.Subset (columnToTextVec, exclude, rowsAtIndices)
import DataFrame.Operators ((.&&.), (.==.))
import DataFrame.SymbolicRegression (SRConfig)

{- | A base estimator @cfg@ wrapped to fit one model per categorical
value-combination. @segOn@ picks the columns ('Nothing' = auto-detect),
@segMinRows@ the smallest own-model segment, @segPool@ the pooling strength @λ@.
-}
data Segmented cfg = Segmented
    { segBase :: !cfg
    , segOn :: !(Maybe [T.Text])
    , segMaxCard :: !Int
    , segMinRows :: !Int
    , segPool :: !Double
    }
    deriving (Eq, Show)

-- | Wrap a base config with the defaults: auto-detect, cap 32, min 30 rows, no pooling.
segmented :: cfg -> Segmented cfg
segmented base = Segmented base Nothing 32 30 0

-- | Segment only on the named columns (each must be Text), overriding auto-detect.
segmentOn :: Segmented cfg -> [T.Text] -> Segmented cfg
segmentOn s cols = s{segOn = Just cols}

-- | Set the pooling strength @λ@ (shrink segments toward the reference).
pooled :: Segmented cfg -> Double -> Segmented cfg
pooled s lam = s{segPool = lam}

-- | One fitted segment: its categorical key, row count, and base model.
data Segment model = Segment
    { segKey :: ![T.Text]
    , segN :: !Int
    , segModel :: !model
    }
    deriving (Show)

{- | A fitted segmented model: the columns segmented on, the per-segment models
(ascending key order), the observed combinations that fell back (key + row
count), the global fallback model, and the compiled routing expression.
-}
data SegmentedModel a model = SegmentedModel
    { smCatCols :: ![T.Text]
    , smSegments :: ![Segment model]
    , smFellBack :: ![([T.Text], Int)]
    , smFallback :: !model
    , smExpr :: !(Expr a)
    }
    deriving (Show)

{- | How a base estimator fits its per-segment models under pooling strength @λ@.
The default fits each segment independently and rejects @λ > 0@; the linear
instance overrides it with closed-form shrinkage. Every base model needs an instance.
-}
class (Fit cfg (Expr a)) => SegmentFit cfg a where
    -- | Fit the qualifying segments (each a numeric-only frame, in order).
    fitSegments :: cfg -> Double -> Expr a -> [DataFrame] -> [ModelOf cfg (Expr a)]
    fitSegments cfg lam target dfs
        | lam == 0 = map (fit cfg target) dfs
        | otherwise =
            error
                "Segmented: pooling (lambda > 0) is not supported for this base model; use lambda = 0 or a linear base."

-- | Logistic segments support independent fitting (lambda = 0) only, for now.
instance (Columnable a, Ord a) => SegmentFit LogisticConfig a

-- | Symbolic-regression segments support independent fitting (lambda = 0) only.
instance SegmentFit SRConfig Double

instance
    ( Fit cfg (Expr a)
    , SegmentFit cfg a
    , Predict (ModelOf cfg (Expr a))
    , Prediction (ModelOf cfg (Expr a)) ~ Expr a
    , Columnable a
    ) =>
    Fit (Segmented cfg) (Expr a)
    where
    type ModelOf (Segmented cfg) (Expr a) = SegmentedModel a (ModelOf cfg (Expr a))
    type FrameReq (Segmented cfg) (Expr a) = 'AnyFrame
    fit = fitSegmented

instance Predict (SegmentedModel a model) where
    type Prediction (SegmentedModel a model) = Expr a
    predict = smExpr

fitSegmented ::
    forall cfg a.
    ( Fit cfg (Expr a)
    , SegmentFit cfg a
    , Predict (ModelOf cfg (Expr a))
    , Prediction (ModelOf cfg (Expr a)) ~ Expr a
    , Columnable a
    ) =>
    Segmented cfg ->
    Expr a ->
    DataFrame ->
    SegmentedModel a (ModelOf cfg (Expr a))
fitSegmented (Segmented base mcols maxCard minRows lam) target df =
    seq (guardNumeric df textCols target) result
  where
    mTarget = case target of
        Col n -> Just n
        _ -> Nothing
    feats = featureNames target df
    textCols = [c | c <- feats, isTextCol df c]
    catCols = resolveCatCols df mTarget textCols mcols maxCard
    numericFrame = exclude textCols
    globalM = fit base target (numericFrame df)
    result
        | null catCols =
            SegmentedModel [] [] [] globalM (predict globalM)
        | otherwise =
            let d = length (feats \\ textCols)
                floor' = max minRows (d + 1)
                grouped = groupByKey df catCols
                (qualifying, undersized) =
                    span' (\(_, ixs) -> VU.length ixs >= floor') grouped
                qualFrames =
                    [numericFrame (rowsAtIndices ixs df) | (_, ixs) <- qualifying]
                segModels = fitSegments base lam target qualFrames
                segments =
                    zipWith
                        (\(k, ixs) m -> Segment k (VU.length ixs) m)
                        qualifying
                        segModels
                fellBack = [(k, VU.length ixs) | (k, ixs) <- undersized]
                expr = buildExpr catCols segments globalM
             in SegmentedModel catCols segments fellBack globalM expr

-- | 'span' over a predicate that need not hold contiguously (a filter partition).
span' :: (b -> Bool) -> [b] -> ([b], [b])
span' p xs = (filter p xs, filter (not . p) xs)

{- | Compile the routing: a right-folded @If@ ladder of @key == value@ conjuncts,
ending in the fallback model's prediction. Keys are disjoint, so order is
immaterial.
-}
buildExpr ::
    (Columnable a, Predict model, Prediction model ~ Expr a) =>
    [T.Text] ->
    [Segment model] ->
    model ->
    Expr a
buildExpr catCols segs fallback =
    foldr
        (\(Segment key _ m) acc -> If (keyCond catCols key) (predict m) acc)
        (predict fallback)
        segs

-- | @col1 == v1 && col2 == v2 && ...@ for a segment's key.
keyCond :: [T.Text] -> [T.Text] -> Expr Bool
keyCond catCols vals =
    foldr1 (.&&.) [(Col c :: Expr T.Text) .==. Lit v | (c, v) <- zip catCols vals]

{- | The Text feature columns to segment on, chosen from the frame's Text features
@textCols@. An explicit list is validated to be all-Text; auto-detect keeps Text
columns with at most @maxCard@ distinct values.
-}
resolveCatCols ::
    DataFrame -> Maybe T.Text -> [T.Text] -> Maybe [T.Text] -> Int -> [T.Text]
resolveCatCols df mTarget textCols mcols maxCard = case mcols of
    Just cols ->
        let bad = filter (\c -> not (isTextCol df c) || Just c == mTarget) cols
         in if null bad
                then cols
                else
                    error
                        ( "Segmented: segmentOn columns must be Text features (not the target); invalid: "
                            ++ show bad
                        )
    Nothing -> [c | c <- textCols, distinctCount df c <= maxCard]

isTextCol :: DataFrame -> T.Text -> Bool
isTextCol df c = hasElemType @T.Text (unsafeGetColumn c df)

distinctCount :: DataFrame -> T.Text -> Int
distinctCount df c =
    Set.size (Set.fromList (V.toList (columnToTextVec (unsafeGetColumn c df))))

{- | Reject feature columns that are neither Text (dropped\/segmented) nor
non-null 'Double', naming each with its fix — clearer than the base fitter's
raw type mismatch.
-}
guardNumeric :: DataFrame -> [T.Text] -> Expr a -> ()
guardNumeric df textCols target =
    case problems of
        [] -> ()
        ps ->
            error
                ( "Segmented: unusable feature column(s):\n"
                    ++ unlines (map fmt ps)
                )
  where
    problems =
        [ (c, r)
        | c <- featureNames target df \\ textCols
        , Just r <- [reason (unsafeGetColumn c df)]
        ]
    fmt (c, r) = "  " ++ T.unpack c ++ ": " ++ r
    reason col
        | isJust (columnBitmap col) =
            Just
                "has missing values — drop them (filterJust / filterAllJust) or model missingness explicitly; imputing risks train/inference skew"
        | isIntegralCol col =
            Just "is an integer column — cast to Double with F.toDouble"
        | not (hasElemType @Double col) =
            Just
                "is not Double — convert to Double (numeric) or segment on it (categorical)"
        | otherwise = Nothing
    isIntegralCol col = case col of
        UnboxedColumn _ (_ :: VU.Vector b) -> case sIntegral @b of
            STrue -> True
            _ -> False
        _ -> False

{- | Group row indices by their composite categorical key, dropping rows whose key
has a null in any segmented column (served by the fallback). Ascending key order.
-}
groupByKey :: DataFrame -> [T.Text] -> [([T.Text], VU.Vector Int)]
groupByKey df catCols =
    map (\(k, is) -> (k, VU.fromList (reverse is))) (M.toAscList grouped)
  where
    n = nRows df
    cols = map (`unsafeGetColumn` df) catCols
    textVecs = map columnToTextVec cols
    bitmaps = map columnBitmap cols
    validRow i = all (maybe True (`bitmapTestBit` i)) bitmaps
    keyOf i = [tv V.! i | tv <- textVecs]
    grouped =
        foldl'
            (\m i -> if validRow i then M.insertWith (++) (keyOf i) [i] m else m)
            M.empty
            [0 .. n - 1]

{- | Linear segments with exact closed-form pooling. @λ = 0@ is independent OLS;
@λ > 0@ shrinks each segment's coefficients toward the @n_g@-weighted mean of the
per-segment fits.
-}
instance SegmentFit LinearConfig Double where
    fitSegments cfg lam target dfs
        | lam == 0 = map (fit cfg target) dfs
        | null dfs = []
        | otherwise = shrinkLinear cfg lam target dfs

shrinkLinear ::
    LinearConfig -> Double -> Expr Double -> [DataFrame] -> [LinearRegressor]
shrinkLinear cfg lam target dfs = map toReg dsegs
  where
    names = case dfs of
        (d0 : _) -> featureNames target d0
        [] -> error "shrinkLinear: no segments"
    d = length names
    mats = [snd (numericMatrix names dframe) | dframe <- dfs]
    ys = [targetDoubles target dframe | dframe <- dfs]
    ns = map V.length mats
    pooledRows = V.concat mats
    nP = V.length pooledRows
    means =
        VU.generate d $ \j ->
            sum [(pooledRows V.! i) VU.! j | i <- [0 .. nP - 1]] / fromIntegral nP
    sds =
        VU.generate d $ \j ->
            let m = means VU.! j
                var =
                    sum [sq ((pooledRows V.! i) VU.! j - m) | i <- [0 .. nP - 1]]
                        / fromIntegral nP
             in sqrt var
    stdValue j x = let s = sds VU.! j in if s == 0 then 0 else (x - means VU.! j) / s
    augStd row = VU.cons 1 (VU.imap stdValue row)
    zMats = [V.map augStd m | m <- mats]
    olsAll = zipWith fitOLS zMats ys
    fitOLS z y = either (const Nothing) Just (qrLeastSquares z y)
    good = [(n, sol) | (n, Just sol) <- zip ns olsAll]
    totW = fromIntegral (sum (map fst good)) :: Double
    dref
        | null good = VU.replicate (d + 1) 0
        | otherwise =
            VU.generate (d + 1) $ \k ->
                sum [fromIntegral n * (sol VU.! k) | (n, sol) <- good] / totW
    dsegs = zipWith solveSeg zMats ys
    solveSeg z y =
        let r = VU.zipWith (-) y (matVec z dref)
            a = addDiagonal lam (gram z)
            rhs = tMatVec z r
         in case choleskySolve a rhs of
                Just eta -> VU.zipWith (+) dref eta
                Nothing -> dref
    toReg dseg =
        let bStd = dseg VU.! 0
            wStd = VU.drop 1 dseg
            rawCoef = VU.imap (\j w -> let s = sds VU.! j in if s == 0 then 0 else w / s) wStd
            adj =
                sum
                    [ let s = sds VU.! j
                       in if s == 0 then 0 else (wStd VU.! j) * (means VU.! j) / s
                    | j <- [0 .. d - 1]
                    ]
         in LinearRegressor rawCoef (bStd - adj) (V.fromList names) (lcPenalty cfg)

sq :: Double -> Double
sq x = x * x

-- | Add @lam@ to the diagonal of a square matrix.
addDiagonal :: Double -> Matrix -> Matrix
addDiagonal lam = V.imap (\i row -> row VU.// [(i, (row VU.! i) + lam)])