okhsl 1.0.0

//! This is an extension of the [`oklab` crate](https://lib.rs/oklab) that adds HSV/HSL equivalents in the Ok Lab color space.

pub use oklab::*;
use std::f64::consts::PI;

/// HSV color
#[derive(Debug, Copy, Clone)]
#[doc(alias = "hsv")]
#[repr(C)]
pub struct Okhsv {
    /// Hue
    pub h: f64,

    /// Saturation
    pub s: f32,

    /// Value
    pub v: f32,
}

/// HSL color
#[derive(Debug, Copy, Clone)]
#[doc(alias = "hsl")]
#[repr(C)]
pub struct Okhsl {
    /// Hue
    pub h: f64,

    /// Saturation
    pub s: f32,

    /// Lightness
    pub l: f32,
}

impl From<Rgb<u8>> for Okhsl {
    fn from(c: Rgb<u8>) -> Self {
        oklab_to_okhsl(srgb_to_oklab(c))
    }
}

impl From<Rgb<u8>> for Okhsv {
    fn from(c: Rgb<u8>) -> Self {
        oklab_to_okhsv(srgb_to_oklab(c))
    }
}

impl From<Oklab> for Okhsl {
    fn from(c: oklab::Oklab) -> Self {
        oklab_to_okhsl(c)
    }
}

impl From<Oklab> for Okhsv {
    fn from(c: oklab::Oklab) -> Self {
        oklab_to_okhsv(c)
    }
}

impl From<Okhsl> for Oklab {
    fn from(c: Okhsl) -> Self {
        okhsl_to_oklab(c)
    }
}

impl From<Okhsv> for Oklab {
    fn from(c: Okhsv) -> Self {
        okhsv_to_oklab(c)
    }
}

impl Okhsv {
    #[must_use]
    pub fn to_oklab(self) -> Oklab {
        okhsv_to_oklab(self)
    }

    #[must_use]
    pub fn to_srgb(self) -> Rgb<u8> {
        oklab_to_srgb(okhsv_to_oklab(self))
    }
}

impl Okhsl {
    #[must_use]
    pub fn to_oklab(self) -> Oklab {
        okhsl_to_oklab(self)
    }

    #[must_use]
    pub fn to_srgb(self) -> Rgb<u8> {
        oklab_to_srgb(okhsl_to_oklab(self))
    }
}

// Finds the maximum saturation possible for a given hue that fits in sRGB
// Saturation here is defined as S = C/L
// a and b must be normalized so a^2 + b^2 == 1
fn compute_max_saturation(a: f32, b: f32) -> f32 {
    // Max saturation will be when one of r, g or b goes below zero.

    // Select different coefficients depending on which component goes below zero first
    let k0: f32;
    let k1: f32;
    let k2: f32;
    let k3: f32;
    let k4: f32;
    let wl: f32;
    let wm: f32;
    let ws: f32;

    if (-1.8817033_f32).mul_add(a, -(0.8093649 * b)) > 1.0 {
        // Red component
        k0 = 1.1908628;
        k1 = 1.7657673;
        k2 = 0.5966264;
        k3 = 0.755152;
        k4 = 0.5677124;
        wl = 4.0767417;
        wm = -3.3077116;
        ws = 0.23096994;
    } else if 1.8144411_f32.mul_add(a, -(1.1944528 * b)) > 1.0 {
        // Green component
        k0 = 0.73956515;
        k1 = -0.45954404;
        k2 = 0.08285427;
        k3 = 0.1254107;
        k4 = 0.14503204;
        wl = -1.268438;
        wm = 2.6097574;
        ws = -0.34131938;
    } else {
        // Blue component
        k0 = 1.3573365;
        k1 = -0.00915799;
        k2 = -1.1513021;
        k3 = -0.50559606;
        k4 = 0.00692167;
        wl = -0.0041960863;
        wm = -0.7034186;
        ws = 1.7076147;
    }

    // Approximate max saturation using a polynomial:
    let mut saturation = (k4 * a)
        .mul_add(b, (k3 * a).mul_add(a, k2.mul_add(b, k1.mul_add(a, k0))));

    // Do one step Halley's method to get closer
    // this gives an error less than 10e6, except for some blue hues where the dS/dh is close to infinite
    // this should be sufficient for most applications, otherwise do two/three steps

    let k_l = 0.39633778_f32.mul_add(a, 0.21580376 * b);
    let k_m = (-0.105561346_f32).mul_add(a, -(0.06385417 * b));
    let k_s = (-0.08948418_f32).mul_add(a, -(1.2914855 * b));

    {
        let l_ = saturation.mul_add(k_l, 1.);
        let m_ = saturation.mul_add(k_m, 1.);
        let s_ = saturation.mul_add(k_s, 1.);

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

        let l_d_s = 3. * (k_l * l_) * l_;
        let m_d_s = 3. * (k_m * m_) * m_;
        let s_d_s = 3. * (k_s * s_) * s_;

        let l_d_s2 = 6. * k_l * (k_l * l_);
        let m_d_s2 = 6. * k_m * (k_m * m_);
        let s_d_s2 = 6. * k_s * (k_s * s_);

        let f = ws.mul_add(s, wl.mul_add(l, wm * m));
        let f1 = ws.mul_add(s_d_s, wl.mul_add(l_d_s, wm * m_d_s));
        let f2 = ws.mul_add(s_d_s2, wl.mul_add(l_d_s2, wm * m_d_s2));

        saturation -= f * f1 / f1.mul_add(f1, -(0.5 * f * f2));
    }

    saturation
}

