nmr-schedule 0.2.1

//! Plotting and visualization in the terminal. Requires the `terminal-viz` feature to be enabled.
//!
//! Colors are chosen in the [OKLCH](https://evilmartians.com/chronicles/oklch-in-css-why-quit-rgb-hsl) color space, ensuring that lightness and hue are orthogonal.

use core::fmt::Display;
use core::ops::Range;

use alloc::vec;

use owo_colors::AnsiColors;
use owo_colors::DynColors;
use owo_colors::OwoColorize;
use owo_colors::Rgb;

use alloc::string::String;
use alloc::vec::Vec;
use owo_colors::Style;
use rustfft::num_complex::Complex;
use rustfft::num_complex::ComplexFloat;

/// Create a plottable frequency list for each number in the array
#[allow(clippy::missing_panics_doc)]
pub fn histogram(items: &[usize]) -> Vec<f64> {
    if items.is_empty() {
        return Vec::new();
    }

    // The array was verified not to be empty
    let max = *items.iter().max().unwrap();

    let mut hist = vec![0.; max + 1];

    for item in items {
        hist[*item] += 1.;
    }

    hist
}

const GAMUT_CLIPPED: (u8, u8, u8) = (255, 0, 0);

// https://bottosson.github.io/posts/oklab/

fn oklch_to_rgb(l: f64, c: f64, h: f64) -> Option<(u8, u8, u8)> {
    let (sin, cos) = h.to_radians().sin_cos();

    oklab_to_rgb(l, c * cos, c * sin)
}

fn linear_to_u8(v: f64) -> Option<u8> {
    // Linear RGB -> RGB
    let v = if v < 0.0031308 {
        v * 12.92
    } else {
        v.powf(2.4.recip()) * 1.055 - 0.055
    };

    let value = (v * 255.).round();

    if !(0. ..=255.).contains(&value) {
        return None;
    }

    Some(value as u8)
}

fn oklab_to_rgb(l: f64, a: f64, b: f64) -> Option<(u8, u8, u8)> {
    let l_ = l + 0.3963377774 * a + 0.2158037573 * b;
    let m_ = l - 0.1055613458 * a - 0.0638541728 * b;
    let s_ = l - 0.0894841775 * a - 1.2914855480 * b;

    let l = l_ * l_ * l_;
    let m = m_ * m_ * m_;
    let s = s_ * s_ * s_;

    let r = 4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s;
    let g = -1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s;
    let b = -0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s;

    Some((linear_to_u8(r)?, linear_to_u8(g)?, linear_to_u8(b)?))
}

/// Convert a complex number in polar coordinates into a color
///
/// The lightness represents the magnitude and the hue represents the phase angle
pub fn polar_to_color((r, θ): (f64, f64), max: f64) -> Rgb {
    let deg = θ * 180. * core::f64::consts::FRAC_1_PI;

    let color = if r > max {
        GAMUT_CLIPPED
    } else {
        oklch_to_rgb(0.76 * r / max, 0.12 * r / max, 142.5 + deg).unwrap_or(GAMUT_CLIPPED)
    };

    Rgb(color.0, color.1, color.2)
}

/// Convert a complex number to a color
///
/// See [`polar_to_color`]
pub fn complex_to_color(complex: Complex<f64>, max: f64) -> Rgb {
    let (r, θ) = complex.to_polar();

    polar_to_color((r, θ), max)
}

/// Options for how to make bar charts
#[must_use]
#[derive(Clone, Copy, Debug)]
pub struct BarChartOptions {
    max: Option<f64>,
    lines: Option<usize>,
    left_pad: usize,
    bg_color: DynColors,
}

impl BarChartOptions {
    /// Create a new `PlotOptions` with default parameters.
    pub const fn new() -> Self {
        BarChartOptions {
            max: None,
            lines: None,
            left_pad: 0,
            bg_color: DynColors::Ansi(AnsiColors::Black),
        }
    }

    /// Set the maximum Y value of the plot. Values above it will be clipped.
    pub const fn with_max(mut self, max: f64) -> Self {
        self.max = Some(max);
        self
    }

