{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}

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

Position adjustments — 'PosStack', 'PosFill', 'PosDodge', 'PosJitter'.
Operates on long-format frames; @stack@ and @fill@ additionally write
a @__ybase@ column for the bar bottom.
-}
module Granite.Position (
    applyPosition,
) where

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

import Granite.Data.Frame (
    Column (..),
    DataFrame (..),
    columnAsNum,
    columnAsText,
    fromColumns,
    lookupColumn,
 )
import Granite.Spec (
    ColumnRef (..),
    Mapping (..),
    Position (..),
 )

applyPosition :: Position -> Mapping -> DataFrame -> DataFrame
applyPosition :: Position -> Mapping -> DataFrame -> DataFrame
applyPosition Position
pos Mapping
m DataFrame
df = case Position
pos of
    Position
PosIdentity -> DataFrame
df
    Position
PosStack -> Bool -> Mapping -> DataFrame -> DataFrame
stackPos Bool
False Mapping
m DataFrame
df
    Position
PosFill -> Bool -> Mapping -> DataFrame -> DataFrame
stackPos Bool
True Mapping
m DataFrame
df
    PosDodge Double
gap -> Double -> Mapping -> DataFrame -> DataFrame
dodgePos Double
gap Mapping
m DataFrame
df
    PosJitter Double
dx Double
dy -> Double -> Double -> Mapping -> DataFrame -> DataFrame
jitterPos Double
dx Double
dy Mapping
m DataFrame
df

stackPos :: Bool -> Mapping -> DataFrame -> DataFrame
stackPos :: Bool -> Mapping -> DataFrame -> DataFrame
stackPos Bool
fill Mapping
m DataFrame
df =
    case ( Maybe ColumnRef -> DataFrame -> Maybe [Double]
numColAt (Mapping -> Maybe ColumnRef
aesX Mapping
m) DataFrame
df
         , Maybe ColumnRef -> DataFrame -> Maybe [Double]
numColAt (Mapping -> Maybe ColumnRef
aesY Mapping
m) DataFrame
df
         ) of
        (Just [Double]
xs, Just [Double]
ys) ->
            let gs :: [Text]
gs = Mapping -> DataFrame -> Int -> [Text]
groupKeys Mapping
m DataFrame
df ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
                rows :: [(Double, Double, Text)]
rows = [Double] -> [Double] -> [Text] -> [(Double, Double, Text)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Double]
xs [Double]
ys [Text]
gs
                uniqXs :: [Double]
uniqXs = [Double] -> [Double]
forall a. Eq a => [a] -> [a]
List.nub [Double]
xs
                stackedPerX :: [(Double, Double, Text, Double, Double)]
stackedPerX =
                    (Double -> [(Double, Double, Text, Double, Double)])
