Typing.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Operations.Typing where

import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector as V

import Control.Applicative (asum)
import Control.Monad (join)
import Data.Maybe (fromMaybe)
import qualified Data.Proxy as P
import Data.Time
import Data.Type.Equality (TestEquality (..))
import DataFrame.Internal.Column (Column (..), fromVector)
import DataFrame.Internal.DataFrame (DataFrame (..), unsafeGetColumn)
import DataFrame.Internal.Parsing
import DataFrame.Internal.Schema
import DataFrame.Operations.Core
import Text.Read
import Type.Reflection

type DateFormat = String

-- | Options controlling how text columns are parsed into typed values.
data ParseOptions = ParseOptions
    { missingValues :: [T.Text]
    -- ^ Values to treat as @Nothing@ when 'parseSafe' is @True@.
    , sampleSize :: Int
    -- ^ Number of rows to inspect when inferring a column's type (0 = all rows).
    , parseSafe :: Bool
    {- ^ When @True@, treat 'missingValues' and nullish strings as @Nothing@.
    When @False@, only empty strings become @Nothing@.
    -}
    , 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 = True
        , parseDateFormat = "%Y-%m-%d"
        }

parseDefaults :: ParseOptions -> DataFrame -> DataFrame
parseDefaults opts df = df{columns = V.map (parseDefault opts) (columns df)}

parseDefault :: ParseOptions -> Column -> Column
parseDefault opts (BoxedColumn (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 c
        Just Refl -> parseFromExamples opts c
parseDefault opts (OptionalColumn (c :: V.Vector (Maybe 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 . fromMaybe "") c)
            Nothing -> BoxedColumn c
        Just Refl -> parseFromExamples opts (V.map (fromMaybe "") c)
parseDefault _ column = column

parseFromExamples :: ParseOptions -> V.Vector T.Text -> Column
parseFromExamples opts cols =
    let
        converter =
            if parseSafe opts then convertNullish (missingValues opts) else convertOnlyEmpty
        examples = V.map converter (V.take (sampleSize opts) cols)
        asMaybeText = V.map converter cols
        dfmt = parseDateFormat opts
     in
        case makeParsingAssumption dfmt examples of
            BoolAssumption -> handleBoolAssumption asMaybeText
            IntAssumption -> handleIntAssumption asMaybeText
            DoubleAssumption -> handleDoubleAssumption asMaybeText
            TextAssumption -> handleTextAssumption asMaybeText
            DateAssumption -> handleDateAssumption dfmt asMaybeText
            NoAssumption -> handleNoAssumption dfmt asMaybeText

handleBoolAssumption :: V.Vector (Maybe T.Text) -> Column
handleBoolAssumption asMaybeText
    | parsableAsBool =
        maybe (fromVector asMaybeBool) fromVector (sequenceA asMaybeBool)
    | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText)
  where
    asMaybeBool = V.map (>>= readBool) asMaybeText
    parsableAsBool = vecSameConstructor asMaybeText asMaybeBool

handleIntAssumption :: V.Vector (Maybe T.Text) -> Column
handleIntAssumption asMaybeText
    | parsableAsInt =
        maybe (fromVector asMaybeInt) fromVector (sequenceA asMaybeInt)
    | parsableAsDouble =
        maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble)
    | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText)
  where
    asMaybeInt = V.map (>>= readInt) asMaybeText
    asMaybeDouble = V.map (>>= readDouble) asMaybeText
    parsableAsInt =
        vecSameConstructor asMaybeText asMaybeInt
            && vecSameConstructor asMaybeText asMaybeDouble
    parsableAsDouble = vecSameConstructor asMaybeText asMaybeDouble

handleDoubleAssumption :: V.Vector (Maybe T.Text) -> Column
handleDoubleAssumption asMaybeText
    | parsableAsDouble =
        maybe (fromVector asMaybeDouble) fromVector (sequenceA asMaybeDouble)
    | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText)
  where
    asMaybeDouble = V.map (>>= readDouble) asMaybeText
    parsableAsDouble = vecSameConstructor asMaybeText asMaybeDouble

handleDateAssumption :: DateFormat -> V.Vector (Maybe T.Text) -> Column
handleDateAssumption dateFormat asMaybeText
    | parsableAsDate =
        maybe (fromVector asMaybeDate) fromVector (sequenceA asMaybeDate)
    | otherwise = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText)
  where
    asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText
    parsableAsDate = vecSameConstructor asMaybeText asMaybeDate

handleTextAssumption :: V.Vector (Maybe T.Text) -> Column
handleTextAssumption asMaybeText = maybe (fromVector asMaybeText) fromVector (sequenceA asMaybeText)

