Typing.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- The inference lattice (sample classification, 'ParsingAssumption',
-- Int -> Double promotion) lives in "DataFrame.Operations.Inference"
-- and is re-exported here for backwards compatibility.
module DataFrame.Operations.Typing (
    module DataFrame.Operations.Typing,
    module DataFrame.Operations.Inference,
) where

import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Applicative (asum)
import Control.Monad (join)
import Control.Monad.ST (runST)
import Data.Maybe (fromMaybe)
import qualified Data.Proxy as P
import Data.Time
import Data.Type.Equality (TestEquality (..))
import DataFrame.Internal.Column (
    Bitmap,
    Column (..),
    Columnable,
    bitmapTestBit,
    ensureOptional,
    finalizeParseResult,
    fromVector,
    materializePacked,
 )
import DataFrame.Internal.DataFrame (
    DataFrame (..),
    insertColumn,
    unsafeGetColumn,
 )
import DataFrame.Internal.Parsing
import DataFrame.Operations.Core ()
import DataFrame.Operations.Inference
import DataFrame.Schema
import Text.Read
import Type.Reflection

{- | How parse failures are surfaced: 'NoSafeRead' throws, 'MaybeRead' yields
@Nothing@ (column wrapped @Maybe a@), 'EitherRead' yields @Left rawText@
(column wrapped @Either Text a@, preserving the original input).
-}
data SafeReadMode
    = NoSafeRead
    | MaybeRead
    | EitherRead
    deriving (Eq, Show, Read)