-> [Double] -> [(Double, Double, Text, Double, Double)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
                        ( \Double
x ->
                            let here :: [(Double, Double, Text)]
here = [(Double
x', Double
y', Text
g') | (Double
x', Double
y', Text
g') <- [(Double, Double, Text)]
rows, Double
x' Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
x]
                                ordered :: [(Double, Double, Text)]
ordered = ((Double, Double, Text) -> Int)
-> [(Double, Double, Text)] -> [(Double, Double, Text)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
List.sortOn (\(Double
_, Double
_, Text
g) -> [Text] -> Text -> Int
seriesIx [Text]
gs Text
g) [(Double, Double, Text)]
here
                                running :: [Double]
running = (Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 [Double
y | (Double
_, Double
y, Text
_) <- [(Double, Double, Text)]
ordered]
                                tops :: [Double]
tops = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
running
                                bots :: [Double]
bots = [Double] -> [Double]
forall a. HasCallStack => [a] -> [a]
init [Double]
running
                                segs :: [((Double, Double, Text), Double, Double)]
segs = [(Double, Double, Text)]
-> [Double]
-> [Double]
-> [((Double, Double, Text), Double, Double)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [(Double, Double, Text)]
ordered [Double]
tops [Double]
bots
                                total :: Double
total = [Double] -> Double
forall a. HasCallStack => [a] -> a
last [Double]
running
                                norm :: Double
norm = if Bool
fill Bool -> Bool -> Bool
&& Double
total Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
0 then Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
total else Double
1
                             in [(Double
xi, Double
yi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
norm, Text
gi, Double
top Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
norm, Double
bot Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
norm) | ((Double
xi, Double
yi, Text
gi), Double
top, Double
bot) <- [((Double, Double, Text), Double, Double)]
segs]
                        )
                        [Double]
uniqXs
                xs' :: [Double]
xs' = [Double
x | (Double
x, Double
_, Text
_, Double
_, Double
_) <- [(Double, Double, Text, Double, Double)]
stackedPerX]
                ys' :: [Double]
ys' = [Double
y | (Double
_, Double
_, Text
_, Double
y, Double
_) <- [(Double, Double, Text, Double, Double)]
stackedPerX]
                bases :: [Double]
bases = [Double
b | (Double
_, Double
_, Text
_, Double
_, Double
b) <- [(Double, Double, Text, Double, Double)]
stackedPerX]
                gs' :: [Text]
gs' = [Text
g | (Double
_, Double
_, Text
g, Double
_, Double
_) <- [(Double, Double, Text, Double, Double)]
stackedPerX]
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
                yName :: Text
yName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesY Mapping
m) Text
"y"
                gName :: Text
gName = case Mapping -> Maybe ColumnRef
groupRef Mapping
m of
                    Just (ColumnRef Text
n) -> Text
n
                    Maybe ColumnRef
Nothing -> Text
"_group"
                base0 :: [(Text, Column)]
base0 = [(Text
xName, [Double] -> Column
ColNum [Double]
xs'), (Text
yName, [Double] -> Column
ColNum [Double]
ys'), (Text
"__ybase", [Double] -> Column
ColNum [Double]
bases)]
                cols :: [(Text, Column)]
cols =
                    if Mapping -> Bool
isJustGroup Mapping
m
                        then [(Text, Column)]
base0 [(Text, Column)] -> [(Text, Column)] -> [(Text, Column)]
forall a. Semigroup a => a -> a -> a
<> [(Text
gName, [Text] -> Column
ColCat [Text]
gs')]
                        else [(Text, Column)]
base0
             in [(Text, Column)] -> DataFrame
fromColumns [(Text, Column)]
cols
        (Maybe [Double], Maybe [Double])
_ -> DataFrame
df

dodgePos :: Double -> Mapping -> DataFrame -> DataFrame
dodgePos :: Double -> Mapping -> DataFrame -> DataFrame
dodgePos Double
gap Mapping
m DataFrame
df =
    case Maybe ColumnRef -> DataFrame -> Maybe [Double]
numColAt (Mapping -> Maybe ColumnRef
aesX Mapping
m) DataFrame
df of
        Just [Double]
xs ->
            let gs :: [Text]
gs = Mapping -> DataFrame -> Int -> [Text]
groupKeys Mapping
m DataFrame
df ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs)
                seriesOrder :: [Text]
seriesOrder = [Text] -> [Text]
forall a. Eq a => [a] -> [a]
List.nub [Text]
gs
                nSeries :: Int
nSeries = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 ([Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
seriesOrder)
                center :: Double
center = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
nSeries Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
                shifted :: [Double]
shifted =
                    [ Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
gap Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Text] -> Text -> Int
seriesIx [Text]
gs Text
g) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
center)
                    | (Double
x, Text
g) <- [Double] -> [Text] -> [(Double, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
xs [Text]
gs
                    ]
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
             in Text -> Column -> DataFrame -> DataFrame
replaceColumn Text
xName ([Double] -> Column
ColNum [Double]
shifted) DataFrame
df
        Maybe [Double]
_ -> DataFrame
df

{- | Deterministic jitter via the golden-ratio low-discrepancy
sequence, so two runs on the same data produce the same offsets
(golden tests stay stable).
-}
jitterPos :: Double -> Double -> Mapping -> DataFrame -> DataFrame
jitterPos :: Double -> Double -> Mapping -> DataFrame -> DataFrame
jitterPos Double
dx Double
dy Mapping
m DataFrame
df =
    let phi :: Double
phi = (Double -> Double
forall a. Floating a => a -> a
sqrt Double
5 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2 :: Double
        wobble :: Double -> p -> Double
wobble Double
seed p
i =
            let r :: Double
r = (p -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral p
i Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
phi Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
seed
             in (Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
r :: Int)) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5
        update :: DataFrame -> Double -> Double -> Maybe ColumnRef -> DataFrame
update DataFrame
mc Double
seedScale Double
axis Maybe ColumnRef
col = case Maybe ColumnRef
col of
            Just (ColumnRef Text
name) -> case Text -> DataFrame -> Maybe Column
lookupColumn Text
name DataFrame
df of
                Just (ColNum [Double]
vs) ->
                    let vs' :: [Double]
vs' = [Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
axis Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Int -> Double
forall {p}. Integral p => Double -> p -> Double
wobble Double
seedScale Int
i | (Int
i, Double
v) <- [Int] -> [Double] -> [(Int, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Double]
vs]
                     in Text -> Column -> DataFrame -> DataFrame
replaceColumn Text
name ([Double] -> Column
ColNum [Double]
vs') DataFrame
mc
                Maybe Column
_ -> DataFrame
mc
            Maybe ColumnRef
Nothing -> DataFrame
mc
        df1 :: DataFrame
df1 = DataFrame -> Double -> Double -> Maybe ColumnRef -> DataFrame
update DataFrame
df Double
0.123 Double
dx (Mapping -> Maybe ColumnRef
aesX Mapping
m)
        df2 :: DataFrame
df2 = DataFrame -> Double -> Double -> Maybe ColumnRef -> DataFrame
update DataFrame
df1 Double
0.789 Double
dy (Mapping -> Maybe ColumnRef
aesY Mapping
m)
     in DataFrame
df2

{- | A 'ColCat' column projects each row to its index in the
unique-value list, so stack / dodge group rows that share a
category.
-}
numColAt :: Maybe ColumnRef -> DataFrame -> Maybe [Double]
numColAt :: Maybe ColumnRef -> DataFrame -> Maybe [Double]
numColAt Maybe ColumnRef
Nothing DataFrame
_ = Maybe [Double]
forall a. Maybe a
Nothing
numColAt (Just (ColumnRef Text
n)) DataFrame
df =
    case Text -> DataFrame -> Maybe Column
lookupColumn Text
n DataFrame
df of
        Just (ColCat [Text]
xs) ->
            let uniques :: [Text]
uniques = [Text] -> [Text]
forall a. Eq a => [a] -> [a]
List.nub [Text]
xs
                indexOf :: Text -> b
indexOf Text
x = b -> (Int -> b) -> Maybe Int -> b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe b
0 Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Text -> [Text] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
List.elemIndex Text
x [Text]
uniques)
             in [Double] -> Maybe [Double]
forall a. a -> Maybe a
Just ((Text -> Double) -> [Text] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Double
forall {b}. Num b => Text -> b
indexOf [Text]
xs)
        Just Column
c -> Column -> Maybe [Double]
columnAsNum Column
c
        Maybe Column
Nothing -> Maybe [Double]
forall a. Maybe a
Nothing

groupRef :: Mapping -> Maybe ColumnRef
groupRef :: Mapping -> Maybe ColumnRef
groupRef Mapping
m = case Mapping -> Maybe ColumnRef
aesGroup Mapping
m of
    Just ColumnRef
r -> ColumnRef -> Maybe ColumnRef
forall a. a -> Maybe a
Just ColumnRef
r
    Maybe ColumnRef
Nothing -> case Mapping -> Maybe ColumnRef
aesColor Mapping
m of
        Just ColumnRef
r -> ColumnRef -> Maybe ColumnRef
forall a. a -> Maybe a
Just ColumnRef
r
        Maybe ColumnRef
Nothing -> Mapping -> Maybe ColumnRef
aesFill Mapping
m

isJustGroup :: Mapping -> Bool
isJustGroup :: Mapping -> Bool
isJustGroup Mapping
m = case Mapping -> Maybe ColumnRef
groupRef Mapping
m of
    Just ColumnRef
_ -> Bool
True
    Maybe ColumnRef
Nothing -> Bool
False

groupKeys :: Mapping -> DataFrame -> Int -> [Text]
groupKeys :: Mapping -> DataFrame -> Int -> [Text]
groupKeys Mapping
m DataFrame
df Int
fallbackLen = case Mapping -> Maybe ColumnRef
groupRef Mapping
m of
    Just (ColumnRef Text
n) -> case Text -> DataFrame -> Maybe Column
lookupColumn Text
n DataFrame
df of
        Just Column
c -> Column -> [Text]
columnAsText Column
c
        Maybe Column
Nothing -> Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate Int
fallbackLen Text
"all"
    Maybe ColumnRef
Nothing -> Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate Int
fallbackLen Text
"all"

seriesIx :: [Text] -> Text -> Int
seriesIx :: [Text] -> Text -> Int
seriesIx [Text]
all_ Text
k =
    let order :: [Text]
order = [Text] -> [Text]
forall a. Eq a => [a] -> [a]
List.nub [Text]
all_
        go :: t -> [Text] -> t
go t
_ [] = t
0
        go t
i (Text
x : [Text]
xs) = if Text
x Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
k then t
i else t -> [Text] -> t
go (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [Text]
xs
     in Int -> [Text] -> Int
forall {t}. Num t => t -> [Text] -> t
go Int
0 [Text]
order

refName :: Maybe ColumnRef -> Text -> Text
refName :: Maybe ColumnRef -> Text -> Text
refName (Just (ColumnRef Text
n)) Text
_ = Text
n
refName Maybe ColumnRef
Nothing Text
fb = Text
fb

replaceColumn :: Text -> Column -> DataFrame -> DataFrame
replaceColumn :: Text -> Column -> DataFrame -> DataFrame
replaceColumn Text
name Column
col (DataFrame [(Text, Column)]
cols)
    | ((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)]
cols =
        [(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)]
cols]
    | Bool
otherwise = [(Text, Column)] -> DataFrame
DataFrame ([(Text, Column)]
cols [(Text, Column)] -> [(Text, Column)] -> [(Text, Column)]
forall a. Semigroup a => a -> a -> a
<> [(Text
name, Column
col)])