#[inline]
fn find_cusp(a: f32, b: f32) -> [f32; 2] {
    // First, find the maximum saturation (saturation S = C/L)
    let s_cusp = compute_max_saturation(a, b);

    // Convert to linear sRGB to find the first point where at least one of r,g or b >= 1:
    let l_cusp = scale_l(1.0, s_cusp, a, b);
    [l_cusp, (l_cusp * s_cusp)]
}

// Finds intersection of the line defined by
// L = L0 * (1. - t) + t * L1;
// C = t * C1;
// a and b must be normalized so a^2 + b^2 == 1
fn find_gamut_intersection(
    a: f32, b: f32, l_1: f32, c_1: f32, l_0: f32, cusp: Option<[f32; 2]>,
) -> f32 {
    // Find the cusp of the gamut triangle
    let [cusp_l, cusp_c] = cusp.unwrap_or_else(|| find_cusp(a, b));

    // Find the intersection for upper and lower half seprately
    let mut t;
    if (l_1 - l_0).mul_add(cusp_c, -((cusp_l - l_0) * c_1)) <= 0.0 {
        // Lower half

        t = cusp_c * l_0 / c_1.mul_add(cusp_l, cusp_c * (l_0 - l_1));
    } else {
        // Upper half

        // First intersect with triangle
        t = cusp_c * (l_0 - 1.0)
            / c_1.mul_add(cusp_l - 1.0, cusp_c * (l_0 - l_1));

        // Then one step Halley's method
        {
            let d_l = l_1 - l_0;
            let d_c = c_1;

            let k_l = 0.39633778_f32.mul_add(a, 0.21580376 * b);
            let k_m = (-0.105561346_f32).mul_add(a, -(0.06385417 * b));
            let k_s = (-0.08948418_f32).mul_add(a, -(1.2914855 * b));

            let l_dt = d_c.mul_add(k_l, d_l);
            let m_dt = d_c.mul_add(k_m, d_l);
            let s_dt = d_c.mul_add(k_s, d_l);

            // If higher accuracy is required, 2 or 3 iterations of the following block can be used:
            {
                let l = l_0.mul_add(1. - t, t * l_1);
                let c = t * c_1;

                let l_ = c.mul_add(k_l, l);
                let m_ = c.mul_add(k_m, l);
                let s_ = c.mul_add(k_s, l);

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

                let ldt = 3. * (l_dt * l_) * l_;
                let mdt = 3. * (m_dt * m_) * m_;
                let sdt = 3. * (s_dt * s_) * s_;

                let ldt2 = 6. * l_dt * (l_dt * l_);
                let mdt2 = 6. * m_dt * (m_dt * m_);
                let sdt2 = 6. * s_dt * (s_dt * s_);

                let r = 0.23096994_f32
                    .mul_add(s, 4.0767417_f32.mul_add(l, -(3.3077116 * m)))
                    - 1.;
                let r1 = 0.23096994_f32
                    .mul_add(sdt, 4.0767417_f32.mul_add(ldt, -(3.3077116 * mdt)));
                let r2 = 0.23096994_f32.mul_add(
                    sdt2,
                    4.0767417_f32.mul_add(ldt2, -(3.3077116 * mdt2)),
                );

                let u_r = r1 / r1.mul_add(r1, -(0.5 * r * r2));
                let t_r = -r * u_r;

                let g = 0.34131938_f32
                    .mul_add(-s, (-1.268438_f32).mul_add(l, 2.6097574 * m))
                    - 1.;
                let g1 = 0.34131938_f32
                    .mul_add(-sdt, (-1.268438_f32).mul_add(ldt, 2.6097574 * mdt));
                let g2 = 0.34131938_f32.mul_add(
                    -sdt2,
                    (-1.268438_f32).mul_add(ldt2, 2.6097574 * mdt2),
                );

                let u_g = g1 / g1.mul_add(g1, -(0.5 * g * g2));
                let t_g = -g * u_g;

                let b = 1.7076147_f32
                    .mul_add(s, (-0.0041960863_f32).mul_add(l, -(0.7034186 * m)))
                    - 1.;
                let b1 = 1.7076147_f32.mul_add(
                    sdt,
                    (-0.0041960863_f32).mul_add(ldt, -(0.7034186 * mdt)),
                );
                let b2 = 1.7076147_f32.mul_add(
                    sdt2,
                    (-0.0041960863_f32).mul_add(ldt2, -(0.7034186 * mdt2)),
                );

                let u_b = b1 / b1.mul_add(b1, -(0.5 * b * b2));
                let t_b = -b * u_b;

                let t_r = if u_r >= 0.0 { t_r } else { 10e5 };
                let t_g = if u_g >= 0.0 { t_g } else { 10e5 };
                let t_b = if u_b >= 0.0 { t_b } else { 10e5 };

                t += t_r.min(t_g.min(t_b));
            }
        }
    }

    t
}