    /// Set the number of lines the plot will print out.
    pub const fn with_height(mut self, lines: usize) -> Self {
        self.lines = Some(lines);
        self
    }

    /// Set the number of characters of left-padding
    pub const fn with_left_pad(mut self, left_pad: usize) -> Self {
        self.left_pad = left_pad;
        self
    }

    /// Set the background color of the plot
    pub const fn with_bg_color(mut self, bg_color: DynColors) -> Self {
        self.bg_color = bg_color;
        self
    }
}

impl Default for BarChartOptions {
    fn default() -> Self {
        Self::new()
    }
}

/// Options for how to make row plots
#[must_use]
#[derive(Clone, Copy, Debug)]
pub struct RowChartOptions {
    max: Option<f64>,
    show_max: bool,
    left_pad: usize,
}

impl RowChartOptions {
    /// Create a new `RowPlotOptions` with default values.
    pub const fn new() -> Self {
        RowChartOptions {
            max: None,
            show_max: false,
            left_pad: 0,
        }
    }

    /// Set the maximum value.
    ///
    /// Colors of values that exceed the maximum will be rendered as pure red.
    pub const fn with_max(mut self, max: f64) -> Self {
        self.max = Some(max);
        self
    }

    /// Print the maximum value of the plot at the end of the row.
    pub const fn show_max(mut self) -> Self {
        self.show_max = true;
        self
    }

    /// Set the number of characters of left-padding.
    pub const fn with_left_pad(mut self, left_pad: usize) -> Self {
        self.left_pad = left_pad;
        self
    }
}

impl Default for RowChartOptions {
    fn default() -> Self {
        Self::new()
    }
}

fn lerp(v: f64, from: Range<f64>, to: Range<f64>) -> f64 {
    (v - from.start) / (from.end - from.start) * (to.end - to.start) + to.start
}

fn do_pad(f: &mut core::fmt::Formatter<'_>, padding: usize) -> core::fmt::Result {
    for _ in 0..padding {
        write!(f, " ")?;
    }

    Ok(())
}

fn character_at(
    (r, θ): (f64, f64),
    height_at: f64,
    row_height: f64,
    upside_down: bool,
    bg_color: DynColors,
) -> impl Display {
    const BLOCKS: [char; 9] = [' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█'];

    let char_idx = (((r - height_at) / row_height) * 8.).round().clamp(0., 8.) as usize;

    let char = &BLOCKS[if upside_down { 8 - char_idx } else { char_idx }];

    let mut style = Style::new()
        .color(polar_to_color((1., θ), 1.))
        .on_color(bg_color);

    if upside_down {
        style = style.reversed();
    }

    char.style(style)
}

/// A list of complex numbers that can be plotted in the terminal
pub trait Plot {
    /// Displays a bar chart when printed
    type BarChart<'a>: Display
    where
        Self: 'a;
    /// Displays a row chart when printed
    type RowChart<'a>: Display
    where
        Self: 'a;

    /// Plot the list as a bar chart.
    ///
    /// The height of each bar represents the magnitude of the number and the color represents the phase angle. The scale is printed above the graph.
    #[must_use = "Creating a bar chart has no effect without printing it."]
    fn bar_chart(&self, options: BarChartOptions) -> Self::BarChart<'_>;

    /// Plot the list in a single row.
    ///
    /// The magnitude is represented by lightness and the phase angle is represented by hue.
    #[must_use = "Creating a row chart has no effect without printing it."]
    fn row_chart(&self, options: RowChartOptions) -> Self::RowChart<'_>;
}

impl Plot for [Complex<f64>] {
    type BarChart<'a>
        = ComplexBarChart<'a>
    where
        Self: 'a;

    type RowChart<'a>
        = ComplexRowChart<'a>
    where
        Self: 'a;

    fn bar_chart(&self, options: BarChartOptions) -> Self::BarChart<'_> {
        ComplexBarChart(self, options)
    }

    fn row_chart(&self, options: RowChartOptions) -> Self::RowChart<'_> {
        ComplexRowChart(self, options)
    }
}

