content: {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE ExplicitNamespaces #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE InstanceSigs #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  
  module DataFrame.Internal.Row where
  
  import qualified Data.List as L
  import qualified Data.Map as M
  import qualified Data.Set as S
  import qualified Data.Text as T
  import qualified Data.Vector as V
  import qualified Data.Vector.Algorithms.Merge as VA
  import qualified Data.Vector.Generic as VG
  import qualified Data.Vector.Unboxed as VU
  
  import Control.Exception (throw)
  import Control.Monad.ST (runST)
  import Data.Function (on)
  import Data.Maybe (fromMaybe)
  import Data.Type.Equality (TestEquality (..))
  import Data.Typeable (type (:~:) (..))
  import DataFrame.Errors (DataFrameException (..))
  import DataFrame.Internal.Column
  import DataFrame.Internal.DataFrame
  import Text.ParserCombinators.ReadPrec (ReadPrec)
  import Text.Read (
      Lexeme (Ident),
      lexP,
      parens,
      readListPrec,
      readListPrecDefault,
      readPrec,
   )
  import Type.Reflection (typeOf, typeRep)
  
  data Any where
      Value :: (Columnable a) => a -> Any
  
  instance Eq Any where
      (==) :: Any -> Any -> Bool
      (Value a) == (Value b) = fromMaybe False $ do
          Refl <- testEquality (typeOf a) (typeOf b)
          return $ a == b
  
  instance Ord Any where
      (<=) :: Any -> Any -> Bool
      (Value a) <= (Value b) = fromMaybe False $ do
          Refl <- testEquality (typeOf a) (typeOf b)
          return $ a <= b
  
  instance Show Any where
      show :: Any -> String
      show (Value a) = T.unpack (showValue a)
  
  showValue :: forall a. (Columnable a) => a -> T.Text
  showValue v = case testEquality (typeRep @a) (typeRep @T.Text) of
      Just Refl -> v
      Nothing -> case testEquality (typeRep @a) (typeRep @String) of
          Just Refl -> T.pack v
          Nothing -> (T.pack . show) v
  
  instance Read Any where
      readListPrec :: ReadPrec [Any]
      readListPrec = readListPrecDefault
  
      readPrec :: ReadPrec Any
      readPrec = parens $ do
          Ident "Value" <- lexP
          readPrec
  
  -- | Wraps a value into an \Any\ type. This helps up represent rows as heterogenous lists.
  toAny :: forall a. (Columnable a) => a -> Any
  toAny = Value
  
  -- | Unwraps a value from an \Any\ type.
  fromAny :: forall a. (Columnable a) => Any -> Maybe a
  fromAny (Value (v :: b)) = do
      Refl <- testEquality (typeRep @a) (typeRep @b)
      pure v
  
  type Row = V.Vector Any
  
  (!?) :: [a] -> Int -> Maybe a
  (!?) [] _ = Nothing
  (!?) (x : _) 0 = Just x
  (!?) (x : xs) n = (!?) xs (n - 1)
  
  mkColumnFromRow :: Int -> [[Any]] -> Column
  mkColumnFromRow i rows = case rows of
      [] -> fromList ([] :: [T.Text])
      (row : _) -> case row !? i of
          Nothing -> fromList ([] :: [T.Text])
          Just (Value (v :: a)) -> fromList $ reverse $ L.foldl' addToList [v] (drop 1 rows)
            where
              addToList acc r = case r !? i of
                  Nothing -> acc
                  Just (Value (v' :: b)) -> case testEquality (typeRep @a) (typeRep @b) of
                      Nothing -> acc
                      Just Refl -> v' : acc
  
  {- | Converts the entire dataframe to a list of rows.
  
  Each row contains all columns in the dataframe, ordered by their column indices.
  The rows are returned in their natural order (from index 0 to n-1).
  
  ==== __Examples__
  
  >>> toRowList df
  [Row {name = "Alice", age = 25, ...}, Row {name = "Bob", age = 30, ...}, ...]
  
  ==== __Performance note__
  
  This function materializes all rows into a list, which may be memory-intensive
  for large dataframes. Consider using 'toRowVector' if you need random access
  or streaming operations.
  -}
  toRowList :: DataFrame -> [Row]
  toRowList df =
      let
          nameSet =
              S.fromList (map fst (L.sortBy (compare `on` snd) $ M.toList (columnIndices df)))
       in
          map (mkRowRep df nameSet) [0 .. (fst (dataframeDimensions df) - 1)]
  
  {- | Converts the dataframe to a vector of rows with only the specified columns.
  
  Each row will contain only the columns named in the @names@ parameter.
  This is useful when you only need a subset of columns or want to control
  the column order in the resulting rows.
  
  ==== __Parameters__
  
  [@names@] List of column names to include in each row. The order of names
            determines the order of fields in the resulting rows.
  
  [@df@] The dataframe to convert.
  
  ==== __Examples__
  
  >>> toRowVector ["name", "age"] df
  Vector of rows with only name and age fields
  
  >>> toRowVector [] df  -- Empty column list
  Vector of empty rows (one per dataframe row)
  -}
  toRowVector :: [T.Text] -> DataFrame -> V.Vector Row
  toRowVector names df =
      let
          nameSet = S.fromList names
       in
          V.generate (fst (dataframeDimensions df)) (mkRowRep df nameSet)
  
  mkRowFromArgs :: [T.Text] -> DataFrame -> Int -> Row
  mkRowFromArgs names df i = V.map get (V.fromList names)
    where
      get name = case getColumn name df of
          Nothing ->
              throw $
                  ColumnNotFoundException
                      name
                      "[INTERNAL] mkRowFromArgs"
                      (M.keys $ columnIndices df)
          Just (BoxedColumn column) -> toAny (column V.! i)
          Just (UnboxedColumn column) -> toAny (column VU.! i)
          Just (OptionalColumn column) -> toAny (column V.! i)
  
  mkRowRep :: DataFrame -> S.Set T.Text -> Int -> Row
  mkRowRep df names i = V.generate (S.size names) (\index -> get (names' V.! index))
    where
      inOrderIndexes = map fst $ L.sortBy (compare `on` snd) $ M.toList (columnIndices df)
      names' = V.fromList [n | n <- inOrderIndexes, S.member n names]
      throwError name =
          error $
              "Column "
                  ++ T.unpack name
                  ++ " has less items than "
                  ++ "the other columns at index "
                  ++ show i
      get name = case getColumn name df of
          Just (BoxedColumn c) -> case c V.!? i of
              Just e -> toAny e
              Nothing -> throwError name
          Just (OptionalColumn c) -> case c V.!? i of
              Just e -> toAny e
              Nothing -> throwError name
          Just (UnboxedColumn c) -> case c VU.!? i of
              Just e -> toAny e
              Nothing -> throwError name
          Nothing ->
              throw $ ColumnNotFoundException name "mkRowRep" (M.keys $ columnIndices df)
  
  sortedIndexes' :: Bool -> V.Vector Row -> VU.Vector Int
  sortedIndexes' asc rows = runST $ do
      withIndexes <- VG.thaw (V.indexed rows)
      VA.sortBy ((if asc then compare else flip compare) `on` snd) withIndexes
      sorted <- VG.unsafeFreeze withIndexes
      return $ VU.generate (VG.length rows) (\i -> fst (sorted VG.! i))
  