fn toe(x: f32) -> f32 {
    let k_1: f32 = 0.206;
    let k_2: f32 = 0.03;
    let k_3: f32 = (1. + k_1) / (1. + k_2);

    0.5 * (k_3.mul_add(x, -k_1)
        + k_3
            .mul_add(x, -k_1)
            .mul_add(k_3.mul_add(x, -k_1), 4. * k_2 * (k_3 * x))
            .sqrt())
}

fn toe_inv(x: f32) -> f32 {
    let k_1 = 0.206;
    let k_2 = 0.03;
    let k_3 = (1. + k_1) / (1. + k_2);
    x.mul_add(x, k_1 * x) / (k_3 * (x + k_2))
}

fn st_mid(a_: f32, b_: f32) -> [f32; 2] {
    let s_mid = 0.11516993
        + 1. / a_.mul_add(
            a_.mul_add(
                a_.mul_add(
                    4.69891_f32.mul_add(a_, 5.387708_f32.mul_add(b_, -4.2489457)),
                    10.02301_f32.mul_add(-b_, -2.1370494),
                ),
                1.751984_f32.mul_add(b_, -2.1955736),
            ),
            4.1590123_f32.mul_add(b_, 7.4477897),
        );

    let t_mid = 0.11239642
        + 1. / a_.mul_add(
            a_.mul_add(
                a_.mul_add(
                    0.14661872_f32
                        .mul_add(-a_, 0.45399568_f32.mul_add(-b_, 0.00299215)),
                    0.6122399_f32.mul_add(b_, -0.27087943),
                ),
                0.9014812_f32.mul_add(b_, 0.40370612),
            ),
            0.6812438_f32.mul_add(-b_, 1.6132032),
        );
    [s_mid, t_mid]
}

fn st_max(a_: f32, b_: f32, cusp: Option<[f32; 2]>) -> [f32; 2] {
    let [l, c] = cusp.unwrap_or_else(|| find_cusp(a_, b_));

    [c / l, c / (1. - l)]
}

fn get_cs(l: f32, a_: f32, b_: f32) -> [f32; 3] {
    let cusp = find_cusp(a_, b_);

    let c_max = find_gamut_intersection(a_, b_, l, 1.0, l, Some(cusp));
    let [s_max, t_max] = st_max(a_, b_, Some(cusp));
    let [s_mid, t_mid] = st_mid(a_, b_);

    let k = c_max / (l * s_max).min((1. - l) * t_max);

    let c_mid = {
        let c_a = l * s_mid;
        let c_b = (1. - l) * t_mid;
        let ca4 = (c_a * c_a) * (c_a * c_a);
        let cb4 = (c_b * c_b) * (c_b * c_b);

        0.9 * k * ((1. / (1. / ca4 + 1. / cb4)).sqrt()).sqrt()
    };

    let c_0 = {
        let c_a = l * 0.4;
        let c_b = (1. - l) * 0.8;

        (1. / (1. / (c_a * c_a) + 1. / (c_b * c_b))).sqrt()
    };

    [c_0, c_mid, c_max]
}

