use std::sync::LazyLock;
use wide::f64x4;

pub trait Color {
    /// Computes the Euclidean distance between `self` and `other`.
    ///
    /// This method finds the difference of each component, squares each
    /// difference, sums the squares, then takes the square root of the
    /// sum.
    fn euclidian_distance(&self, other: &Self) -> f64;
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Rgb {
    pub r: u8,
    pub g: u8,
    pub b: u8,
}

impl Color for Rgb {
    fn euclidian_distance(&self, other: &Self) -> f64 {
        let diff_r = (self.r as f64) - (other.r as f64);
        let diff_g = (self.g as f64) - (other.g as f64);
        let diff_b = (self.b as f64) - (other.b as f64);
        (diff_r * diff_r + diff_g * diff_g + diff_b * diff_b).sqrt()
    }
}

impl Rgb {
    /// Blends `self` and `other` by `alpha`.
    ///
    /// This method finds a color between `self` and `other`. An `alpha`
    /// of `0.0` returns `self`. An `alpha` of `1.0` returns `other`. A
    /// value between `0.0` and `1.0` mixes the two colors in that
    /// proportion.
    ///
    /// The name `lerp` is short for "linear interpolation," a common term
    /// in graphics code.
    pub fn lerp(&self, other: &Rgb, alpha: f64) -> Rgb {
        let r = (1.0 - alpha) * self.r as f64 + alpha * other.r as f64;
        let g = (1.0 - alpha) * self.g as f64 + alpha * other.g as f64;
        let b = (1.0 - alpha) * self.b as f64 + alpha * other.b as f64;
        Rgb {
            r: r as u8,
            g: g as u8,
            b: b as u8,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Lab {
    pub l: f64,
    pub a: f64,
    pub b: f64,
}

impl Color for Lab {
    fn euclidian_distance(&self, other: &Self) -> f64 {
        let diff_l = self.l - other.l;
        let diff_a = self.a - other.a;
        let diff_b = self.b - other.b;
        (diff_l * diff_l + diff_a * diff_a + diff_b * diff_b).sqrt()
    }
}

/// Constants for the sRGB transfer function. `srgb_to_linear` uses these
/// to convert a normalized (`0.0..=1.0`) sRGB channel value to linear-light
/// intensity:
///
/// ```txt
/// srgb_channel > SRGB_LINEAR_THRESHOLD
///    ? ((srgb_channel + SRGB_ALPHA) / (1.0 + SRGB_ALPHA)) ^ SRGB_GAMMA
///    : srgb_channel / SRGB_LINEAR_DIVISOR
/// ```
///
/// The sRGB standard (IEC 61966-2-1) fixes all four values. They are not
/// tunable. A pure power curve has an infinite slope at zero. This causes
/// numerical instability for very dark values. The sRGB standard uses a
/// plain linear segment near black instead. `SRGB_LINEAR_THRESHOLD` and
/// `SRGB_LINEAR_DIVISOR` make the linear segment match the power-law
/// segment at the threshold. Both value and slope match. This leaves no
/// visible seam between the two pieces.
const SRGB_LINEAR_THRESHOLD: f64 = 0.04045;
const SRGB_ALPHA: f64 = 0.055;
const SRGB_GAMMA: f64 = 2.4;
const SRGB_LINEAR_DIVISOR: f64 = 12.92;

/// The linear sRGB to CIE 1931 XYZ conversion matrix (D65 white point).
///
/// Each row multiplies a `(linear_r, linear_g, linear_b)` triple to produce
/// one XYZ tristimulus value. Row 0 produces X. Row 1 produces Y (luminance).
/// Row 2 produces Z. See `linear_rgb_to_lab` and `linear_rgb_to_lab_simd`,
/// which compute `XYZ_MATRIX[i][0] * linear_r + XYZ_MATRIX[i][1] * linear_g +
/// XYZ_MATRIX[i][2] * linear_b` for each row `i`. This is the middle step
/// of RGB to Lab. It must use linear-light RGB (after `SRGB_TO_LINEAR`).
/// It must not use gamma-encoded values. The transform is only physically
/// meaningful between two linear color spaces.
///
/// The values come from the sRGB standard (IEC 61966-2-1). The standard
/// defines three primary chromaticities and the D65 white point. The matrix
/// solves for mixing the primaries at full intensity to reproduce D65 white.
/// This is a standard, publicly fixed matrix. It matches the `nordify`
/// Python reference implementation.
const XYZ_MATRIX: [[f64; 3]; 3] = [
    [0.4124564, 0.3575761, 0.1804375],
    [0.2126729, 0.7151522, 0.0721750],
    [0.0193339, 0.1191920, 0.9503041],
];

/// The CIE standard illuminant D65 reference white point in CIE XYZ
/// coordinates. `XYZ_MATRIX` produces XYZ values scaled so the sRGB
/// primaries mix to D65 white. Dividing by `WHITE_POINT` (in
/// `linear_rgb_to_lab` and `linear_rgb_to_lab_simd`, right before
/// `xyz_compress`) re-normalizes so white maps to `(1.0, 1.0, 1.0)`.
/// This makes the Lab convention "white has L* = 100" work correctly.
/// `Y` is exactly `1.0` by definition. Y represents luminance,
/// normalized so white has luminance 1. `X` and `Z` are not `1.0`
/// because D65 white does not have equal energy in those directions.
const WHITE_POINT: [f64; 3] = [0.95047, 1.00000, 1.08883];

/// Constants for the CIE Lab forward compression function `f(t)`.
/// `xyz_compress` and `xyz_compress_simd` use these to convert a
/// white-point-normalized XYZ component into the perceptually-uniform
/// value that the L*/a*/b* formulas use:
///
/// ```txt
/// t > LAB_CUBE_THRESHOLD ? cbrt(t) : LAB_LINEAR_SLOPE * t + LAB_LINEAR_INTERCEPT
/// ```
///
/// This piecewise split avoids the numerical instability of a cube root
/// infinite slope near zero. CIE fixes all three values so the linear
/// segment meets the cube-root curve smoothly at `LAB_CUBE_THRESHOLD`.
/// `LAB_LINEAR_INTERCEPT` (`16/116`) has a second role. The final
/// `L = 116 * f(Y) - 16` formula subtracts this value back out. A
/// pure-black input (`f(0) = 16/116`) maps to `L = 0`.
const LAB_CUBE_THRESHOLD: f64 = 0.008856;
const LAB_LINEAR_SLOPE: f64 = 7.787;
const LAB_LINEAR_INTERCEPT: f64 = 16.0 / 116.0;

/// Converts a normalized (`0.0..=1.0`) gamma-corrected sRGB channel
/// value to its linear-light equivalent. This reverses the sRGB transfer
/// function.
fn srgb_to_linear(srgb_channel: f64) -> f64 {
    if srgb_channel > SRGB_LINEAR_THRESHOLD {
        ((srgb_channel + SRGB_ALPHA) / (1.0 + SRGB_ALPHA)).powf(SRGB_GAMMA)
    } else {
        srgb_channel / SRGB_LINEAR_DIVISOR
    }
}

/// Precomputed `srgb_to_linear` result for every possible `u8` channel
/// value (`0..=255`). This table builds once from the scalar formula so
/// it cannot drift out of sync. An sRGB channel takes 256 distinct values.
/// This table turns "linearize a channel" into an O(1) array read
/// (`SRGB_TO_LINEAR[value as usize]`) instead of a `powf` call. The SIMD
/// chunk path (`linearize_chunk_simd`) uses it with 4 lookups packed into
/// an `f64x4`. Scalar code (`From<Rgb> for Lab`, `convert_from` remainder
/// tail) also uses it. This table avoids `wide`'s `powf_simd`, whose
/// precision is unspecified and can vary by platform. This table gives
/// deterministic results everywhere. This matters for tolerance-based tests.
static SRGB_TO_LINEAR: LazyLock<[f64; 256]> = LazyLock::new(|| {
    let mut table = [0.0f64; 256];
    for (channel_value, linear) in table.iter_mut().enumerate() {
        *linear = srgb_to_linear(channel_value as f64 / 255.0);
    }
    table
});

/// Applies the CIE Lab forward compression function `f(t)` to a single
/// white-point-normalized XYZ tristimulus component. Use this to derive
/// L*/a*/b* from X/Y/Z.
fn xyz_compress(xyz_component: f64) -> f64 {
    if xyz_component > LAB_CUBE_THRESHOLD {
        xyz_component.cbrt()
    } else {
        LAB_LINEAR_SLOPE * xyz_component + LAB_LINEAR_INTERCEPT
    }
}

fn xyz_compress_simd(xyz_component: f64x4) -> f64x4 {
    let cube_root = xyz_component.cbrt();
    let linear =
        f64x4::splat(LAB_LINEAR_SLOPE) * xyz_component + f64x4::splat(LAB_LINEAR_INTERCEPT);
    xyz_component
        .simd_gt(f64x4::splat(LAB_CUBE_THRESHOLD))
        .select(cube_root, linear)
}

/// Computes Lab from linear-light RGB (after sRGB decoding). The scalar
/// `From<Rgb> for Lab` and `convert_from` scalar remainder tail share this
/// function.
fn linear_rgb_to_lab(linear_r: f64, linear_g: f64, linear_b: f64) -> Lab {
    let x = xyz_compress(
        (XYZ_MATRIX[0][0] * linear_r + XYZ_MATRIX[0][1] * linear_g + XYZ_MATRIX[0][2] * linear_b)
            / WHITE_POINT[0],
    );
    let y = xyz_compress(
        (XYZ_MATRIX[1][0] * linear_r + XYZ_MATRIX[1][1] * linear_g + XYZ_MATRIX[1][2] * linear_b)
            / WHITE_POINT[1],
    );
    let z = xyz_compress(
        (XYZ_MATRIX[2][0] * linear_r + XYZ_MATRIX[2][1] * linear_g + XYZ_MATRIX[2][2] * linear_b)
            / WHITE_POINT[2],
    );

    Lab {
        l: 116.0 * y - 16.0,
        a: 500.0 * (x - y),
        b: 200.0 * (y - z),
    }
}

/// Converts 4 linear-light RGB triples at once. This is the vectorized
/// counterpart to `linear_rgb_to_lab`. It returns `(l, a, b)` lane vectors.
fn linear_rgb_to_lab_simd(
    linear_r: f64x4,
    linear_g: f64x4,
    linear_b: f64x4,
) -> (f64x4, f64x4, f64x4) {
    let x = xyz_compress_simd(
        (f64x4::splat(XYZ_MATRIX[0][0]) * linear_r
            + f64x4::splat(XYZ_MATRIX[0][1]) * linear_g
            + f64x4::splat(XYZ_MATRIX[0][2]) * linear_b)
            / f64x4::splat(WHITE_POINT[0]),
    );
    let y = xyz_compress_simd(
        (f64x4::splat(XYZ_MATRIX[1][0]) * linear_r
            + f64x4::splat(XYZ_MATRIX[1][1]) * linear_g
            + f64x4::splat(XYZ_MATRIX[1][2]) * linear_b)
            / f64x4::splat(WHITE_POINT[1]),
    );
    let z = xyz_compress_simd(
        (f64x4::splat(XYZ_MATRIX[2][0]) * linear_r
            + f64x4::splat(XYZ_MATRIX[2][1]) * linear_g
            + f64x4::splat(XYZ_MATRIX[2][2]) * linear_b)
            / f64x4::splat(WHITE_POINT[2]),
    );

    let l = f64x4::splat(116.0) * y - f64x4::splat(16.0);
    let a = f64x4::splat(500.0) * (x - y);
    let b = f64x4::splat(200.0) * (y - z);

    (l, a, b)
}

/// Linearizes 4 raw (`0..=255`) sRGB channel values at once. This uses 4
/// independent `SRGB_TO_LINEAR` lookups packed into one vector. Table
/// lookups avoid `powf_simd` documented precision non-determinism.
fn linearize_chunk_simd(raw_channel_chunk: &[f64]) -> f64x4 {
    f64x4::new([
        SRGB_TO_LINEAR[raw_channel_chunk[0] as usize],
        SRGB_TO_LINEAR[raw_channel_chunk[1] as usize],
        SRGB_TO_LINEAR[raw_channel_chunk[2] as usize],
        SRGB_TO_LINEAR[raw_channel_chunk[3] as usize],
    ])
}

impl From<Rgb> for Lab {
    /// Converts an 8-bit sRGB color to CIE Lab. The conversion uses
    /// linear-light RGB and XYZ with a D65 white point. This gives
    /// perceptually accurate color distance.
    fn from(value: Rgb) -> Self {
        let linear_r = SRGB_TO_LINEAR[value.r as usize];
        let linear_g = SRGB_TO_LINEAR[value.g as usize];
        let linear_b = SRGB_TO_LINEAR[value.b as usize];
        linear_rgb_to_lab(linear_r, linear_g, linear_b)
    }
}

pub trait PlanarColor<T: Color> {
    fn len(&self) -> usize;
    fn is_empty(&self) -> bool;
    fn with_capacity(n: usize) -> Self;
    fn clear(&mut self);
    fn push(&mut self, color: T);
}

/// Planar (struct-of-arrays) RGB storage. One contiguous `Vec<f64>` holds
/// each channel. Channel values are raw (`0..=255`, matching the `u8`
/// inputs). They are not normalized to `0..1`. This keeps `SRGB_TO_LINEAR`
/// indexing (`value as usize`) exact. Reconstructing an integer index from
/// a normalized float would lose precision. Distance and argmin comparisons
/// work the same either way. Scaling every coordinate by the same constant
/// preserves relative distances.
pub struct RgbPlanar {
    r: Vec<f64>,
    g: Vec<f64>,
    b: Vec<f64>,
}

impl RgbPlanar {
    pub fn new(r: Vec<f64>, g: Vec<f64>, b: Vec<f64>) -> Self {
        assert_eq!(r.len(), g.len());
        assert_eq!(r.len(), b.len());
        Self { r, g, b }
    }

    pub fn r(&self) -> &[f64] {
        &self.r
    }
    pub fn g(&self) -> &[f64] {
        &self.g
    }
    pub fn b(&self) -> &[f64] {
        &self.b
    }
}

impl From<&[Rgb]> for RgbPlanar {
    fn from(value: &[Rgb]) -> Self {
        // let collected: Vec<(u8, u8, u8)> = value.iter().map(|rgb| (rgb.r, rgb.g, rgb.b)).collect();
        let r: Vec<f64> = value.iter().map(|rgb| rgb.r as f64).collect();
        let g: Vec<f64> = value.iter().map(|rgb| rgb.g as f64).collect();
        let b: Vec<f64> = value.iter().map(|rgb| rgb.b as f64).collect();
        Self { r, g, b }
    }
}

impl PlanarColor<Rgb> for RgbPlanar {
    fn len(&self) -> usize {
        self.r.len()
    }

    fn is_empty(&self) -> bool {
        self.r.is_empty()
    }

    fn with_capacity(n: usize) -> Self {
        Self {
            r: Vec::with_capacity(n),
            g: Vec::with_capacity(n),
            b: Vec::with_capacity(n),
        }
    }

    fn clear(&mut self) {
        self.r.clear();
        self.g.clear();
        self.b.clear();
    }

    fn push(&mut self, color: Rgb) {
        self.r.push(f64::from(color.r));
        self.g.push(f64::from(color.g));
        self.b.push(f64::from(color.b));
    }
}

pub struct LabPlanar {
    l: Vec<f64>,
    a: Vec<f64>,
    b: Vec<f64>,
}

impl LabPlanar {
    pub fn new(l: Vec<f64>, a: Vec<f64>, b: Vec<f64>) -> Self {
        assert_eq!(l.len(), a.len());
        assert_eq!(l.len(), b.len());
        Self { l, a, b }
    }

    pub fn l(&self) -> &[f64] {
        &self.l
    }
    pub fn a(&self) -> &[f64] {
        &self.a
    }
    pub fn b(&self) -> &[f64] {
        &self.b
    }

    /// Repopulates `self` from `rgb` per-color RGB values. This converts
    /// to Lab via linear-light RGB and XYZ. It reuses the existing
    /// allocation in `self`. It clears and refills in place instead of
    /// allocating a new `LabPlanar`. This processes 4 colors at a time
    /// via `wide::f64x4`. A scalar tail handles `rgb.len() % 4` leftover
    /// colors.
    pub fn convert_from(&mut self, rgb: &RgbPlanar) {
        self.clear();
        let len = rgb.len();
        let chunks = len / 4;

        for chunk_index in 0..chunks {
            let i = chunk_index * 4;
            let linear_r = linearize_chunk_simd(&rgb.r()[i..i + 4]);
            let linear_g = linearize_chunk_simd(&rgb.g()[i..i + 4]);
            let linear_b = linearize_chunk_simd(&rgb.b()[i..i + 4]);

            let (l, a, b) = linear_rgb_to_lab_simd(linear_r, linear_g, linear_b);
            let (l, a, b) = (l.to_array(), a.to_array(), b.to_array());

            for lane in 0..4 {
                self.push(Lab {
                    l: l[lane],
                    a: a[lane],
                    b: b[lane],
                });
            }
        }

        for i in (chunks * 4)..len {
            let linear_r = SRGB_TO_LINEAR[rgb.r()[i] as usize];
            let linear_g = SRGB_TO_LINEAR[rgb.g()[i] as usize];
            let linear_b = SRGB_TO_LINEAR[rgb.b()[i] as usize];
            self.push(linear_rgb_to_lab(linear_r, linear_g, linear_b));
        }
    }
}

impl PlanarColor<Lab> for LabPlanar {
    fn len(&self) -> usize {
        self.l.len()
    }

    fn is_empty(&self) -> bool {
        self.l.is_empty()
    }

    fn with_capacity(n: usize) -> Self {
        Self {
            l: Vec::with_capacity(n),
            a: Vec::with_capacity(n),
            b: Vec::with_capacity(n),
        }
    }

    fn clear(&mut self) {
        self.l.clear();
        self.a.clear();
        self.b.clear();
    }

    fn push(&mut self, color: Lab) {
        self.l.push(color.l);
        self.a.push(color.a);
        self.b.push(color.b);
    }
}

impl From<&RgbPlanar> for LabPlanar {
    fn from(value: &RgbPlanar) -> Self {
        let mut out = Self::with_capacity(value.len());
        out.convert_from(value);
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::palette::NORD_PALETTE;

    const TOLERANCE: f64 = 1e-6;

    fn assert_lab_close(actual: Lab, expected: Lab) {
        assert!(
            (actual.l - expected.l).abs() < TOLERANCE,
            "l: {} vs {}",
            actual.l,
            expected.l
        );
        assert!(
            (actual.a - expected.a).abs() < TOLERANCE,
            "a: {} vs {}",
            actual.a,
            expected.a
        );
        assert!(
            (actual.b - expected.b).abs() < TOLERANCE,
            "b: {} vs {}",
            actual.b,
            expected.b
        );
    }

    fn edge_case_colors() -> Vec<Rgb> {
        let mut colors: Vec<Rgb> = NORD_PALETTE.to_vec();
        colors.push(Rgb { r: 0, g: 0, b: 0 });
        colors.push(Rgb {
            r: 255,
            g: 255,
            b: 255,
        });
        colors.push(Rgb { r: 255, g: 0, b: 0 });
        colors
    }

    fn to_rgb_planar(colors: &[Rgb]) -> RgbPlanar {
        let mut planar = RgbPlanar::with_capacity(colors.len());
        for &color in colors {
            planar.push(color);
        }
        planar
    }

    fn assert_lab_planar_matches_scalar(colors: &[Rgb], lab_planar: &LabPlanar) {
        assert_eq!(lab_planar.len(), colors.len());
        for (index, &color) in colors.iter().enumerate() {
            let expected = Lab::from(color);
            let actual = Lab {
                l: lab_planar.l()[index],
                a: lab_planar.a()[index],
                b: lab_planar.b()[index],
            };
            assert_lab_close(actual, expected);
        }
    }

    #[test]
    fn lab_planar_from_rgb_planar_matches_scalar() {
        let colors = edge_case_colors();
        let rgb_planar = to_rgb_planar(&colors);
        let lab_planar = LabPlanar::from(&rgb_planar);

        assert_lab_planar_matches_scalar(&colors, &lab_planar);
    }

    #[test]
    fn lab_planar_convert_from_handles_non_multiple_of_lane_width() {
        for len in [5, 7] {
            let colors: Vec<Rgb> = edge_case_colors().into_iter().take(len).collect();
            let rgb_planar = to_rgb_planar(&colors);
            let lab_planar = LabPlanar::from(&rgb_planar);

            assert_lab_planar_matches_scalar(&colors, &lab_planar);
        }
    }

    #[test]
    fn lab_planar_convert_from_reuses_allocation() {
        let long_colors = edge_case_colors();
        let short_colors: Vec<Rgb> = long_colors.iter().take(3).copied().collect();

        let long_rgb = to_rgb_planar(&long_colors);
        let short_rgb = to_rgb_planar(&short_colors);

        let mut lab_planar = LabPlanar::with_capacity(long_colors.len());
        lab_planar.convert_from(&long_rgb);
        lab_planar.convert_from(&short_rgb);

        assert_lab_planar_matches_scalar(&short_colors, &lab_planar);
    }
}