{-# LANGUAGE Strict #-}

{- |
Module      : Granite.Data.Frame
Copyright   : (c) 2025
License     : MIT
Maintainer  : mschavinda@gmail.com

Column-oriented data frame used as input to the chart IR.
-}
module Granite.Data.Frame (
    Column (..),
    columnLength,
    columnAsNum,
    columnAsText,
    DataFrame (..),
    emptyFrame,
    addColumn,
    lookupColumn,
    columnNames,
    frameLength,
    fromColumns,
    filterByRows,
    uniqueText,
    pickRows,
) where

import Data.Int (Int64)
import Data.List qualified as List
import Data.Text (Text)
import Data.Text qualified as Text

data Column
    = ColNum ![Double]
    | ColCat ![Text]
    | ColTime ![Int64]
    | ColBool ![Bool]
    deriving (Column -> Column -> Bool
(Column -> Column -> Bool)
-> (Column -> Column -> Bool) -> Eq Column
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Column -> Column -> Bool
== :: Column -> Column -> Bool
$c/= :: Column -> Column -> Bool
/= :: Column -> Column -> Bool
Eq, Int -> Column -> ShowS
[Column] -> ShowS
Column -> String
(Int -> Column -> ShowS)
-> (Column -> String) -> ([Column] -> ShowS) -> Show Column
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Column -> ShowS
showsPrec :: Int -> Column -> ShowS
$cshow :: Column -> String
show :: Column -> String
$cshowList :: [Column] -> ShowS
showList :: [Column] -> ShowS
Show, ReadPrec [Column]
ReadPrec Column
Int -> ReadS Column
ReadS [Column]
(Int -> ReadS Column)
-> ReadS [Column]
-> ReadPrec Column
-> ReadPrec [Column]
-> Read Column
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Column
readsPrec :: Int -> ReadS Column
$creadList :: ReadS [Column]
readList :: ReadS [Column]
$creadPrec :: ReadPrec Column
readPrec :: ReadPrec Column
$creadListPrec :: ReadPrec [Column]
readListPrec :: ReadPrec [Column]
Read)

columnLength :: Column -> Int
columnLength :: Column -> Int
columnLength Column
col = case Column
col of
    ColNum [Double]
xs -> [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
    ColCat [Text]
xs -> [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
xs
    ColTime [Int64]
xs -> [Int64] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int64]
xs
    ColBool [Bool]
xs -> [Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
xs

columnAsNum :: Column -> Maybe [Double]
columnAsNum :: Column -> Maybe [Double]
columnAsNum Column
col = case Column
col of
    ColNum [Double]
xs -> [Double] -> Maybe [Double]
forall a. a -> Maybe a
Just [Double]
xs
    ColBool [Bool]
xs -> [Double] -> Maybe [Double]
forall a. a -> Maybe a
Just ((Bool -> Double) -> [Bool] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Bool
b -> if Bool
b then Double
1 else Double
0) [Bool]
xs)
    ColTime [Int64]
xs -> [Double] -> Maybe [Double]
forall a. a -> Maybe a
Just ((Int64 -> Double) -> [Int64] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int64]
xs)
    ColCat [Text]
_ -> Maybe [Double]
forall a. Maybe a
Nothing

columnAsText :: Column -> [Text]
columnAsText :: Column -> [Text]
columnAsText Column
col = case Column
col of
    ColCat [Text]
xs -> [Text]
xs
    ColNum [Double]
xs -> (Double -> Text) -> [Double] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Text
showD [Double]
xs
    ColTime [Int64]
xs -> (Int64 -> Text) -> [Int64] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Text
Text.pack (String -> Text) -> (Int64 -> String) -> Int64 -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> String
forall a. Show a => a -> String
show) [Int64]
xs
    ColBool [Bool]
xs -> (Bool -> Text) -> [Bool] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Text
Text.pack (String -> Text) -> (Bool -> String) -> Bool -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> String
forall a. Show a => a -> String
show) [Bool]
xs
  where
    showD :: Double -> Text
    showD :: Double -> Text
showD Double
d = String -> Text
Text.pack (Double -> String
forall a. Show a => a -> String
show Double
d)