#[must_use]
/// Convert [`Okhsl`] to [`Oklab`]
pub fn okhsl_to_oklab(Okhsl { h, s, l }: Okhsl) -> Oklab {
    if l == 0.0 {
        return srgb_f32_to_oklab([0.0, 0., 0.].into());
    }

    let a_ = (2. * PI * h).cos() as f32;
    let b_ = (2. * PI * h).sin() as f32;
    let l = toe_inv(l);

    let [c_0, c_mid, c_max] = get_cs(l, a_, b_);
    let t;
    let k_0;
    let k_1;
    let k_2;
    if s < 0.8 {
        t = 1.25 * s;
        k_0 = 0.;
        k_1 = 0.8 * c_0;
        k_2 = 1. - k_1 / c_mid;
    } else {
        t = 5. * (s - 0.8);
        k_0 = c_mid;
        k_1 = 0.2 * (c_mid * c_mid) * (1.25 * 1.25) / c_0;
        k_2 = 1. - k_1 / (c_max - c_mid);
    }

    let c = k_0 + t * k_1 / k_2.mul_add(-t, 1.);

    // If we would only use one of the Cs:
    //C = s*c_0;
    //C = s*1.25*c_mid;
    //C = s*c_max;

    Oklab {
        l,
        a: c * a_,
        b: c * b_,
    }
}

#[must_use]
#[doc(alias = "oklab_to_hsl")]
#[doc(alias = "lab_to_hsl")]
/// Convert [`Oklab`] to [`Okhsl`]
pub fn oklab_to_okhsl(Oklab { l, a, b }: Oklab) -> Okhsl {
    let (h, a_, b_, c) = hue(b.into(), a.into());

    let [c_0, c_mid, c_max] = get_cs(l, a_, b_);

    let s = if c < c_mid {
        let k_0 = 0.;
        let k_1 = 0.8 * c_0;
        let k_2 = 1. - k_1 / c_mid;

        let t = (c - k_0) / k_2.mul_add(c - k_0, k_1);
        t * 0.8
    } else {
        let k_0 = c_mid;
        let k_1 = 0.2 * (c_mid * c_mid) * (1.25 * 1.25) / c_0;
        let k_2 = 1. - k_1 / (c_max - c_mid);

        let t = (c - k_0) / k_2.mul_add(c - k_0, k_1);
        0.2_f32.mul_add(t, 0.8)
    };

    Okhsl { h, s, l: toe(l) }
}

fn hue(b: f64, a: f64) -> (f64, f32, f32, f32) {
    let h = (0.5 * (-b).atan2(-a)).mul_add(1. / PI, 0.5);
    let c = a.hypot(b);
    let a_ = (a * (1. / c)) as f32;
    let b_ = (b * (1. / c)) as f32;
    (h, a_, b_, c as f32)
}

#[must_use]
/// Convert [`Okhsv`] to [`Oklab`]
pub fn okhsv_to_oklab(Okhsv { h, s, v }: Okhsv) -> Oklab {
    let a_ = (2. * PI * h).cos() as f32;
    let b_ = (2. * PI * h).sin() as f32;

    let [s_max, t_max] = st_max(a_, b_, None);
    let s_0 = 0.5;
    let t = t_max;
    let k = 1. - s_0 / s_max;

    let l_v = 1. - s * s_0 / (t * k).mul_add(-s, s_0 + t);
    let c_v = s * t * s_0 / (t * k).mul_add(-s, s_0 + t);

    let mut l = v * l_v;
    let mut c = v * c_v;

    // to present steps along the way
    //L = v;
    //C = v*s*s_max;
    //L = v*(1. - s*s_max/(s_max+T));
    //C = v*s*s_max*T/(s_max+T);

    let l_vt = toe_inv(l_v);
    let c_vt = c_v * l_vt / l_v;

    let l_new = toe_inv(l); // * l_v/l_vt;
    c = c * l_new / l;
    l = l_new;

    let scale_l = scale_l(l_vt, c_vt, a_, b_);

    // remove to see effect without rescaling
    l *= scale_l;
    c *= scale_l;

    Oklab {
        l,
        a: c * a_,
        b: c * b_,
    }
}

