{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Typed CSV reading.

The schema is the whole specification of the read: it names the columns to
fetch and gives their types, so these readers touch only the columns @cols@
declares and never infer a type the schema already knows.

@
type Customer = '[ '(\"customer_id\", Int), '(\"customer_name\", Text)]

customers <- readCsv \@Customer \"customers.csv\"  -- reads 2 columns, whatever
                                                 -- else the file holds
@
-}
module DataFrame.Typed.IO.CSV (
    readCsv,
    readCsvWithError,
    readTsv,
    readCsvWithOpts,
    writeCsv,
    writeTsv,
) where

import Control.Applicative ((<|>))
import Control.Exception (SomeException, try)
import qualified Data.Text as T

import DataFrame.IO.CSV (ReadOptions (..), TypeSpec (..), defaultReadOptions)
import qualified DataFrame.IO.CSV as CSV
import DataFrame.Schema (RuntimeSchema (..), elements)
import DataFrame.Typed.Freeze (freezeOrThrow, freezeWithError, thaw)
import DataFrame.Typed.Schema (KnownSchema)
import DataFrame.Typed.Types (TypedDataFrame)

import qualified Data.Map.Strict as M

{- | Fold a type-level schema into read options: fetch exactly its columns,
and parse them at its types. Anything else the caller asked for is kept,
including a wider 'typeSpec' for columns the schema does not name.
-}
schemaOptions :: forall cols. (RuntimeSchema cols) => ReadOptions -> ReadOptions
schemaOptions :: forall (cols :: [(Symbol, *)]).
RuntimeSchema cols =>
ReadOptions -> ReadOptions
schemaOptions ReadOptions
opts =
    ReadOptions
opts
        { typeSpec = SpecifyTypes (M.toList declared) (typeSpec opts)
        , readColumns = readColumns opts <|> Just (M.keys declared)
        }
  where
    declared :: Map Text SchemaType
declared = Schema -> Map Text SchemaType
elements (forall (cols :: [(Symbol, *)]). RuntimeSchema cols => Schema
runtimeSchema @cols)

{- | Read a CSV file into a typed DataFrame, throwing on schema mismatch.
Reads only the columns @cols@ names, typed as @cols@ says.

==== __Example__
@
ghci> type Customer = '[ '(\"id\", Int), '(\"name\", Text)]
ghci> customers <- readCsv \@Customer \"customers.csv\"
@
-}
readCsv ::
    forall cols.
    (KnownSchema cols, RuntimeSchema cols) =>
    FilePath -> IO (TypedDataFrame cols)
readCsv :: forall (cols :: [(Symbol, *)]).
(KnownSchema cols, RuntimeSchema cols) =>
FilePath -> IO (TypedDataFrame cols)
readCsv = forall (cols :: [(Symbol, *)]).
(KnownSchema cols, RuntimeSchema cols) =>
ReadOptions -> FilePath -> IO (TypedDataFrame cols)
readCsvWithOpts @cols ReadOptions
defaultReadOptions

{- | Read a CSV file, returning a descriptive error on schema mismatch or a
missing column instead of throwing.

==== __Example__
@
ghci> readCsvWithError \@Customer \"customers.csv\"
Right (TDF ...)
@
-}
readCsvWithError ::
    forall cols.
    (KnownSchema cols, RuntimeSchema cols) =>
    FilePath -> IO (Either T.Text (TypedDataFrame cols))
readCsvWithError :: forall (cols :: [(Symbol, *)]).
(KnownSchema cols, RuntimeSchema cols) =>
FilePath -> IO (Either Text (TypedDataFrame cols))
readCsvWithError FilePath
path = do
    -- The projection can fail before the freeze does (a column the file does
    -- not have); this reader reports rather than throws either way.
    Either SomeException DataFrame
r <- IO DataFrame -> IO (Either SomeException DataFrame)
forall e a. Exception e => IO a -> IO (Either e a)
try (ReadOptions -> FilePath -> IO DataFrame
CSV.readCsvWithOpts (forall (cols :: [(Symbol, *)]).
RuntimeSchema cols =>
ReadOptions -> ReadOptions
schemaOptions @cols ReadOptions
defaultReadOptions) FilePath
path)
    Either Text (TypedDataFrame cols)
-> IO (Either Text (TypedDataFrame cols))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text (TypedDataFrame cols)
 -> IO (Either Text (TypedDataFrame cols)))
-> Either Text (TypedDataFrame cols)
-> IO (Either Text (TypedDataFrame cols))
forall a b. (a -> b) -> a -> b
$ case Either SomeException DataFrame
r of
        Left (SomeException
e :: SomeException) -> Text -> Either Text (TypedDataFrame cols)
forall a b. a -> Either a b
Left (FilePath -> Text
T.pack (SomeException -> FilePath
forall a. Show a => a -> FilePath
show SomeException
e))
        Right DataFrame
df -> DataFrame -> Either Text (TypedDataFrame cols)
forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> Either Text (TypedDataFrame cols)
freezeWithError DataFrame
df

{- | Read a tab-separated file into a typed DataFrame, throwing on mismatch.

==== __Example__
@
ghci> customers <- readTsv \@Customer \"customers.tsv\"
@
-}
readTsv ::
    forall cols.
    (KnownSchema cols, RuntimeSchema cols) =>
    FilePath -> IO (TypedDataFrame cols)
readTsv :: forall (cols :: [(Symbol, *)]).
(KnownSchema cols, RuntimeSchema cols) =>
FilePath -> IO (TypedDataFrame cols)
readTsv = forall (cols :: [(Symbol, *)]).
(KnownSchema cols, RuntimeSchema cols) =>
ReadOptions -> FilePath -> IO (TypedDataFrame cols)
readCsvWithOpts @cols ReadOptions
defaultReadOptions{columnSeparator = '\t'}

{- | Read a CSV file with custom options, throwing on schema mismatch. The
schema still supplies the column selection and types; an explicit
'readColumns' takes precedence, and is then checked by the freeze.

==== __Example__
@
ghci> customers <- readCsvWithOpts \@Customer defaultReadOptions{safeRead = MaybeRead} \"customers.csv\"
@
-}
readCsvWithOpts ::
    forall cols.
    (KnownSchema cols, RuntimeSchema cols) =>
    ReadOptions -> FilePath -> IO (TypedDataFrame cols)
readCsvWithOpts :: forall (cols :: [(Symbol, *)]).
(KnownSchema cols, RuntimeSchema cols) =>
ReadOptions -> FilePath -> IO (TypedDataFrame cols)
readCsvWithOpts ReadOptions
opts FilePath
path =
    ReadOptions -> FilePath -> IO DataFrame
CSV.readCsvWithOpts (forall (cols :: [(Symbol, *)]).
RuntimeSchema cols =>
ReadOptions -> ReadOptions
schemaOptions @cols ReadOptions
opts) FilePath
path IO DataFrame
-> (DataFrame -> IO (TypedDataFrame cols))
-> IO (TypedDataFrame cols)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (cols :: [(Symbol, *)]).
KnownSchema cols =>
DataFrame -> IO (TypedDataFrame cols)
freezeOrThrow @cols

{- | Write a typed DataFrame to a CSV file.

==== __Example__
@
ghci> writeCsv \"customers.csv\" customers
@
-}
writeCsv :: FilePath -> TypedDataFrame cols -> IO ()
writeCsv :: forall (cols :: [(Symbol, *)]).
FilePath -> TypedDataFrame cols -> IO ()
writeCsv FilePath
path = FilePath -> DataFrame -> IO ()
CSV.writeCsv FilePath
path (DataFrame -> IO ())
-> (TypedDataFrame cols -> DataFrame)
-> TypedDataFrame cols
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
thaw

{- | Write a typed DataFrame to a tab-separated file.

==== __Example__
@
ghci> writeTsv \"customers.tsv\" customers
@
-}
writeTsv :: FilePath -> TypedDataFrame cols -> IO ()
writeTsv :: forall (cols :: [(Symbol, *)]).
FilePath -> TypedDataFrame cols -> IO ()
writeTsv FilePath
path = FilePath -> DataFrame -> IO ()
CSV.writeTsv FilePath
path (DataFrame -> IO ())
-> (TypedDataFrame cols -> DataFrame)
-> TypedDataFrame cols
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
thaw