handleNoAssumption :: DateFormat -> V.Vector (Maybe T.Text) -> Column
handleNoAssumption dateFormat asMaybeText
    -- No need to check for null values. If we are in this condition, that
    -- means that the examples consisted only of null values, so we can
    -- confidently know that this column must be an OptionalColumn
    | V.all (== Nothing) asMaybeText = fromVector asMaybeText
    | parsableAsBool = fromVector asMaybeBool
    | parsableAsInt = fromVector asMaybeInt
    | parsableAsDouble = fromVector asMaybeDouble
    | parsableAsDate = fromVector asMaybeDate
    | otherwise = fromVector asMaybeText
  where
    asMaybeBool = V.map (>>= readBool) asMaybeText
    asMaybeInt = V.map (>>= readInt) asMaybeText
    asMaybeDouble = V.map (>>= readDouble) asMaybeText
    asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText
    parsableAsBool = vecSameConstructor asMaybeText asMaybeBool
    parsableAsInt =
        vecSameConstructor asMaybeText asMaybeInt
            && vecSameConstructor asMaybeText asMaybeDouble
    parsableAsDouble = vecSameConstructor asMaybeText asMaybeDouble
    parsableAsDate = vecSameConstructor asMaybeText asMaybeDate

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

parseTimeOpt :: DateFormat -> T.Text -> Maybe Day
parseTimeOpt dateFormat s =
    parseTimeM {- Accept leading/trailing whitespace -}
        True
        defaultTimeLocale
        dateFormat
        (T.unpack s)

unsafeParseTime :: DateFormat -> T.Text -> Day
unsafeParseTime dateFormat s =
    parseTimeOrError {- Accept leading/trailing whitespace -}
        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

makeParsingAssumption ::
    DateFormat -> V.Vector (Maybe T.Text) -> ParsingAssumption
makeParsingAssumption dateFormat asMaybeText
    -- All the examples are "NA", "Null", "", so we can't make any shortcut
    -- assumptions and just have to go the long way.
    | V.all (== Nothing) asMaybeText = NoAssumption
    -- After accounting for nulls, parsing for Ints and Doubles results in the
    -- same corresponding positions of Justs and Nothings, so we assume
    -- that the best way to parse is Int
    | vecSameConstructor asMaybeText asMaybeBool = BoolAssumption
    | vecSameConstructor asMaybeText asMaybeInt
        && vecSameConstructor asMaybeText asMaybeDouble =
        IntAssumption
    -- After accounting for nulls, the previous condition fails, so some (or none) can be parsed as Ints
    -- and some can be parsed as Doubles, so we make the assumpotion of doubles.
    | vecSameConstructor asMaybeText asMaybeDouble = DoubleAssumption
    -- After accounting for nulls, parsing for Dates results in the same corresponding
    -- positions of Justs and Nothings, so we assume that the best way to parse is Date.
    | vecSameConstructor asMaybeText asMaybeDate = DateAssumption
    | otherwise = TextAssumption
  where
    asMaybeBool = V.map (>>= readBool) asMaybeText
    asMaybeInt = V.map (>>= readInt) asMaybeText
    asMaybeDouble = V.map (>>= readDouble) asMaybeText
    asMaybeDate = V.map (>>= parseTimeOpt dateFormat) asMaybeText

data ParsingAssumption
    = BoolAssumption
    | IntAssumption
    | DoubleAssumption
    | DateAssumption
    | NoAssumption
    | TextAssumption

parseWithTypes :: Bool -> M.Map T.Text SchemaType -> DataFrame -> DataFrame
parseWithTypes safe ts df
    | M.null ts = df
    | otherwise =
        M.foldrWithKey
            (\k v d -> insertColumn k (asType v (unsafeGetColumn k d)) d)
            df
            ts
  where
    asType :: SchemaType -> Column -> Column
    asType (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 ->
                                if safe
                                    then fromVector (V.map ((readMaybe @a) . T.unpack) col)
                                    else fromVector (V.map ((read @a) . T.unpack) col)
                            Nothing ->
                                if safe
                                    then fromVector (V.map ((readMaybe @a) . show) col)
                                    else fromVector (V.map ((read @a) . show) col)
                _ -> c
        _ -> case testEquality (typeRep @a) (typeRep @b) of
            Just Refl -> c
            Nothing -> case testEquality (typeRep @T.Text) (typeRep @b) of
                Just Refl ->
                    if safe
                        then fromVector (V.map ((readMaybe @a) . T.unpack) col)
                        else fromVector (V.map ((read @a) . T.unpack) col)
                Nothing ->
                    if safe
                        then fromVector (V.map ((readMaybe @a) . show) col)
                        else fromVector (V.map ((read @a) . show) col)
    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