newtype DataFrame = DataFrame {DataFrame -> [(Text, Column)]
unDataFrame :: [(Text, Column)]}
    deriving (DataFrame -> DataFrame -> Bool
(DataFrame -> DataFrame -> Bool)
-> (DataFrame -> DataFrame -> Bool) -> Eq DataFrame
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DataFrame -> DataFrame -> Bool
== :: DataFrame -> DataFrame -> Bool
$c/= :: DataFrame -> DataFrame -> Bool
/= :: DataFrame -> DataFrame -> Bool
Eq, Int -> DataFrame -> ShowS
[DataFrame] -> ShowS
DataFrame -> String
(Int -> DataFrame -> ShowS)
-> (DataFrame -> String)
-> ([DataFrame] -> ShowS)
-> Show DataFrame
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DataFrame -> ShowS
showsPrec :: Int -> DataFrame -> ShowS
$cshow :: DataFrame -> String
show :: DataFrame -> String
$cshowList :: [DataFrame] -> ShowS
showList :: [DataFrame] -> ShowS
Show, ReadPrec [DataFrame]
ReadPrec DataFrame
Int -> ReadS DataFrame
ReadS [DataFrame]
(Int -> ReadS DataFrame)
-> ReadS [DataFrame]
-> ReadPrec DataFrame
-> ReadPrec [DataFrame]
-> Read DataFrame
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS DataFrame
readsPrec :: Int -> ReadS DataFrame
$creadList :: ReadS [DataFrame]
readList :: ReadS [DataFrame]
$creadPrec :: ReadPrec DataFrame
readPrec :: ReadPrec DataFrame
$creadListPrec :: ReadPrec [DataFrame]
readListPrec :: ReadPrec [DataFrame]
Read)

emptyFrame :: DataFrame
emptyFrame :: DataFrame
emptyFrame = [(Text, Column)] -> DataFrame
DataFrame []

addColumn :: Text -> Column -> DataFrame -> DataFrame
addColumn :: Text -> Column -> DataFrame -> DataFrame
addColumn Text
name Column
col (DataFrame [(Text, Column)]
xs)
    | ((Text, Column) -> Bool) -> [(Text, Column)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
name) (Text -> Bool)
-> ((Text, Column) -> Text) -> (Text, Column) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Column) -> Text
forall a b. (a, b) -> a
fst) [(Text, Column)]
xs =
        [(Text, Column)] -> DataFrame
DataFrame [(Text
n, if Text
n Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
name then Column
col else Column
c) | (Text
n, Column
c) <- [(Text, Column)]
xs]
    | Bool
otherwise = [(Text, Column)] -> DataFrame
DataFrame ([(Text, Column)]
xs [(Text, Column)] -> [(Text, Column)] -> [(Text, Column)]
forall a. [a] -> [a] -> [a]
++ [(Text
name, Column
col)])

lookupColumn :: Text -> DataFrame -> Maybe Column
lookupColumn :: Text -> DataFrame -> Maybe Column
lookupColumn Text
name (DataFrame [(Text, Column)]
xs) = Text -> [(Text, Column)] -> Maybe Column
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
name [(Text, Column)]
xs

columnNames :: DataFrame -> [Text]
columnNames :: DataFrame -> [Text]
columnNames (DataFrame [(Text, Column)]
xs) = ((Text, Column) -> Text) -> [(Text, Column)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Column) -> Text
forall a b. (a, b) -> a
fst [(Text, Column)]
xs

frameLength :: DataFrame -> Int
frameLength :: DataFrame -> Int
frameLength (DataFrame [(Text, Column)]
xs) = ((Text, Column) -> Int -> Int) -> Int -> [(Text, Column)] -> Int
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int -> Int -> Int)
-> ((Text, Column) -> Int) -> (Text, Column) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Column -> Int
columnLength (Column -> Int)
-> ((Text, Column) -> Column) -> (Text, Column) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Column) -> Column
forall a b. (a, b) -> b
snd) Int
0 [(Text, Column)]
xs

fromColumns :: [(Text, Column)] -> DataFrame
fromColumns :: [(Text, Column)] -> DataFrame
fromColumns = [(Text, Column)] -> DataFrame
DataFrame

pickRows :: [Int] -> Column -> Column
pickRows :: [Int] -> Column -> Column
pickRows [Int]
ixs Column
col = case Column
col of
    ColNum [Double]
xs -> [Double] -> Column
ColNum ([Double] -> [Double]
forall a. [a] -> [a]
pick [Double]
xs)
    ColCat [Text]
xs -> [Text] -> Column
ColCat ([Text] -> [Text]
forall a. [a] -> [a]
pick [Text]
xs)
    ColTime [Int64]
xs -> [Int64] -> Column
ColTime ([Int64] -> [Int64]
forall a. [a] -> [a]
pick [Int64]
xs)
    ColBool [Bool]
xs -> [Bool] -> Column
ColBool ([Bool] -> [Bool]
forall a. [a] -> [a]
pick [Bool]
xs)
  where
    pick :: [a] -> [a]
    pick :: forall a. [a] -> [a]
pick [a]
xs = [[a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
i <- [Int]
ixs, Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs]

filterByRows :: [Int] -> DataFrame -> DataFrame
filterByRows :: [Int] -> DataFrame -> DataFrame
filterByRows [Int]
ixs (DataFrame [(Text, Column)]
cols) =
    [(Text, Column)] -> DataFrame
DataFrame [(Text
n, [Int] -> Column -> Column
pickRows [Int]
ixs Column
c) | (Text
n, Column
c) <- [(Text, Column)]
cols]

uniqueText :: [Text] -> [Text]
uniqueText :: [Text] -> [Text]
uniqueText = [Text] -> [Text]
forall a. Eq a => [a] -> [a]
List.nub