impl Plot for [f64] {
    type BarChart<'a>
        = RealBarChart<'a>
    where
        Self: 'a;

    type RowChart<'a>
        = RealRowChart<'a>
    where
        Self: 'a;

    fn bar_chart(&self, options: BarChartOptions) -> Self::BarChart<'_> {
        RealBarChart(self, options)
    }

    fn row_chart(&self, options: RowChartOptions) -> Self::RowChart<'_> {
        RealRowChart(self, options)
    }
}

#[doc(hidden)]
#[derive(Debug)]
pub struct ComplexBarChart<'a>(&'a [Complex<f64>], BarChartOptions);

impl Display for ComplexBarChart<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let ComplexBarChart(values, options) = self;

        if values.is_empty() {
            writeln!(f, "No values to plot")?;
            return Ok(());
        }

        let values = values.iter().map(|v| v.to_polar()).collect::<Vec<_>>();

        let max = options.max.unwrap_or_else(|| {
            values
                .iter()
                .max_by(|a, b| a.0.total_cmp(&b.0))
                .expect("the list is not empty to have been checked previously")
                .0
        });

        let rows = options.lines.unwrap_or_else(|| (values.len() / 4).max(10));

        do_pad(f, options.left_pad)?;
        writeln!(
            f,
            "Y-Axis: 0 - {max} |  1:{}  i:{} -1:{} -i:{}",
            '▩'.color(complex_to_color(Complex::new(1., 0.), 1.)),
            '▩'.color(complex_to_color(Complex::new(0., 1.), 1.)),
            '▩'.color(complex_to_color(Complex::new(-1., 0.), 1.)),
            '▩'.color(complex_to_color(Complex::new(0., -1.), 1.)),
        )?;

        let rows_f = rows as f64;
        let row_height = max / rows_f;

        for row in 0..rows {
            do_pad(f, options.left_pad)?;
            write!(f, "|")?;
            for value in &values {
                write!(
                    f,
                    "{}",
                    character_at(
                        *value,
                        lerp(row as f64, (0.)..(rows_f - 1.), (max - row_height)..0.),
                        row_height,
                        false,
                        options.bg_color,
                    )
                )?;
            }
            writeln!(f,)?;
        }

        do_pad(f, options.left_pad)?;
        writeln!(
            f,
            "{}",
            (0..values.len() + 1).map(|_| '‾').collect::<String>()
        )?;

        Ok(())
    }
}

#[doc(hidden)]
#[derive(Debug)]
pub struct ComplexRowChart<'a>(&'a [Complex<f64>], RowChartOptions);

impl Display for ComplexRowChart<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        plot_complex_row(f, self.0.iter().copied(), self.1)
    }
}

#[doc(hidden)]
#[derive(Debug)]
pub struct RealBarChart<'a>(&'a [f64], BarChartOptions);

impl Display for RealBarChart<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let RealBarChart(values, options) = self;

        if values.is_empty() {
            writeln!(f, "No values to plot")?;
            return Ok(());
        }

        let max_pos = options.max.unwrap_or_else(|| {
            values
                .iter()
                .max_by(|a, b| a.total_cmp(b))
                .expect("the list is not empty to have been checked previously")
                .max(0.)
        });

        let max_neg = options
            .max
            .map(|v| -v)
            .unwrap_or_else(|| values.iter().min_by(|a, b| a.total_cmp(b)).unwrap().min(0.));

        let rows = options.lines.unwrap_or_else(|| (values.len() / 4).max(10));

        do_pad(f, options.left_pad)?;
        writeln!(f, "Y-Axis: {max_neg:.2} - {max_pos:.2}",)?;

        let rows_f = rows as f64;
        let row_height = (max_pos - max_neg) / rows_f;

        for row in 0..rows {
            do_pad(f, options.left_pad)?;
            write!(f, "|")?;
            for value in *values {
                let row_pos = lerp(
                    row as f64,
                    (0.)..(rows_f - 1.),
                    (max_pos - row_height)..(max_neg + row_height),
                );

                write!(
                    f,
                    "{}",
                    character_at(
                        (
                            (value * row_pos.signum()).max(0.),
                            if *value > 0. {
                                0.
                            } else {
                                core::f64::consts::PI
                            }
                        ),
                        row_pos.abs(),
                        row_height,
                        row_pos.is_sign_negative(),
                        options.bg_color,
                    )
                )?;
            }
            writeln!(f,)?;
        }

        do_pad(f, options.left_pad)?;
        writeln!(
            f,
            "{}",
            (0..values.len() + 1).map(|_| '‾').collect::<String>()
        )?;

        Ok(())
    }
}