#[must_use]
#[doc(alias = "oklab_to_hsv")]
#[doc(alias = "lab_to_hsv")]
/// Convert [`Oklab`] to [`Okhsv`]
pub fn oklab_to_okhsv(Oklab { l, a, b }: Oklab) -> Okhsv {
    let (h, a_, b_, c) = hue(b.into(), a.into());

    let [cusp_l, cusp_c] = find_cusp(a_, b_);
    let s_max = cusp_c / cusp_l;
    let t_max = cusp_c / (1. - cusp_l);

    let s_0 = 0.5;
    let t = t_max / l.mul_add(t_max, c);
    let k = 1. - s_0 / s_max;

    let l_v = t * l;
    let c_v = t * c;
    let l_vt = toe_inv(l_v);

    let c_vt = c_v * l_vt / l_v;
    let scale_l = scale_l(l_vt, c_vt, a_, b_);

    let mut l = l / scale_l;
    // C /= scale_l;
    // C = C * toe(l) / l;
    l = toe(l);

    let v = l / l_v;
    let s = (s_0 + t_max) * c_v / t_max.mul_add(s_0, t_max * k * c_v);
    Okhsv { h, s, v }
}

fn scale_l(l_vt: f32, c_vt: f32, a_: f32, b_: f32) -> f32 {
    let rgb_scale = (Oklab {
        l: l_vt,
        a: a_ * c_vt,
        b: b_ * c_vt,
    })
    .to_linear_rgb();
    let rgb_max = rgb_scale.r.max(rgb_scale.g).max(rgb_scale.b.max(0.));
    (1. / rgb_max).cbrt()
}

#[test]
fn hsl_roundtrips() {
    use rgb::Rgb;

    let mut total_diff = 0.;
    let mut max_diff = 0.;
    let mut skipped = 0;
    for r in 0..256u16 {
        for g in 0..256u16 {
            for b in 0..256u16 {
                let p1 = Rgb::new(
                    f32::from(r) / 255.0,
                    f32::from(g) / 255.0,
                    f32::from(b) / 255.0,
                );
                let hsl = oklab_to_okhsl(srgb_f32_to_oklab(p1));
                assert!(hsl.h.is_finite());
                assert!(hsl.l.is_finite());
                if !hsl.s.is_finite() {
                    skipped += 1;
                    continue;
                }
                let p2 = oklab_to_srgb_f32(okhsl_to_oklab(hsl));
                let diff = (p1.b - p2.b).mul_add(
                    p1.b - p2.b,
                    (p1.g - p2.g).mul_add(p1.g - p2.g, (p1.r - p2.r).powi(2)),
                );
                if diff > max_diff {
                    max_diff = diff;
                    eprintln!("{p1:?} {p2:?} diff {max_diff:0.5} {hsl:?}");
                }
                total_diff += f64::from(diff);
            }
        }
    }
    assert!(max_diff < 0.0006, "{max_diff}");
    assert!(total_diff < 1.3, "{total_diff}");
    assert!(skipped <= 2);
}

#[test]
fn hsv_roundtrips() {
    use rgb::Rgb;

    let mut total_diff = 0.;
    let mut max_diff = 0.;
    let mut skipped = 0;
    for r in 0..256u16 {
        for g in 0..256u16 {
            for b in 0..256u16 {
                let p1 = Rgb::new(
                    f32::from(r) / 255.0,
                    f32::from(g) / 255.0,
                    f32::from(b) / 255.0,
                );
                let hsv = oklab_to_okhsv(srgb_f32_to_oklab(p1));
                if !hsv.h.is_finite() || !hsv.s.is_finite() || !hsv.v.is_finite() {
                    eprintln!("{hsv:?}");
                    skipped += 1;
                    continue;
                }
                let p2 = oklab_to_srgb_f32(okhsv_to_oklab(hsv));
                let diff = (p1.b - p2.b).mul_add(
                    p1.b - p2.b,
                    (p1.g - p2.g).mul_add(p1.g - p2.g, (p1.r - p2.r).powi(2)),
                );
                if diff > max_diff {
                    max_diff = diff;
                    eprintln!("{p1:.4?} {p2:.4?} diff {max_diff:0.5} {hsv:?}");
                }
                total_diff += f64::from(diff);
            }
        }
    }
    assert!(max_diff < 0.0000025, "{max_diff}");
    assert!(total_diff < 0.0001, "{total_diff}"); // ooooof
    assert!(skipped <= 1, "{skipped}");
}