granite-0.4.0.0: src/Granite.hs

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

{- |
Module      : Granite
Copyright   : (c) 2025
License     : MIT
Maintainer  : mschavinda@gmail.com
Stability   : experimental
Portability : POSIX

A terminal-based plotting library that renders beautiful charts using Unicode
Braille characters and ANSI colors. Granite provides a variety of chart types
including scatter plots, line graphs, bar charts, pie charts, histograms,
heatmaps, and box plots.

= Basic Usage

Create a simple scatter plot:

@
{\-# LANGUAGE OverloadedStrings #-\}
import Granite
import Data.Text.IO as T

main = do
  let points = [(x, sin x) | x <- [0, 0.1 .. 6.28]]
      chart = scatter [series "sin(x)" points] defPlot
  T.putStrLn chart
@

= Customization

Plots can be customized using record update syntax:

@
let customPlot = defPlot
      { widthChars = 80
      , heightChars = 30
      , plotTitle = "My Chart"
      , legendPos = LegendBottom
      }
@

= Terminal Requirements

This library requires a terminal that supports:

  * Unicode (specifically Braille patterns U+2800-U+28FF)
  * ANSI color codes
  * Monospace font with proper Braille character rendering
-}
module Granite (
    -- * Plot Configuration
    Plot (..),
    defPlot,
    LegendPos (..),

    -- * Formatting
    Color (..),
    LabelFormatter,
    AxisEnv (..),

    -- * Data Preparation
    series,
    bins,
    Bins (..),

    -- * Chart Types
    histogram,
    bars,
    scatter,
    pie,
    stackedBars,
    heatmap,
    lineGraph,
    boxPlot,
) where

import Data.Bits (xor, (.&.), (.|.))
import Data.Char (chr)
import Data.Function (on)
import Data.List qualified as List
import Data.Maybe
import Data.Text (Text)
import Data.Text qualified as Text
import Numeric (showEFloat, showFFloat)
import Text.Printf

-- | Position of the legend in the plot.
data LegendPos
    = -- | Display legend on the right side of the plot
      LegendRight
    | -- | Display legend below the plot
      LegendBottom
    | -- | Do not display legend.
      LegendNone
    deriving (Eq, Show)

{- | Plot configuration parameters.

Controls the appearance and layout of generated charts.
-}
data Plot = Plot
    { widthChars :: Int
    -- ^ Width of the plot area in terminal characters (default: 60)
    , heightChars :: Int
    -- ^ Height of the plot area in terminal characters (default: 20)
    , leftMargin :: Int
    -- ^ Space reserved for y-axis labels (default: 6)
    , bottomMargin :: Int
    -- ^ Space reserved for x-axis labels (default: 2)
    , titleMargin :: Int
    -- ^ Space above the plot for the title (default: 1)
    , xBounds :: (Maybe Double, Maybe Double)
    {- ^ Optional manual x-axis bounds (min, max).
    'Nothing' uses automatic bounds with 5% padding.
    -}
    , yBounds :: (Maybe Double, Maybe Double)
    {- ^ Optional manual y-axis bounds (min, max).
    'Nothing' uses automatic bounds with 5% padding.
    -}
    , plotTitle :: Text
    -- ^ Title displayed above the plot (default: empty)
    , legendPos :: LegendPos
    -- ^ Position of the legend (default: 'LegendRight')
    , colorPalette :: [Color]
    -- ^ Color palette that'll be used by the plot.
    , xFormatter :: LabelFormatter
    -- ^ Formatter for x-axis labels.
    , yFormatter :: LabelFormatter
    -- ^ Formatter for y-axis labels.
    , xNumTicks :: Int
    -- ^ Number of ticks on the x axis.
    , yNumTicks :: Int
    -- ^ Number of ticks on the y axis.
    }

{- | Default plot configuration.

Creates a 60×20 character plot with reasonable defaults:

@
defPlot = Plot
  { widthChars   = 60
  , heightChars  = 20
  , leftMargin   = 6
  , bottomMargin = 2
  , titleMargin  = 1
  , xBounds      = (Nothing, Nothing)
  , yBounds      = (Nothing, Nothing)
  , plotTitle    = ""
  , legendPos    = LegendRight
  , colorPalette = [ BrightBlue, BrightMagenta, BrightCyan, BrightGreen, BrightYellow, BrightRed, BrightWhite, BrightBlack]
  , xFormatter   = \\ _ _ v -> show v
  , yFormatter   = \\ _ _ v -> show v
  , xNumTicks    = 2
  , yNumTicks    = 2
  }
@
-}
defPlot :: Plot
defPlot =
    Plot
        { widthChars = 60
        , heightChars = 20
        , leftMargin = 6
        , bottomMargin = 2
        , titleMargin = 1
        , xBounds = (Nothing, Nothing)
        , yBounds = (Nothing, Nothing)
        , plotTitle = ""
        , legendPos = LegendRight
        , colorPalette = paletteColors
        , xFormatter = fmt
        , yFormatter = fmt
        , xNumTicks = 3
        , yNumTicks = 3
        }

{- | Axis-aware, width-limited, tick-label formatter.

Given:

   * axis context
   * a per-tick width budget (in terminal cells)
   * and the raw tick value.

returns the label to render.
-}
type LabelFormatter =
    -- | Axis context (domain, tick index/count, etc)
    AxisEnv ->
    -- | Slot width budget in characters for this tick.
    Int ->
    -- | Raw data value for the tick
    Double ->
    -- | Rendered label (if it doesn't fit in the slot it will be truncated)
    Text.Text

-- | What the formatter gets to know about the axis/ticks
data AxisEnv = AxisEnv
    { domain :: (Double, Double)
    -- ^ min/max of the axis in data space
    , tickIndex :: Int
    -- ^ index of THIS tick [0..tickCount-1]
    , tickCount :: Int
    -- ^ total number of ticks
    }

-- | Supported ANSI colo(u)rs.
data Color
    = Default
    | Black
    | Red
    | Green
    | Yellow
    | Blue
    | Magenta
    | Cyan
    | White
    | BrightBlack
    | BrightRed
    | BrightGreen
    | BrightYellow
    | BrightBlue
    | BrightMagenta
    | BrightCyan
    | BrightWhite
    deriving (Eq, Show)

{- | Create a named data series for multi-series plots.

@
let s1 = series "Dataset A" [(1,2), (2,4), (3,6)]
    s2 = series "Dataset B" [(1,3), (2,5), (3,7)]
    chart = scatter [s1, s2] defPlot
@
-}
series ::
    -- | Name of the series (appears in legend)
    Text ->
    -- | List of (x, y) data points
    [(Double, Double)] ->
    (Text, [(Double, Double)])
series = (,)

{- | Create a scatter plot from multiple data series.

Each series is rendered with a different color and pattern.
Points are plotted using Braille characters for sub-character resolution.

==== __Example__

@
let points1 = [(x, x^2) | x <- [-3, -2.5 .. 3]]
    points2 = [(x, 2*x + 1) | x <- [-3, -2.5 .. 3]]
    chart = scatter [series "y = x²" points1,
                     series "y = 2x + 1" points2] defPlot
@
-}
scatter ::
    -- | List of named data series
    [(Text, [(Double, Double)])] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
scatter sers cfg =
    let wC = widthChars cfg
        hC = heightChars cfg
        plotC = newCanvas wC hC
        (xmin, xmax, ymin, ymax) = boundsXY cfg (concatMap snd sers)
        sx x =
            clamp 0 (wC * 2 - 1) $
                round ((x - xmin) / (xmax - xmin + eps) * fromIntegral (wC * 2 - 1))
        sy y =
            clamp 0 (hC * 4 - 1) $
                round ((ymax - y) / (ymax - ymin + eps) * fromIntegral (hC * 4 - 1))
        pats = cycle palette
        cols = cycle (colorPalette cfg)
        withSty = zipWith3 (\(n, ps) p c -> (n, ps, p, c)) sers pats cols
        drawOne c0 (_name, pts, pat, col) =
            List.foldl'
                ( \c (x, y) ->
                    let xd = sx x; yd = sy y
                     in if ink pat xd yd then setDotC c xd yd (Just col) else c
                )
                c0
                pts
        cDone = List.foldl' drawOne plotC withSty
        ax = axisify cfg cDone (xmin, xmax) (ymin, ymax)
        legend =
            legendBlock
                (legendPos cfg)
                (leftMargin cfg + widthChars cfg)
                [(n, p, col) | (n, _, p, col) <- withSty]
     in drawFrame cfg ax legend

{- | Create a line graph connecting data points.

Similar to 'scatter' but connects consecutive points with lines.
Points are automatically sorted by x-coordinate before connecting.

==== __Example__

@
let sine = [(x, sin x) | x <- [0, 0.1 .. 2*pi]]
    cosine = [(x, cos x) | x <- [0, 0.1 .. 2*pi]]
    chart = lineGraph [series "sin" sine, series "cos" cosine] defPlot
@
-}
lineGraph ::
    -- | List of named data series
    [(Text, [(Double, Double)])] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
lineGraph sers cfg =
    let wC = widthChars cfg
        hC = heightChars cfg
        plotC = newCanvas wC hC
        (xmin, xmax, ymin, ymax) = boundsXY cfg (concatMap snd sers)
        sx x =
            clamp 0 (wC * 2 - 1) $
                round ((x - xmin) / (xmax - xmin + eps) * fromIntegral (wC * 2 - 1))
        sy y =
            clamp 0 (hC * 4 - 1) $
                round ((ymax - y) / (ymax - ymin + eps) * fromIntegral (hC * 4 - 1))

        cols = cycle (colorPalette cfg)
        withSty = zip sers cols

        drawSeries ((_name, pts), col) c0 =
            let sortedPts = List.sortOn fst pts
                dotPairs = zip sortedPts (drop 1 sortedPts)
             in List.foldl'
                    ( \c ((x1, y1), (x2, y2)) ->
                        lineDotsC (sx x1, sy y1) (sx x2, sy y2) (Just col) c
                    )
                    c0
                    dotPairs

        cDone = List.foldl' (flip drawSeries) plotC withSty
        ax :: Text
        ax = axisify cfg cDone (xmin, xmax) (ymin, ymax)
        legend :: Text
        legend =
            legendBlock
                (legendPos cfg)
                (leftMargin cfg + widthChars cfg)
                [(n, Solid, col) | ((n, _), col) <- withSty]
     in drawFrame cfg ax legend

{- | Create a bar chart from categorical data.

Each bar is colored differently and labeled with its category name.

==== __Example__

@
let fruits = [(\"Apple\", 45.2), (\"Banana\", 38.1), (\"Orange\", 52.7)]
    chart = bars fruits defPlot { plotTitle = "Fruit Sales" }
@
-}
bars ::
    -- | List of (category, value) pairs
    [(Text, Double)] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
bars kvs cfg =
    let wC = widthChars cfg
        hC = heightChars cfg
        (catNames, vals) = unzip kvs
        vmax = maximum' (map abs vals)

        cats :: [(Text, Double, Color)]
        cats =
            [ (name, abs v / vmax, col)
            | ((name, v), col) <- zip kvs (cycle (colorPalette cfg))
            ]

        nCats = min wC (length cats)

        (base, extra) =
            if nCats == 0 then (0, 0) else (wC `div` nCats, wC - wC `div` nCats * nCats)
        widths = [base + (if i < extra then 1 else 0) | i <- [0 .. nCats - 1]]

        catGroups :: [[(String, Maybe Color)]]
        catGroups =
            [ replicate w (colGlyphs hC f, Just col)
            | ((_, f, col), w) <- zip cats widths
            ]

        gutterCol = (replicate hC ' ', Nothing)
        columns = List.intercalate [gutterCol] catGroups

        grid :: [[(Char, Maybe Color)]]
        grid =
            [ [(glyphs !! y, mc) | (glyphs, mc) <- columns]
            | y <- [0 .. hC - 1]
            ]

        ax =
            axisifyGrid
                cfg
                grid
                (0, fromIntegral (max 1 nCats))
                (0, vmax)
                catNames
                (fmap (+ 1) (listToMaybe widths))
        legendWidth = leftMargin cfg + 1 + gridWidth grid
        legend =
            legendBlock
                (legendPos cfg)
                legendWidth
                [(name, Checker, col) | (name, _, col) <- cats]
     in drawFrame cfg ax legend

{- | Create a stacked bar chart.

Each category can have multiple stacked components.

==== __Example__

@
let sales = [(\"Q1\", [(\"Product A\", 100), (\"Product B\", 150)]),
             (\"Q2\", [(\"Product A\", 120), (\"Product B\", 180)])]
    chart = stackedBars sales defPlot
@
-}
stackedBars ::
    -- | Categories with stacked components
    [(Text, [(Text, Double)])] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
stackedBars categories cfg =
    let wC = widthChars cfg
        hC = heightChars cfg

        seriesNames = case categories of
            [] -> []
            (c : _) -> map fst (snd c)

        totals = [sum (map snd series') | (_, series') <- categories]
        maxHeight = maximum (1e-12 : totals)

        nCats = length categories
        (base, extra) =
            if nCats == 0
                then (0, 0)
                else (wC `div` nCats, wC - wC `div` nCats * nCats)
        widths = [base + (if i < extra then 1 else 0) | i <- [0 .. nCats - 1]]

        cols = cycle (colorPalette cfg)
        seriesColors = zip seriesNames cols

        makeBar (_, series') width =
            let cumHeights = scanl (+) 0 [v / maxHeight | (_, v) <- series']
                segments = zip3 (map fst series') cumHeights (drop 1 cumHeights)

                makeColumn :: [(Char, Maybe Color)]
                makeColumn =
                    [ let heightFromBottom = fromIntegral (hC - y) / fromIntegral hC
                          findSegment [] = (' ', Nothing)
                          findSegment ((name, bottom, top) : rest) =
                            if heightFromBottom > bottom && heightFromBottom <= top
                                then ('█', lookup name seriesColors)
                                else findSegment rest
                       in findSegment segments
                    | y <- [0 .. hC - 1]
                    ]
             in replicate width makeColumn

        gutterCol = replicate hC (' ', Nothing)
        allBars = zipWith makeBar categories widths
        columns = List.intercalate [gutterCol] allBars

        grid = [[col !! y | col <- columns] | y <- [0 .. hC - 1]]

        ax :: Text
        ax =
            axisifyGrid
                cfg
                grid
                (0, fromIntegral (max 1 nCats))
                (0, maxHeight)
                (map fst categories)
                (fmap (+ 1) (listToMaybe widths))
        legend :: Text
        legend =
            legendBlock
                (legendPos cfg)
                ( leftMargin cfg
                    + 1
                    + gridWidth grid
                )
                [(name, Solid, col) | (name, col) <- seriesColors]
     in drawFrame cfg ax legend

-- | Defines the binning parameters.
data Bins = Bins
    { nBins :: Int
    , lo :: Double
    , hi :: Double
    }
    deriving (Eq, Show)

{- | Create a bin configuration for histograms.

@
bins 10 0 100  -- 10 bins from 0 to 100
bins 20 (-5) 5 -- 20 bins from -5 to 5
@
-}
bins :: Int -> Double -> Double -> Bins
bins n a b = Bins (max 1 n) (min a b) (max a b)

{- | Create a histogram from numerical data.

Data is binned according to the provided 'Bins' configuration.

==== __Example__

@
import System.Random

-- Generate random normal-like distribution
let values = take 1000 $ randomRs (0, 100) gen
    chart = histogram (bins 20 0 100) values defPlot
@
-}
histogram ::
    -- | Binning configuration
    Bins ->
    -- | Raw data values to bin
    [Double] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
histogram (Bins n a b) xs cfg =
    let step = (b - a) / fromIntegral n
        binIx x = clamp 0 (n - 1) $ floor ((x - a) / step)
        counts =
            List.foldl'
                ( \acc x ->
                    if x < a || x > b
                        then acc
                        else addAt acc (binIx x) 1
                )
                (replicate n 0 :: [Int])
                xs
        maxC = fromIntegral (maximum (1 : counts))
        fracs0 = [fromIntegral c / maxC | c <- counts]

        wData = widthChars cfg
        hC = heightChars cfg
        colsF = resampleToWidth wData fracs0

        dataCols = [(colGlyphs hC f, Just BrightCyan) | f <- colsF]
        grid :: [[(Char, Maybe Color)]]
        grid =
            [ [(fst col !! y, snd col) | col <- dataCols]
            | y <- [0 .. hC - 1]
            ]

        ax =
            axisifyGrid cfg grid (a, b) (0, maxC) [] Nothing
        legendWidth = leftMargin cfg + 1 + gridWidth grid
        legend = legendBlock (legendPos cfg) legendWidth [("count", Solid, BrightCyan)]
     in drawFrame cfg ax legend

{- | Create a pie chart showing proportions.

Values are normalized to sum to 100%. Negative values are treated as zero.

==== __Example__

@
let browsers = [(\"Chrome\", 65), (\"Firefox\", 20), (\"Safari\", 10), (\"Other\", 5)]
    chart = pie browsers defPlot { plotTitle = "Browser Market Share" }
@
-}
pie ::
    -- | List of (category, value) pairs
    [(Text, Double)] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
pie parts0 cfg =
    let parts = normalize parts0
        wC = widthChars cfg
        hC = heightChars cfg
        plotC = newCanvas wC hC
        wDots = wC * 2
        hDots = hC * 4
        r = min (wDots `div` 2 - 2) (hDots `div` 2 - 2)
        cx = wDots `div` 2
        cy = hDots `div` 2
        toAng p = p * 2 * pi
        wedges = scanl (\a (_, p) -> a + toAng p) 0 parts
        angles = zip wedges (drop 1 wedges)
        names = map fst parts
        cols = cycle pieColors
        withP :: [(Text, (Double, Double), Color)]
        withP = zip3 names angles cols

        drawOne (_name, (a0, a1), col) c0 =
            let inside x y =
                    let dx = fromIntegral (x - cx)
                        dy = fromIntegral (cy - y)
                        rr2 = dx * dx + dy * dy
                        r2 = fromIntegral (r * r)
                        ang = atan2 dy dx `mod'` (2 * pi)
                     in rr2 <= r2 && angleWithin ang a0 a1
             in fillDotsC (cx - r, cy - r) (cx + r, cy + r) inside (Just col) c0

        cDone = List.foldl' (flip drawOne) plotC withP
        ax = axisify cfg cDone (0, 1) (0, 1)
        legend =
            legendBlock
                (legendPos cfg)
                (leftMargin cfg + widthChars cfg)
                [(n, Solid, col) | (n, _, col) <- withP]
     in drawFrame cfg ax legend

{- | Create a heatmap visualization of a 2D matrix.

Values are mapped to a color gradient from blue (low) to red (high).

==== __Example__

@
let matrix = [[x * y | x <- [1..10]] | y <- [1..10]]
    chart = heatmap matrix defPlot { plotTitle = "Multiplication Table" }
@
-}
heatmap ::
    -- | 2D matrix of values (rows × columns)
    [[Double]] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
heatmap matrix cfg =
    let rows = length matrix
        cols = gridWidth matrix

        allVals = concat matrix
        vmin = if null allVals then 0 else minimum' allVals
        vmax = if null allVals then 1 else maximum' allVals
        vrange = vmax - vmin

        intensityColors =
            [ Blue
            , BrightBlue
            , Cyan
            , BrightCyan
            , Green
            , BrightGreen
            , Yellow
            , BrightYellow
            , Magenta
            , BrightRed
            , Red
            ]

        colorForValue v =
            if vrange < eps
                then Green
                else
                    let norm = clamp 0 1 ((v - vmin) / vrange)
                        idx = floor (norm * fromIntegral (length intensityColors - 1))
                        idx' = clamp 0 (length intensityColors - 1) idx
                     in intensityColors !! idx'

        displayGrid =
            [ [ let
                    matrixRow = min (rows - 1) ((plotH - 1 - i) * rows `div` plotH)
                    matrixCol = min (cols - 1) (j * cols `div` plotW)
                    val = matrix !! matrixRow !! matrixCol
                 in
                    ('█', Just (colorForValue val))
              | j <- [0 .. plotW - 1]
              ]
            | i <- [0 .. plotH - 1]
            ]

        plotW = widthChars cfg
        plotH = heightChars cfg

        colLabels = [Text.pack (show i) | i <- [0 .. cols - 1]]
        rowLabels = [Text.pack (show i) | i <- reverse [0 .. rows - 1]]

        cellHeight = fromIntegral plotH / fromIntegral rows
        yTicks =
            [ ( round @Double @Int (fromIntegral i * cellHeight + cellHeight / 2)
              , rowLabels !! i
              )
            | i <- [0 .. rows - 1]
            ]

        left = leftMargin cfg
        baseLbl = replicate plotH (Text.replicate left " ")
        yLabels =
            List.foldl'
                ( \acc (row, lbl) ->
                    setAt acc row (justifyRight left (ellipsisize left lbl))
                )
                baseLbl
                yTicks

        renderRow cells =
            Text.concat
                (fmap (\(ch, mc) -> maybe (Text.singleton ch) (`paint` ch) mc) cells)
        attachY = zipWith (\lbl cells -> lbl <> "│" <> renderRow cells) yLabels displayGrid

        xBar = Text.replicate left " " <> "└" <> Text.replicate plotW "─"
        xLine =
            placeGridLabels
                (Text.replicate (left + 1) " ")
                (plotW `div` cols)
                colLabels

        ax = Text.unlines (attachY <> [xBar, xLine])

        gradientLegend =
            Text.pack (printf "%.2f " vmin)
                <> Text.concat (fmap (`paint` '█') intensityColors)
                <> Text.pack (printf " %.2f" vmax)
     in drawFrame cfg ax gradientLegend

{- | Create a box plot showing statistical distributions.

Displays quartiles, median, and min/max values for each dataset.

==== __Example__

@
let data1 = [1.2, 2.3, 2.1, 3.4, 2.8, 4.1, 3.9]
    data2 = [5.1, 4.8, 6.2, 5.9, 7.1, 6.5, 5.5]
    chart = boxPlot [("Group A", data1), ("Group B", data2)] defPlot
@

The box plot displays:

  * Box: First quartile (Q1) to third quartile (Q3)
  * Line inside box: Median (Q2)
  * Whiskers: Minimum and maximum values
-}
boxPlot ::
    -- | Named datasets
    [(Text, [Double])] ->
    -- | Plot configuration
    Plot ->
    -- | Rendered chart as Text
    Text
boxPlot datasets cfg =
    let wC = widthChars cfg
        hC = heightChars cfg

        stats = [(name, quartiles vals) | (name, vals) <- datasets]

        allVals = concatMap snd datasets
        ymin = if null allVals then 0 else minimum' allVals - abs (minimum' allVals) * 0.1
        ymax = if null allVals then 1 else maximum' allVals + abs (maximum' allVals) * 0.1

        nBoxes = length datasets
        boxWidth = if nBoxes == 0 then 1 else max 1 (wC `div` (nBoxes * 2))
        spacing = if nBoxes <= 1 then 0 else (wC - boxWidth * nBoxes) `div` (nBoxes - 1)

        scaleY v =
            clamp 0 (hC - 1) $
                round ((ymax - v) / (ymax - ymin + eps) * fromIntegral (hC - 1))

        emptyGrid = replicate hC (replicate wC (' ', Nothing))

        drawBox grid (idx, (_name, (minV, q1, median, q3, maxV))) =
            let xStart = idx * (boxWidth + spacing)
                xMid = xStart + boxWidth `div` 2
                xEnd = xStart + boxWidth - 1

                minRow = scaleY minV
                q1Row = scaleY q1
                medRow = scaleY median
                q3Row = scaleY q3
                maxRow = scaleY maxV

                col = pieColors !! (idx `mod` length pieColors)

                grid1 = drawVLine grid xMid minRow q1Row '│' (Just col)
                grid2 = drawVLine grid1 xMid q3Row maxRow '│' (Just col)

                grid3 = drawHLine grid2 xStart xEnd q1Row '─' (Just col)
                grid4 = drawHLine grid3 xStart xEnd q3Row '─' (Just col)
                grid5 = drawVLine grid4 xStart q1Row q3Row '│' (Just col)
                grid6 = drawVLine grid5 xEnd q1Row q3Row '│' (Just col)

                grid7 = drawHLine grid6 xStart xEnd medRow '═' (Just col)

                grid8 = setGridChar grid7 xMid minRow '┴' (Just col)
                grid9 = setGridChar grid8 xMid maxRow '┬' (Just col)
             in grid9

        finalGrid = List.foldl' drawBox emptyGrid (zip [0 ..] stats)

        left = leftMargin cfg
        baseLbl = replicate hC (Text.replicate left " ")

        yTicks = ticks1D hC (yNumTicks cfg) (ymin, ymax) True
        yEnv n = AxisEnv (ymin, ymax) n (yNumTicks cfg)
        ySlot = max 1 left
        yLabels =
            List.foldl'
                ( \acc (row, v) ->
                    setAt acc row . ellipsisize left . justifyRight left $
                        yFormatter cfg (yEnv row) ySlot v
                )
                baseLbl
                yTicks

        renderRow cells =
            Text.concat
                (fmap (\(ch, mc) -> maybe (Text.singleton ch) (`paint` ch) mc) cells)
        attachY = zipWith (\lbl cells -> lbl <> "│" <> renderRow cells) yLabels finalGrid

        xBar = Text.replicate left " " <> "└" <> Text.replicate wC "─"

        xLine =
            List.foldl'
                ( \acc (idx, name) ->
                    let boxCenter = idx * (boxWidth + spacing) + boxWidth `div` 2
                        lblWidth = wcswidth name
                        lblStart = left + 1 + boxCenter - lblWidth `div` 2
                     in Text.take lblStart acc <> name <> Text.drop (lblStart + wcswidth name) acc
                )
                (Text.replicate (left + 1 + wC) " ")
                (zip [0 ..] (map fst datasets))

        ax = Text.unlines (attachY <> [xBar, xLine])

        legend =
            legendBlock
                (legendPos cfg)
                (leftMargin cfg + widthChars cfg)
                [ (name, Solid, pieColors !! (i `mod` length pieColors))
                | (i, (name, _)) <- zip [0 ..] stats
                ]
     in drawFrame cfg ax legend
  where
    drawVLine grid x y1 y2 ch col =
        let yStart = min y1 y2
            yEnd = max y1 y2
         in List.foldl' (\g y -> setGridChar g x y ch col) grid [yStart .. yEnd]

    drawHLine grid x1 x2 y ch col =
        let xStart = min x1 x2
            xEnd = max x1 x2
         in List.foldl' (\g x -> setGridChar g x y ch col) grid [xStart .. xEnd]

    setGridChar grid x y ch col =
        updateAt grid y (\row -> setAt row x (ch, col))

ansiCode :: Color -> Int
ansiCode Black = 30
ansiCode Red = 31
ansiCode Green = 32
ansiCode Yellow = 33
ansiCode Blue = 34
ansiCode Magenta = 35
ansiCode Cyan = 36
ansiCode White = 37
ansiCode BrightBlack = 90
ansiCode BrightRed = 91
ansiCode BrightGreen = 92
ansiCode BrightYellow = 93
ansiCode BrightBlue = 94
ansiCode BrightMagenta = 95
ansiCode BrightCyan = 96
ansiCode BrightWhite = 97
ansiCode Default = 39

ansiOn :: Color -> Text
ansiOn c = "\ESC[" <> Text.pack (show (ansiCode c)) <> "m"

ansiOff :: Text
ansiOff = "\ESC[0m"

paint :: Color -> Char -> Text
paint c ch = if ch == ' ' then " " else ansiOn c <> Text.singleton ch <> ansiOff

paletteColors :: [Color]
paletteColors =
    [ BrightBlue
    , BrightMagenta
    , BrightCyan
    , BrightGreen
    , BrightYellow
    , BrightRed
    , BrightWhite
    , BrightBlack
    ]

pieColors :: [Color]
pieColors =
    [ BrightRed
    , BrightGreen
    , BrightYellow
    , BrightBlue
    , BrightMagenta
    , BrightCyan
    , BrightWhite
    , BrightBlack
    ]

data Pat = Solid | Checker | DiagA | DiagB | Sparse deriving (Eq, Show)

ink :: Pat -> Int -> Int -> Bool
ink Solid _ _ = True
ink Checker x y = (x `xor` y) .&. 1 == 0
ink DiagA x y = (x + y) `mod` 3 /= 1
ink DiagB x y = (x - y) `mod` 3 /= 1
ink Sparse x y = x .&. 1 == 0 && y `mod` 3 == 0

palette :: [Pat]
palette = [Solid, Checker, DiagA, DiagB, Sparse]

data Array2D a = A2D Int Int (Arr a)

getA2D :: Array2D a -> Int -> Int -> a
getA2D (A2D w _ xs) x y = indexA xs (y * w + x)

setA2D :: Array2D a -> Int -> Int -> a -> Array2D a
setA2D (A2D w h xs) x y v =
    let i = y * w + x
     in A2D w h (setA xs i v)

newA2D :: Int -> Int -> a -> Array2D a
newA2D w h v = A2D w h (fromList (replicate (w * h) v))

toBit :: Int -> Int -> Int
toBit ry rx = case (ry, rx) of
    (0, 0) -> 1
    (1, 0) -> 2
    (2, 0) -> 4
    (3, 0) -> 64
    (0, 1) -> 8
    (1, 1) -> 16
    (2, 1) -> 32
    (3, 1) -> 128
    _ -> 0

data Canvas = Canvas
    { cW :: Int
    , cH :: Int
    , buffer :: Array2D Int
    , cbuf :: Array2D (Maybe Color)
    }

newCanvas :: Int -> Int -> Canvas
newCanvas w h = Canvas w h (newA2D w h 0) (newA2D w h Nothing)

setDotC :: Canvas -> Int -> Int -> Maybe Color -> Canvas
setDotC c xDot yDot mcol
    | xDot < 0 || yDot < 0 || xDot >= cW c * 2 || yDot >= cH c * 4 = c
    | otherwise =
        let (cx, rx) = xDot `divMod` 2
            (cy, ry) = yDot `divMod` 4
            b = toBit ry rx
            m = getA2D (buffer c) cx cy
            c' = c{buffer = setA2D (buffer c) cx cy (m .|. b)}
         in case mcol of
                Nothing -> c'
                Just col -> c'{cbuf = setA2D (cbuf c) cx cy (Just col)}

fillDotsC ::
    (Int, Int) ->
    (Int, Int) ->
    (Int -> Int -> Bool) ->
    Maybe Color ->
    Canvas ->
    Canvas
fillDotsC (x0, y0) (x1, y1) p mcol c0 =
    let xs = [max 0 x0 .. min (cW c0 * 2 - 1) x1]
        ys = [max 0 y0 .. min (cH c0 * 4 - 1) y1]
     in List.foldl'
            (\c y -> List.foldl' (\c' x -> if p x y then setDotC c' x y mcol else c') c xs)
            c0
            ys

renderCanvas :: Canvas -> Text
renderCanvas (Canvas w h a colA) =
    let glyph 0 = ' '
        glyph m = chr (0x2800 + m)
        rows =
            fmap
                ( \y -> flip fmap [0 .. w - 1] $ \x ->
                    let m = getA2D a x y
                        ch = glyph m
                        mc = getA2D colA x y
                     in maybe (Text.singleton ch) (`paint` ch) mc
                )
                [0 .. h - 1]
     in Text.unlines (fmap Text.concat rows)

justifyRight :: Int -> Text -> Text
justifyRight n s = Text.replicate (max 0 (n - wcswidth s)) " " <> s

wcswidth :: Text -> Int
wcswidth = go 0
  where
    go acc xs
        | Text.null xs = acc
        | Text.isPrefixOf "\ESC[" xs =
            let
                rest' = Text.dropWhile (/= 'm') xs
             in
                if Text.null rest' then acc else go acc (Text.tail rest')
        | otherwise = go (acc + 1) (Text.tail xs)

fmt :: AxisEnv -> Int -> Double -> Text
fmt _ _ v
    | abs v >= 10000 || abs v < 0.01 && v /= 0 =
        Text.pack (showEFloat (Just 1) v "")
    | otherwise = Text.pack (showFFloat (Just 1) v "")

drawFrame :: Plot -> Text -> Text -> Text
drawFrame cfg contentWithAxes legendBlockStr =
    Text.unlines $
        filter
            (not . Text.null)
            [plotTitle cfg, contentWithAxes, legendBlockStr]

{- | Evenly spaced tick positions in screen space paired with data values.
  If invertY = True, 0 maps to ymax (top row) and 1 maps to ymin (bottom).
-}
ticks1D ::
    -- | screen length in chars (width for X, height for Y)
    Int ->
    -- | requested number of ticks (will clamp to >= 2)
    Int ->
    -- | (vmin, vmax) in data space
    (Double, Double) ->
    -- | invertY? (True for Y axis so row 0 = ymax)
    Bool ->
    -- | (screenPos, dataValue)
    [(Int, Double)]
ticks1D screenLen want (vmin, vmax) invertY =
    let n = max 2 want
        lastIx = max 0 (screenLen - 1)
        toVal t =
            if invertY
                then vmax - t * (vmax - vmin)
                else vmin + t * (vmax - vmin)
        mk' k =
            let t = if n == 1 then 0 else fromIntegral k / fromIntegral (n - 1)
                pos = round (t * fromIntegral lastIx)
             in (pos, toVal t)
        raw = [mk' k | k <- [0 .. n - 1]]
        dedup = List.nubBy ((==) `on` fst) raw
     in dedup

slotBudget :: Int -> Int -> Int
slotBudget plotPixels numTicks =
    max 1 (plotPixels `div` max 1 numTicks)

axisify :: Plot -> Canvas -> (Double, Double) -> (Double, Double) -> Text
axisify cfg c (xmin, xmax) (ymin, ymax) =
    let plotW = cW c
        plotH = cH c
        left = leftMargin cfg
        pad = Text.replicate left " "

        yTicks :: [(Int, Double)]
        yTicks = ticks1D plotH (yNumTicks cfg) (ymin, ymax) True

        baseLbl :: [Text]
        baseLbl = replicate plotH pad

        yEnv n = AxisEnv (ymin, ymax) n 3
        ySlot = max 1 left
        yLabels =
            List.foldl'
                ( \acc (row, v) ->
                    setAt acc row
                        . ellipsisize left
                        . justifyRight left
                        $ yFormatter cfg (yEnv row) ySlot v
                )
                baseLbl
                yTicks

        canvasLines = Text.lines (renderCanvas c)
        attachY = zipWith (\lbl line -> lbl <> "│" <> line) yLabels canvasLines

        xBar = pad <> "└" <> Text.replicate plotW "─"

        xTicks :: [(Int, Double)]
        xTicks = ticks1D plotW (xNumTicks cfg) (xmin, xmax) False

        xEnv n = AxisEnv (xmin, xmax) n 3
        slotW = slotBudget plotW (max 1 (length xTicks))
        xLine =
            placeLabels
                (Text.replicate (left + 1 + plotW) " ")
                (left + 1)
                [(x, xFormatter cfg (xEnv i) slotW v) | (i, (x, v)) <- zip [0 ..] xTicks]
     in Text.unlines (attachY <> [xBar, xLine])

axisifyGrid ::
    Plot ->
    [[(Char, Maybe Color)]] ->
    (Double, Double) ->
    (Double, Double) ->
    [Text] ->
    Maybe Int ->
    Text
axisifyGrid cfg grid (xmin, xmax) (ymin, ymax) categories w =
    let plotH = length grid
        plotW = gridWidth grid
        left = leftMargin cfg
        pad = Text.replicate left " "

        yTicks = ticks1D plotH (yNumTicks cfg) (ymin, ymax) True
        baseLbl = List.replicate plotH pad

        yEnv n = AxisEnv (ymin, ymax) n (yNumTicks cfg)
        ySlot = max 1 left
        yLabels =
            List.foldl'
                ( \acc (row, v) ->
                    setAt acc row
                        . ellipsisize left
                        . justifyRight left
                        $ yFormatter cfg (yEnv row) ySlot v
                )
                baseLbl
                yTicks

        renderRow :: [(Char, Maybe Color)] -> Text
        renderRow cells =
            Text.concat
                (fmap (\(ch, mc) -> maybe (Text.singleton ch) (`paint` ch) mc) cells)

        attachY = zipWith (\lbl cells -> lbl <> "│" <> renderRow cells) yLabels grid

        xBar = pad <> "└" <> Text.replicate plotW "─"

        hasCategories = not (null categories) && not (all Text.null categories)

        xLine =
            if hasCategories
                then
                    let slotW =
                            fromMaybe
                                ( slotBudget
                                    plotW
                                    (max 1 (xNumTicks cfg))
                                )
                                w
                        nSlots = plotW `div` slotW
                        xTicks = ticks1D plotW nSlots (xmin, xmax) False
                     in placeGridLabels
                            (Text.replicate (left + 1) " ")
                            slotW
                            (keepPercentiles (length categories) (length xTicks + 1) categories)
                else
                    let xTicks = ticks1D plotW (xNumTicks cfg) (xmin, xmax) False
                        xEnv i = AxisEnv (xmin, xmax) i (xNumTicks cfg)
                        slotW = slotBudget plotW (max 1 (length xTicks))
                     in placeLabels
                            (Text.replicate (left + 1 + plotW) " ")
                            (left + 1)
                            [(x, xFormatter cfg (xEnv i) slotW v) | (i, (x, v)) <- zip [0 ..] xTicks]
     in Text.unlines (attachY <> [xBar, xLine])

keepPercentiles :: Int -> Int -> [Text] -> [Text]
keepPercentiles n k xs
    | k <= 0 = []
    | null xs = replicate k ""
    | n <= 1 = replicate k ""
    | otherwise = map valueAt [0 .. k - 2] ++ [last xs]
  where
    m = length xs
    pairs :: [(Int, Text)]
    pairs =
        [ ( slotIx
          , xs !! srcIx
          )
        | i <- [0 .. n - 2]
        , let srcIx = (i * (m - 1)) `div` (n - 1)
        , let slotIx = (i * (k - 1)) `div` (n - 1)
        ]

    valueAt :: Int -> Text
    valueAt i = fromMaybe "" $ List.lookup i pairs

placeLabels :: Text -> Int -> [(Int, Text)] -> Text
placeLabels base off = List.foldl' place base
  where
    place :: Text -> (Int, Text) -> Text
    place acc (x, s) =
        let i = off + x
         in Text.take i acc <> s <> Text.drop (i + wcswidth s) acc

placeGridLabels :: Text -> Int -> [Text] -> Text
placeGridLabels base slotW = List.foldl' place base
  where
    place :: Text -> Text -> Text
    place acc s =
        let lblWidth = wcswidth s
            padding = max 0 ((slotW - lblWidth) `div` 2)
            centered = Text.replicate padding " " <> s
         in acc <> Text.take slotW (centered <> Text.replicate slotW " ")

legendBlock :: LegendPos -> Int -> [(Text, Pat, Color)] -> Text
legendBlock LegendBottom width entries =
    let cells = [sample pat col <> " " <> name | (name, pat, col) <- entries]
        line = Text.intercalate "   " cells
        pad =
            let vis = wcswidth line
             in if vis < width then Text.replicate ((width - vis) `div` 2) " " else ""
     in pad <> line
legendBlock LegendRight _ entries =
    Text.unlines $
        fmap (\(name, pat, col) -> sample pat col <> " " <> name) entries
legendBlock LegendNone _ _ = ""

sample :: Pat -> Color -> Text
sample p col =
    let c =
            List.foldl'
                ( \cv (dx, dy) -> if ink p dx dy then setDotC cv (dx `mod` 2) (dy `mod` 4) (Just col) else cv
                )
                (newCanvas 1 1)
                [(x, y) | y <- [0 .. 3], x <- [0 .. 1]]
        s = renderCanvas c
     in Text.dropWhileEnd (== '\n') s

clamp :: (Ord a) => a -> a -> a -> a
clamp low high x = max low (min high x)

eps :: Double
eps = 1e-12

boundsXY :: Plot -> [(Double, Double)] -> (Double, Double, Double, Double)
boundsXY cfg pts =
    let (xs, ys) = unzip pts
        xmin = minimum' xs
        xmax = maximum' xs
        ymin = minimum' ys
        ymax = maximum' ys
        padx = (xmax - xmin) * 0.05 + 1e-9
        pady = (ymax - ymin) * 0.05 + 1e-9
     in ( fromMaybe (xmin - padx) (fst (xBounds cfg))
        , fromMaybe (xmax + padx) (snd (xBounds cfg))
        , fromMaybe (ymin - pady) (fst (yBounds cfg))
        , fromMaybe (ymax + pady) (snd (yBounds cfg))
        )

mod' :: Double -> Double -> Double
mod' a m = a - fromIntegral (floor (a / m) :: Int) * m

blockChar :: Int -> Char
blockChar n = case clamp 0 8 n of
    0 -> ' '
    1 -> '▁'
    2 -> '▂'
    3 -> '▃'
    4 -> '▄'
    5 -> '▅'
    6 -> '▆'
    7 -> '▇'
    _ -> '█'

colGlyphs :: Int -> Double -> String
colGlyphs hC frac =
    let total = hC * 8
        ticks = clamp 0 total (round (frac * fromIntegral total))
        full = ticks `div` 8
        rem8 = ticks - full * 8
        topPad = hC - full - (if rem8 > 0 then 1 else 0)
        middle = [blockChar rem8 | rem8 > 0]
     in replicate topPad ' ' <> middle <> replicate full '█'

resampleToWidth :: Int -> [Double] -> [Double]
resampleToWidth w xs
    | w <= 0 = []
    | null xs = replicate w 0
    | n == w = xs
    | n > w = avgGroup (ceiling (fromIntegral n / (fromIntegral w :: Double)))
    | otherwise = replicateOut
  where
    n = length xs
    avgGroup g =
        [avg (take g (drop (i * g) xs)) | i <- [0 .. w - 1]]
      where
        avg ys = if null ys then 0 else sum ys / fromIntegral (length ys)
    replicateOut =
        let base = w `div` n
            extra = w - base * n
         in concat
                [ replicate (base + (if i < extra then 1 else 0)) v
                | (i, v) <- zip [0 ..] xs
                ]

addAt :: [Int] -> Int -> Int -> [Int]
addAt xs i v = updateAt xs i (+ v)

normalize :: [(Text, Double)] -> [(Text, Double)]
normalize xs =
    let s = sum (map (abs . snd) xs) + 1e-12
     in [(n, max 0 (v / s)) | (n, v) <- xs]

angleWithin :: Double -> Double -> Double -> Bool
angleWithin ang a0 a1
    | a1 >= a0 = ang >= a0 && ang <= a1
    | otherwise = ang >= a0 || ang <= a1

lineDotsC :: (Int, Int) -> (Int, Int) -> Maybe Color -> Canvas -> Canvas
lineDotsC (x0, y0) (x1, y1) mcol c0 =
    let dx = abs (x1 - x0)
        sx = if x0 < x1 then 1 else -1
        dy = negate (abs (y1 - y0))
        sy = if y0 < y1 then 1 else -1
        go x y err c
            | x == x1 && y == y1 = setDotC c x y mcol
            | otherwise =
                let e2 = 2 * err
                    (x', err') = if e2 >= dy then (x + sx, err + dy) else (x, err)
                    (y', err'') = if e2 <= dx then (y + sy, err' + dx) else (y, err')
                 in go x' y' err'' (setDotC c x y mcol)
     in go x0 y0 (dx + dy) c0

quartiles :: [Double] -> (Double, Double, Double, Double, Double)
quartiles [] = (0, 0, 0, 0, 0) -- Idk. Maybe throw an error here???
quartiles xs =
    let sorted = List.sort xs
        n = length sorted
        q1Idx = n `div` 4
        q2Idx = n `div` 2
        q3Idx = 3 * n `div` 4
        getIdx i = if i < n then sorted !! i else last sorted
     in if n < 5
            then let m = sum xs / fromIntegral n in (m, m, m, m, m)
            else
                ( fromMaybe 0 (listToMaybe sorted)
                , getIdx q1Idx
                , getIdx q2Idx
                , getIdx q3Idx
                , last sorted
                )

gridWidth :: [[a]] -> Int
gridWidth [] = 0
gridWidth (x : _) = length x

-- | Min and max function for axis bounds which defaults to 0 and 1 when empty.
minimum', maximum' :: [Double] -> Double
minimum' [] = 0
minimum' xs = minimum xs
maximum' [] = 1
maximum' xs = maximum xs

-- AVL Tree we'll use as an array.
-- This improves upon the previous implementation that relies
-- on linked list for indexing and update (both O(n)) while keeping
-- the dependencies very light (wouldn't want to install all of containers
-- just to get an int map).
data Arr a
    = E
    | N Int Int (Arr a) a (Arr a)

size :: Arr a -> Int
size E = 0
size (N sz _ _ _ _) = sz

height :: Arr a -> Int
height E = 0
height (N _ h _ _ _) = h

mk :: Arr a -> a -> Arr a -> Arr a
mk l x r = N sz h l x r
  where
    sl = size l
    sr = size r
    hl = height l
    hr = height r
    sz = 1 + sl + sr
    h = 1 + max hl hr

rotateL :: Arr a -> Arr a
rotateL (N _ _ l x (N _ _ rl y rr)) = mk (mk l x rl) y rr
rotateL _ = error "rotateL: malformed tree"

rotateR :: Arr a -> Arr a
rotateR (N _ _ (N _ _ ll y lr) x r) = mk ll y (mk lr x r)
rotateR _ = error "rotateR: malformed tree"

balance :: Arr a -> Arr a
balance t@(N _ _ l x r)
    | height l > height r + 1 =
        case l of
            N _ _ ll _ lr ->
                if height ll >= height lr
                    then rotateR t
                    else rotateR (mk (rotateL l) x r)
            _ -> t
    | height r > height l + 1 =
        case r of
            N _ _ rl _ rr ->
                if height rr >= height rl
                    then rotateL t
                    else rotateL (mk l x (rotateR r))
            _ -> t
    | otherwise = mk l x r
balance t = t

indexA :: Arr a -> Int -> a
indexA t i =
    case t of
        E -> error ("index out of bounds: " <> show i)
        N _ _ l x r ->
            let sl = size l
             in if i < 0 || i >= 1 + sl + size r
                    then error ("index out of bounds: " <> show i)
                    else
                        if i < sl
                            then indexA l i
                            else
                                if i == sl
                                    then x
                                    else indexA r (i - sl - 1)

setA :: Arr a -> Int -> a -> Arr a
setA t i y =
    case t of
        E -> error ("index out of bounds when setting: " <> show i)
        N _ _ l x r ->
            let sl = size l
             in if i < 0 || i >= 1 + sl + size r
                    then error ("index out of bounds: " <> show i)
                    else
                        if i < sl
                            then balance (mk (setA l i y) x r)
                            else
                                if i == sl
                                    then mk l y r
                                    else balance (mk l x (setA r (i - sl - 1) y))

fromList :: [a] -> Arr a
fromList xs = fst (build (length xs) xs)
  where
    build :: Int -> [a] -> (Arr a, [a])
    build 0 ys = (E, ys)
    build n ys =
        let (l, ys1) = build (n `div` 2) ys
            (x, ys2) = case ys1 of
                [] -> error "IMPOSSIBLE"
                (v : vs) -> (v, vs)
            (r, ys3) = build (n - n `div` 2 - 1) ys2
         in (mk l x r, ys3)

-- >>> setAt "abc" (-1) 'x'
-- "abc"
-- >>> setAt "abc" 0 'x'
-- "xbc"
-- >>> setAt "abc" 2 'x'
-- "abx"
-- >>> setAt "abc" 10 'x'
-- "abc"
setAt :: [a] -> Int -> a -> [a]
setAt xs i v = updateAt xs i (const v)

updateAt :: [a] -> Int -> (a -> a) -> [a]
updateAt xs i f
    | i < 0 = xs
    | otherwise = go xs i
  where
    go [] _ = []
    go (x : rest) 0 = f x : rest
    go (x : rest) n = x : go rest (n - 1)

{- | Ensure the text fits within maxWidth. If it doesn't, truncate and append an ellipsis.
>>> ellipsisize 5 "Hello, World!"
"Hell\8230"
>>> ellipsisize 1 "Hi"
"\8230"
>>> ellipsisize 0 "Hello, World!"
""
>>> ellipsisize 20 "Hello, World!"
"Hello, World!"
-}
ellipsisize :: Int -> Text -> Text
ellipsisize maxWidth lbl
    | maxWidth <= 0 = ""
    | wcswidth lbl > maxWidth = Text.take (maxWidth - 1) lbl <> "…"
    | otherwise = lbl