#[doc(hidden)]
#[derive(Debug)]
pub struct RealRowChart<'a>(&'a [f64], RowChartOptions);

impl Display for RealRowChart<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        plot_complex_row(f, self.0.iter().map(|v| Complex::new(*v, 0.)), self.1)
    }
}

fn plot_complex_row(
    f: &mut core::fmt::Formatter<'_>,
    values: impl Iterator<Item = Complex<f64>> + Clone,
    options: RowChartOptions,
) -> core::fmt::Result {
    do_pad(f, options.left_pad)?;

    let max = match options.max.or_else(|| {
        values
            .clone()
            .map(|v| v.abs())
            .max_by(|a, b| a.total_cmp(b))
    }) {
        Some(v) => v,
        None => return Ok(()), // No values to plot
    };

    for v in values {
        write!(f, "{}", "█".color(complex_to_color(v, max)))?;
    }

    if options.show_max {
        write!(f, " 0.00-{:.2}", max)?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use owo_colors::Rgb;
    use rustfft::num_complex::Complex;

    use crate::terminal_viz::{complex_to_color, oklch_to_rgb};

    use super::histogram;

    #[test]
    fn test_histogram() {
        let data = [1, 2, 3, 1, 2, 8, 6, 4, 1];

        let hist = histogram(&data);

        assert_eq!(&*hist, &[0., 3., 2., 1., 1., 0., 1., 0., 1.]);
    }

    #[test]
    fn test_oklch() {
        // https://oklch.com
        assert_eq!(oklch_to_rgb(0.7, 0.1, 72.), Some((197, 148, 85)));
        assert_eq!(oklch_to_rgb(0.4185, 0.1698, 303.97), Some((98, 40, 149)));
        assert_eq!(oklch_to_rgb(0.873, 0.0967, 158.66), Some((157, 233, 190)));
        assert_eq!(oklch_to_rgb(0.873, 0.0967, 158.66), Some((157, 233, 190)));
        assert_eq!(oklch_to_rgb(0.4876, 0.0428, 122.55), Some((91, 100, 73)));

        assert_eq!(oklch_to_rgb(0.1739, 0.1, 72.), None);
        assert_eq!(oklch_to_rgb(0.8168, 0.1004, 276.4), None);
        assert_eq!(oklch_to_rgb(0.5494, 0.1104, 199.03), None);
        assert_eq!(oklch_to_rgb(0.3246, 0.1104, 124.33), None);
        assert_eq!(oklch_to_rgb(0.2011, 0.0729, 241.72), None);
    }

    #[test]
    fn test_complex_to_color() {
        assert_eq!(
            complex_to_color(Complex::new(1., 0.), 1.),
            Rgb(131, 197, 125)
        );

        assert_eq!(
            complex_to_color(Complex::new(1., 0.0001), 1.),
            Rgb(255, 0, 0)
        );

        assert_eq!(complex_to_color(Complex::new(0., 0.5), 1.), Rgb(28, 72, 93));

        assert_eq!(
            complex_to_color(Complex::new(0.5, 0.5), 1.),
            Rgb(31, 126, 119)
        );

        assert_eq!(
            complex_to_color(Complex::new(0., -0.5), 1.),
            Rgb(91, 57, 35)
        );

        assert_eq!(
            complex_to_color(Complex::new(-1. / 6., 0.), 1.),
            Rgb(10, 5, 11)
        );
    }
}