-- | Options controlling how text columns are parsed into typed values.
data ParseOptions = ParseOptions
    { missingValues :: [T.Text]
    -- ^ Values to treat as @Nothing@ when the effective mode is 'MaybeRead'.
    , sampleSize :: Int
    -- ^ Number of rows to inspect when inferring a column's type (0 = all rows).
    , parseSafe :: SafeReadMode
    -- ^ Default 'SafeReadMode' for columns without a 'parseSafeOverrides' entry.
    , parseSafeOverrides :: [(T.Text, SafeReadMode)]
    {- ^ Per-column overrides taking precedence over 'parseSafe' — e.g. strict
    IDs (@NoSafeRead@) alongside lenient fields (@MaybeRead@/@EitherRead@).
    -}
    , parseDateFormat :: DateFormat
    -- ^ Date format string as accepted by "Data.Time.Format" (e.g. @\"%Y-%m-%d\"@).
    }

{- | Sensible out-of-the-box parse options: infer from the first 100 rows,
  treat common nullish strings as missing, and expect ISO 8601 dates.
-}
defaultParseOptions :: ParseOptions
defaultParseOptions =
    ParseOptions
        { missingValues = []
        , sampleSize = 100
        , parseSafe = MaybeRead
        , parseSafeOverrides = []
        , parseDateFormat = "%Y-%m-%d"
        }

{- | Resolve a column's effective 'SafeReadMode': the override if present,
otherwise the default.
-}
effectiveSafeRead ::
    SafeReadMode -> [(T.Text, SafeReadMode)] -> T.Text -> SafeReadMode
effectiveSafeRead def overrides name = fromMaybe def (lookup name overrides)

parseDefaults :: ParseOptions -> DataFrame -> DataFrame
parseDefaults opts df = df{columns = V.imap forCol (columns df)}
  where
    nameAt =
        let inverted = M.fromList [(i, n) | (n, i) <- M.toList (columnIndices df)]
         in \i -> M.findWithDefault "" i inverted
    forCol i col =
        let mode =
                effectiveSafeRead
                    (parseSafe opts)
                    (parseSafeOverrides opts)
                    (nameAt i)
         in parseDefault opts{parseSafe = mode, parseSafeOverrides = []} col

parseDefault :: ParseOptions -> Column -> Column
parseDefault opts (BoxedColumn Nothing (c :: V.Vector a)) =
    case (typeRep @a) `testEquality` (typeRep @T.Text) of
        Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of
            Just Refl -> parseFromExamples opts (V.map T.pack c)
            Nothing -> BoxedColumn Nothing c
        Just Refl -> parseFromExamples opts c
parseDefault opts (BoxedColumn (Just bm) (c :: V.Vector a)) =
    case (typeRep @a) `testEquality` (typeRep @T.Text) of
        Nothing -> case (typeRep @a) `testEquality` (typeRep @String) of
            Just Refl ->
                parseFromExamples
                    opts
                    (V.imap (\i x -> if bitmapTestBit bm i then T.pack x else "") c)
            Nothing -> BoxedColumn (Just bm) c
        Just Refl ->
            parseFromExamples opts (V.imap (\i x -> if bitmapTestBit bm i then x else "") c)
parseDefault _ column = column

parseFromExamples :: ParseOptions -> V.Vector T.Text -> Column
parseFromExamples opts cols =
    let isNull = case parseSafe opts of
            NoSafeRead -> T.null
            _ -> isNullishOrMissing (missingValues opts)
        examples = V.map (classify isNull) (V.take (sampleSize opts) cols)
        dfmt = parseDateFormat opts
        assumption = makeParsingAssumption dfmt examples
     in case parseSafe opts of
            EitherRead -> handleEitherAssumption dfmt assumption cols
            mode ->
                let result = case assumption of
                        BoolAssumption -> handleBoolAssumption isNull cols
                        IntAssumption -> handleIntAssumption isNull cols
                        DoubleAssumption -> handleDoubleAssumption isNull cols
                        TextAssumption -> handleTextAssumption isNull cols
                        DateAssumption -> handleDateAssumption dfmt isNull cols
                        NoAssumption -> handleNoAssumption dfmt isNull cols
                 in if mode == MaybeRead then ensureOptional result else result
  where
    classify p t = if p t then Nothing else Just t

{- | For 'EitherRead' mode: parse under the chosen assumption into an
@Either Text a@ column. Successful parses become @Right@; failures (including
null/missing cells) become @Left@ carrying the raw input verbatim.
-}
handleEitherAssumption ::
    DateFormat -> ParsingAssumption -> V.Vector T.Text -> Column
handleEitherAssumption dfmt assumption raw = case assumption of
    BoolAssumption -> fromVector (V.map (toEither readBool) raw)
    IntAssumption -> fromVector (V.map (toEither readInt) raw)
    DoubleAssumption -> fromVector (V.map (toEither readDouble) raw)
    DateAssumption -> fromVector (V.map (toEither (parseTimeOpt dfmt)) raw)
    TextAssumption -> fromVector (V.map textToEither raw)
    NoAssumption -> fromVector (V.map textToEither raw)
  where
    toEither :: (T.Text -> Maybe a) -> T.Text -> Either T.Text a
    toEither p t = maybe (Left t) Right (p t)

    textToEither :: T.Text -> Either T.Text T.Text
    textToEither t = if T.null t then Left t else Right t

parseUnboxedColumnWithPred ::
    forall src a.
    (VU.Unbox a) =>
    a ->
    (src -> Bool) ->
    (src -> Maybe a) ->
    V.Vector src ->
    Maybe (Maybe Bitmap, VU.Vector a)
parseUnboxedColumnWithPred nullValue isNull parser vec = runST $ do
    let n = V.length vec
    values <- VUM.unsafeNew n
    vmask <- VUM.unsafeNew n
    let go !i !anyNull
            | i >= n = finalizeParseResult values vmask anyNull
            | otherwise =
                let !src = V.unsafeIndex vec i
                 in if isNull src
                        then do
                            VUM.unsafeWrite vmask i 0
                            VUM.unsafeWrite values i nullValue
                            go (i + 1) True
                        else case parser src of
                            Just v -> do
                                VUM.unsafeWrite vmask i 1
                                VUM.unsafeWrite values i v
                                go (i + 1) anyNull
                            Nothing -> return Nothing
    go 0 False
{-# INLINE parseUnboxedColumnWithPred #-}

-- | Wrap a successful 'parseUnboxedColumnWithPred' result as a 'Column'.
unboxedOrFallback ::
    (Columnable a, VU.Unbox a) =>
    Maybe (Maybe Bitmap, VU.Vector a) ->
    Column ->
    Column
unboxedOrFallback (Just (mbm, vec)) _ = UnboxedColumn mbm vec
unboxedOrFallback Nothing fallback = fallback

handleBoolAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column
handleBoolAssumption isNull cols =
    unboxedOrFallback
        (parseUnboxedColumnWithPred False isNull readBool cols)
        (handleTextAssumption isNull cols)

{- | Int columns: one fused pass with in-place Int -> Double promotion; a cell
parsing as neither demotes the column to Text. 'readIntStrict' rejects overflow
so a huge integer promotes to 'Double' rather than wrapping.
-}
handleIntAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column
handleIntAssumption isNull cols =
    case promoteIntColumn (\_ t -> isNull t) readIntStrict readDouble cols of
        Just col -> col
        Nothing -> handleTextAssumption isNull cols

handleDoubleAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column
handleDoubleAssumption isNull cols =
    unboxedOrFallback
        (parseUnboxedColumnWithPred 0 isNull readDouble cols)
        (handleTextAssumption isNull cols)

{- | Text columns: no parse, just null-marking. An all-non-null column stays a
plain @V.Vector T.Text@; otherwise it becomes @V.Vector (Maybe T.Text)@.
-}
handleTextAssumption :: (T.Text -> Bool) -> V.Vector T.Text -> Column
handleTextAssumption isNull cols
    | V.any isNull cols =
        fromVector
            (V.map (\t -> if isNull t then Nothing else Just t) cols)
    | otherwise = fromVector cols

{- | Date: single boxed parse pass ('Day' is not unboxable). Bails to
'handleTextAssumption' the moment a non-null cell fails to parse as a 'Day'.
A column with no nulls keeps type 'Day' rather than 'Maybe Day'.
-}
handleDateAssumption ::
    DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column
handleDateAssumption dateFormat isNull cols =
    case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of
        Just (anyNull, vec)
            | anyNull -> fromVector vec
            | otherwise -> fromVector (V.mapMaybe id vec)
        Nothing -> handleTextAssumption isNull cols

parseBoxedMaybeColumn ::
    (T.Text -> Bool) ->
    (T.Text -> Maybe a) ->
    V.Vector T.Text ->
    Maybe (Bool, V.Vector (Maybe a))
parseBoxedMaybeColumn isNull parser cols = runST $ do
    let n = V.length cols
    out <- VM.new n
    let loop !i !anyNull
            | i >= n = do
                frozen <- V.unsafeFreeze out
                return (Just (anyNull, frozen))
            | otherwise =
                let !t = V.unsafeIndex cols i
                 in if isNull t
                        then do
                            VM.unsafeWrite out i Nothing
                            loop (i + 1) True
                        else case parser t of
                            Just v -> do
                                VM.unsafeWrite out i (Just v)
                                loop (i + 1) anyNull
                            Nothing -> return Nothing
    loop 0 False

-- Reached only when the sample was all-null: try each concrete type in turn,
-- falling back to Text. A column with no nulls keeps type 'Day', not 'Maybe Day'.
handleNoAssumption ::
    DateFormat -> (T.Text -> Bool) -> V.Vector T.Text -> Column
handleNoAssumption dateFormat isNull cols
    | V.all isNull cols =
        fromVector (V.map (const (Nothing :: Maybe T.Text)) cols)
    | Just (mbm, vec) <- parseUnboxedColumnWithPred False isNull readBool cols =
        UnboxedColumn mbm vec
    | Just (mbm, vec) <- parseUnboxedColumnWithPred 0 isNull readInt cols =
        UnboxedColumn mbm vec
    | Just (mbm, vec) <- parseUnboxedColumnWithPred 0 isNull readDouble cols =
        UnboxedColumn mbm vec
    | otherwise = case parseBoxedMaybeColumn isNull (parseTimeOpt dateFormat) cols of
        Just (anyNull, vec)
            | anyNull -> fromVector vec
            | otherwise -> fromVector (V.mapMaybe id vec)
        Nothing -> handleTextAssumption isNull cols

{- | True for nullish or explicitly-listed missing strings. ('convertNullish'
and 'convertOnlyEmpty' below are kept only for external callers.)
-}
isNullishOrMissing :: [T.Text] -> T.Text -> Bool
isNullishOrMissing missing v = isNullish v || v `elem` missing

convertNullish :: [T.Text] -> T.Text -> Maybe T.Text
convertNullish missing v = if isNullish v || v `elem` missing then Nothing else Just v

convertOnlyEmpty :: T.Text -> Maybe T.Text
convertOnlyEmpty v = if v == "" then Nothing else Just v

unsafeParseTime :: DateFormat -> T.Text -> Day
unsafeParseTime dateFormat s =
    parseTimeOrError
        True
        defaultTimeLocale
        dateFormat
        (T.unpack s)

hasNullValues :: (Eq a) => V.Vector (Maybe a) -> Bool
hasNullValues = V.any (== Nothing)

vecSameConstructor :: V.Vector (Maybe a) -> V.Vector (Maybe b) -> Bool
vecSameConstructor xs ys = (V.length xs == V.length ys) && V.and (V.zipWith hasSameConstructor xs ys)
  where
    hasSameConstructor :: Maybe a -> Maybe b -> Bool
    hasSameConstructor (Just _) (Just _) = True
    hasSameConstructor Nothing Nothing = True
    hasSameConstructor _ _ = False

{- | Re-type columns of a 'DataFrame' according to a schema map. @resolveMode@
maps a column name to its 'SafeReadMode' (typically via 'effectiveSafeRead').
-}
parseWithTypes ::
    (T.Text -> SafeReadMode) ->
    M.Map T.Text SchemaType ->
    DataFrame ->
    DataFrame
parseWithTypes resolveMode ts df
    | M.null ts = df
    | otherwise =
        M.foldrWithKey
            (\k v d -> insertColumn k (asType (resolveMode k) v (unsafeGetColumn k d)) d)
            df
            ts
  where
    plainType ::
        forall a b.
        (Columnable a, Read a) =>
        SafeReadMode -> V.Vector b -> (b -> String) -> Column
    plainType mode col toStr = case mode of
        NoSafeRead -> fromVector (V.map ((read @a) . toStr) col)
        MaybeRead -> fromVector (V.map ((readMaybe @a) . toStr) col)
        EitherRead -> fromVector (V.map ((readEitherRaw @a) . toStr) col)

    asType :: SafeReadMode -> SchemaType -> Column -> Column
    asType mode st c@(PackedText _ _) = asType mode st (materializePacked c)
    asType mode (SType (_ :: P.Proxy a)) c@(BoxedColumn _ (col :: V.Vector b)) = case typeRep @a of
        App t1 _t2 -> case eqTypeRep t1 (typeRep @Maybe) of
            Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of
                Just Refl -> c
                Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of
                    Just Refl -> fromVector (V.map (join . (readAsMaybe @a) . T.unpack) col)
                    Nothing -> fromVector (V.map (join . (readAsMaybe @a) . show) col)
            Nothing -> case t1 of
                App t1' _t2' -> case eqTypeRep t1' (typeRep @Either) of
                    Just HRefl -> case testEquality (typeRep @a) (typeRep @b) of
                        Just Refl -> c
                        Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of
                            Just Refl -> fromVector (V.map ((readAsEither @a) . T.unpack) col)
                            Nothing -> fromVector (V.map ((readAsEither @a) . show) col)
                    Nothing -> case testEquality (typeRep @a) (typeRep @b) of
                        Just Refl -> c
                        Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of
                            Just Refl -> plainType @a mode col T.unpack
                            Nothing -> plainType @a mode col show
                _ -> c
        _ -> case testEquality (typeRep @a) (typeRep @b) of
            Just Refl -> c
            Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of
                Just Refl -> plainType @a mode col T.unpack
                Nothing -> plainType @a mode col show
    asType _ _ c = c

readAsMaybe :: (Read a) => String -> Maybe a
readAsMaybe s
    | null s = Nothing
    | otherwise = readMaybe $ "Just " <> s

readAsEither :: (Read a) => String -> a
readAsEither v = case asum [readMaybe $ "Left " <> s, readMaybe $ "Right " <> s] of
    Nothing -> error $ "Couldn't read value: " <> s
    Just v' -> v'
  where
    s = if null v then "\"\"" else v

{- | Try 'readMaybe'; on failure return @Left raw@ where @raw@ is the original
input text. Used by 'parseWithTypes' under 'EitherRead'.
-}
readEitherRaw :: forall a. (Read a) => String -> Either T.Text a
readEitherRaw s = case readMaybe s of
    Just v -> Right v
    Nothing -> Left (T.pack s)