content: {-# LANGUAGE OverloadedStrings #-}
  
  module DataFrame.IO.Parquet.Dictionary where
  
  import Control.Monad
  import Data.Bits
  import Data.Char
  import Data.Int
  import Data.Maybe
  import qualified Data.Text as T
  import Data.Time
  import Data.Word
  import DataFrame.IO.Parquet.Binary
  import DataFrame.IO.Parquet.Encoding
  import DataFrame.IO.Parquet.Levels
  import DataFrame.IO.Parquet.Time
  import DataFrame.IO.Parquet.Types
  import qualified DataFrame.Internal.Column as DI
  import GHC.Float
  
  dictCardinality :: DictVals -> Int
  dictCardinality (DBool ds) = length ds
  dictCardinality (DInt32 ds) = length ds
  dictCardinality (DInt64 ds) = length ds
  dictCardinality (DInt96 ds) = length ds
  dictCardinality (DFloat ds) = length ds
  dictCardinality (DDouble ds) = length ds
  dictCardinality (DText ds) = length ds
  
  readDictVals :: ParquetType -> [Word8] -> Maybe Int32 -> DictVals
  readDictVals PBOOLEAN bs (Just count) = DBool (take (fromIntegral count) $ readPageBool bs)
  readDictVals PINT32 bs _ = DInt32 (readPageInt32 bs)
  readDictVals PINT64 bs _ = DInt64 (readPageInt64 bs)
  readDictVals PINT96 bs _ = DInt96 (readPageInt96Times bs)
  readDictVals PFLOAT bs _ = DFloat (readPageFloat bs)
  readDictVals PDOUBLE bs _ = DDouble (readPageWord64 bs)
  readDictVals PBYTE_ARRAY bs _ = DText (readPageBytes bs)
  readDictVals PFIXED_LEN_BYTE_ARRAY bs (Just len) = DText (readPageFixedBytes bs (fromIntegral len))
  readDictVals t _ _ = error $ "Unsupported dictionary type: " ++ show t
  
  readPageInt32 :: [Word8] -> [Int32]
  readPageInt32 [] = []
  readPageInt32 xs = littleEndianInt32 (take 4 xs) : readPageInt32 (drop 4 xs)
  
  readPageWord64 :: [Word8] -> [Double]
  readPageWord64 [] = []
  readPageWord64 xs =
      castWord64ToDouble (littleEndianWord64 (take 8 xs)) : readPageWord64 (drop 8 xs)
  
  readPageBytes :: [Word8] -> [T.Text]
  readPageBytes [] = []
  readPageBytes xs =
      let lenBytes = fromIntegral (littleEndianInt32 $ take 4 xs)
          totalBytesRead = lenBytes + 4
       in T.pack (map (chr . fromIntegral) $ take lenBytes (drop 4 xs))
              : readPageBytes (drop totalBytesRead xs)
  
  readPageBool :: [Word8] -> [Bool]
  readPageBool [] = []
  readPageBool bs =
      concatMap (\b -> map (\i -> (b `shiftR` i) .&. 1 == 1) [0 .. 7]) bs
  
  readPageInt64 :: [Word8] -> [Int64]
  readPageInt64 [] = []
  readPageInt64 xs = fromIntegral (littleEndianWord64 (take 8 xs)) : readPageInt64 (drop 8 xs)
  
  readPageFloat :: [Word8] -> [Float]
  readPageFloat [] = []
  readPageFloat xs =
      castWord32ToFloat (littleEndianWord32 (take 4 xs)) : readPageFloat (drop 4 xs)
  
  readNInt96Times :: Int -> [Word8] -> ([UTCTime], [Word8])
  readNInt96Times 0 bs = ([], bs)
  readNInt96Times k bs =
      let timestamp96 = take 12 bs
          utcTime = int96ToUTCTime timestamp96
          bs' = drop 12 bs
          (times, rest) = readNInt96Times (k - 1) bs'
       in (utcTime : times, rest)
  
  readPageInt96Times :: [Word8] -> [UTCTime]
  readPageInt96Times [] = []
  readPageInt96Times bs =
      let (times, _) = readNInt96Times (length bs `div` 12) bs
       in times
  
  readPageFixedBytes :: [Word8] -> Int -> [T.Text]
  readPageFixedBytes [] _ = []
  readPageFixedBytes xs len =
      let chunk = take len xs
          text = T.pack (map (chr . fromIntegral) chunk)
       in text : readPageFixedBytes (drop len xs) len
  
  decodeDictV1 :: Maybe DictVals -> Int -> [Int] -> Int -> [Word8] -> IO DI.Column
  decodeDictV1 dictValsM maxDef defLvls nPresent bytes =
      case dictValsM of
          Nothing -> error "Dictionary-encoded page but dictionary is missing"
          Just dictVals ->
              let (idxs, _rest) = decodeDictIndicesV1 nPresent (dictCardinality dictVals) bytes
               in do
                      when (length idxs /= nPresent) $
                          error $
                              "dict index count mismatch: got "
                                  ++ show (length idxs)
                                  ++ ", expected "
                                  ++ show nPresent
                      case dictVals of
                          DBool ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeBool maxDef defLvls values)
                          DInt32 ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeInt32 maxDef defLvls values)
                          DInt64 ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeInt64 maxDef defLvls values)
                          DInt96 ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeUTCTime maxDef defLvls values)
                          DFloat ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeFloat maxDef defLvls values)
                          DDouble ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeDouble maxDef defLvls values)
                          DText ds -> do
                              let values = [ds !! i | i <- idxs]
                              pure (toMaybeText maxDef defLvls values)
  
  toMaybeInt32 :: Int -> [Int] -> [Int32] -> DI.Column
  toMaybeInt32 maxDef def xs =
      let filled = stitchNullable maxDef def xs
       in if all isJust filled
              then DI.fromList (map (fromMaybe 0) filled)
              else DI.fromList filled
  
  toMaybeDouble :: Int -> [Int] -> [Double] -> DI.Column
  toMaybeDouble maxDef def xs =
      let filled = stitchNullable maxDef def xs
       in if all isJust filled
              then DI.fromList (map (fromMaybe 0) filled)
              else DI.fromList filled
  
  toMaybeText :: Int -> [Int] -> [T.Text] -> DI.Column
  toMaybeText maxDef def xs =
      let filled = stitchNullable maxDef def xs
       in if all isJust filled
              then DI.fromList (map (fromMaybe "") filled)
              else DI.fromList filled
  
  toMaybeBool :: Int -> [Int] -> [Bool] -> DI.Column
  toMaybeBool maxDef def xs =
      let filled = stitchNullable maxDef def xs
       in if all isJust filled
              then DI.fromList (map (fromMaybe False) filled)
              else DI.fromList filled
  
  toMaybeInt64 :: Int -> [Int] -> [Int64] -> DI.Column
  toMaybeInt64 maxDef def xs =
      let filled = stitchNullable maxDef def xs
       in if all isJust filled
              then DI.fromList (map (fromMaybe 0) filled)
              else DI.fromList filled
  
  toMaybeFloat :: Int -> [Int] -> [Float] -> DI.Column
  toMaybeFloat maxDef def xs =
      let filled = stitchNullable maxDef def xs
       in if all isJust filled
              then DI.fromList (map (fromMaybe 0.0) filled)
              else DI.fromList filled
  
  toMaybeUTCTime :: Int -> [Int] -> [UTCTime] -> DI.Column
  toMaybeUTCTime maxDef def times =
      let filled = stitchNullable maxDef def times
          defaultTime = UTCTime (fromGregorian 1970 1 1) (secondsToDiffTime 0)
       in if all isJust filled
              then DI.fromList (map (fromMaybe defaultTime) filled)
              else DI.fromList filled
  