content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE InstanceSigs #-}
  {-# LANGUAGE NumericUnderscores #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module DataFrame.Lazy.Internal.DataFrame where
  
  import Control.Monad (foldM)
  import qualified Data.List as L
  import qualified Data.Text as T
  import qualified DataFrame.Internal.Column as C
  import qualified DataFrame.Internal.DataFrame as D
  import qualified DataFrame.Internal.Expression as E
  import qualified DataFrame.Lazy.IO.CSV as D
  import DataFrame.Operations.Merge ()
  import qualified DataFrame.Operations.Subset as D
  import qualified DataFrame.Operations.Transformations as D
  
  data LazyOperation where
      Derive :: (C.Columnable a) => T.Text -> E.Expr a -> LazyOperation
      Select :: [T.Text] -> LazyOperation
      Filter :: E.Expr Bool -> LazyOperation
  
  instance Show LazyOperation where
      show :: LazyOperation -> String
      show (Derive name expr) = T.unpack name ++ " := " ++ show expr
      show (Select columns) = "select(" ++ show columns ++ ")"
      show (Filter expr) = "filter(" ++ show expr ++ ")"
  
  data InputType = ICSV deriving (Show)
  
  data LazyDataFrame = LazyDataFrame
      { inputPath :: FilePath
      , inputType :: InputType
      , operations :: [LazyOperation]
      , batchSize :: Int
      }
      deriving (Show)
  
  eval :: LazyOperation -> D.DataFrame -> D.DataFrame
  eval (Derive name expr) = D.derive name expr
  eval (Select columns) = D.select columns
  eval (Filter expr) = D.filterWhere expr
  
  runDataFrame :: forall a. (C.Columnable a) => LazyDataFrame -> IO D.DataFrame
  runDataFrame df = do
      let path = inputPath df
      totalRows <- D.countRows ',' path
      let batches = batchRanges totalRows (batchSize df)
      (df', _) <-
          foldM
              ( \(accDf, (pos, unused, r)) (start, end) -> do
                  mapM_
                      putStr
                      [ "Scanning: "
                      , show start
                      , " to "
                      , show end
                      , " rows out of "
                      , show totalRows
                      , "\n"
                      ]
  
                  (sdf, (pos', unconsumed, rowsRead)) <-
                      D.readSeparated
                          ','
                          ( D.defaultOptions
                              { D.rowRange = Just (start, batchSize df)
                              , D.totalRows = Just totalRows
                              , D.seekPos = pos
                              , D.rowsRead = r
                              , D.leftOver = unused
                              }
                          )
                          path
                  let rdf = L.foldl' (flip eval) sdf (operations df)
                  return (accDf <> rdf, (Just pos', unconsumed, rowsRead + r))
              )
              (D.empty, (Nothing, "", 0))
              batches
      return df'
  
  batchRanges :: Int -> Int -> [(Int, Int)]
  batchRanges n inc = go n [0, inc .. n]
    where
      go _ [] = []
      go n [x] = [(x, n)]
      go n (f : s : rest) = (f, s) : go n (s : rest)
  
  scanCsv :: T.Text -> LazyDataFrame
  scanCsv path = LazyDataFrame (T.unpack path) ICSV [] 512_000
  
  addOperation :: LazyOperation -> LazyDataFrame -> LazyDataFrame
  addOperation op df = df{operations = operations df ++ [op]}
  
  derive ::
      (C.Columnable a) => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame
  derive name expr = addOperation (Derive name expr)
  
  select :: (C.Columnable a) => [T.Text] -> LazyDataFrame -> LazyDataFrame
  select columns = addOperation (Select columns)
  
  filter :: (C.Columnable a) => E.Expr Bool -> LazyDataFrame -> LazyDataFrame
  filter cond = addOperation (Filter cond)
  