use crate::core::geometry::Point2f;
use crate::core::spectrum::Spectrum;
use crate::utils::find_interval;
use crate::utils::math::{clamp, evaluate_polynomial, lerp, SquareMatrix, SquareMatrix3f};
use crate::{Float, GVec, Ptr};
use core::any::TypeId;
use core::fmt;
use core::ops::{
    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};
use enum_dispatch::enum_dispatch;
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct XYZ {
    pub x: Float,
    pub y: Float,
    pub z: Float,
}

impl From<(Float, Float, Float)> for XYZ {
    fn from(triplet: (Float, Float, Float)) -> Self {
        XYZ::new(triplet.0, triplet.1, triplet.2)
    }
}

impl From<[Float; 3]> for XYZ {
    fn from(triplet: [Float; 3]) -> Self {
        XYZ::new(triplet[0], triplet[1], triplet[2])
    }
}

impl<'a> IntoIterator for &'a XYZ {
    type Item = &'a Float;
    type IntoIter = core::array::IntoIter<&'a Float, 3>;

    fn into_iter(self) -> Self::IntoIter {
        [&self.x, &self.y, &self.z].into_iter()
    }
}

impl XYZ {
    pub fn new(x: Float, y: Float, z: Float) -> Self {
        Self { x, y, z }
    }

    pub fn x(&self) -> Float {
        self.x
    }

    pub fn y(&self) -> Float {
        self.y
    }

    pub fn z(&self) -> Float {
        self.z
    }

    pub fn average(&self) -> Float {
        (self.x + self.y + self.z) / 3.0
    }

    pub fn xy(&self) -> Point2f {
        let sum = self.x + self.y + self.z;
        Point2f::new(self.x / sum, self.y / sum)
    }

    pub fn from_xyy(xy: Point2f, y: Option<Float>) -> Self {
        let scale: Float;
        if let Some(y_val) = y {
            scale = y_val;
        } else {
            scale = 1.;
        }
        if xy.y() == 0.0 {
            return Self {
                x: 0.0,
                y: 0.0,
                z: 0.0,
            };
        }
        Self::new(
            xy.x() * scale / xy.y(),
            scale,
            (1.0 - xy.x() - xy.y()) * scale / xy.y(),
        )
    }
}

impl Index<u32> for XYZ {
    type Output = Float;
    fn index(&self, index: u32) -> &Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &self.x,
            1 => &self.y,
            _ => &self.z,
        }
    }
}

impl IndexMut<u32> for XYZ {
    fn index_mut(&mut self, index: u32) -> &mut Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &mut self.x,
            1 => &mut self.y,
            _ => &mut self.z,
        }
    }
}

impl Neg for XYZ {
    type Output = Self;
    fn neg(self) -> Self {
        Self {
            x: -self.x,
            y: -self.y,
            z: -self.z,
        }
    }
}

impl Add for XYZ {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
            z: self.z + rhs.z,
        }
    }
}

impl AddAssign for XYZ {
    fn add_assign(&mut self, rhs: Self) {
        self.x += rhs.x;
        self.y += rhs.y;
        self.z += rhs.z;
    }
}

impl Sub for XYZ {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            x: self.x - rhs.x,
            y: self.y - rhs.y,
            z: self.z - rhs.z,
        }
    }
}

impl SubAssign for XYZ {
    fn sub_assign(&mut self, rhs: Self) {
        self.x -= rhs.x;
        self.y -= rhs.y;
        self.z -= rhs.z;
    }
}

impl Sub<XYZ> for Float {
    type Output = XYZ;
    fn sub(self, rhs: XYZ) -> XYZ {
        XYZ::new(self - rhs.x, self - rhs.y, self - rhs.z)
    }
}

impl Mul for XYZ {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Self {
            x: self.x * rhs.x,
            y: self.y * rhs.y,
            z: self.z * rhs.z,
        }
    }
}

impl MulAssign for XYZ {
    fn mul_assign(&mut self, rhs: Self) {
        self.x *= rhs.x;
        self.y *= rhs.y;
        self.z *= rhs.z;
    }
}

impl Mul<Float> for XYZ {
    type Output = Self;
    fn mul(self, rhs: Float) -> Self {
        Self {
            x: self.x * rhs,
            y: self.y * rhs,
            z: self.z * rhs,
        }
    }
}

impl MulAssign<Float> for XYZ {
    fn mul_assign(&mut self, rhs: Float) {
        self.x *= rhs;
        self.y *= rhs;
        self.z *= rhs;
    }
}

impl Mul<XYZ> for Float {
    type Output = XYZ;
    fn mul(self, rhs: XYZ) -> XYZ {
        rhs * self
    }
}

impl Div for XYZ {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        Self {
            x: self.x / rhs.x,
            y: self.y / rhs.y,
            z: self.z / rhs.z,
        }
    }
}

impl DivAssign for XYZ {
    fn div_assign(&mut self, rhs: Self) {
        self.x /= rhs.x;
        self.y /= rhs.y;
        self.z /= rhs.z;
    }
}

impl Div<Float> for XYZ {
    type Output = Self;
    fn div(self, rhs: Float) -> Self {
        debug_assert_ne!(rhs, 0.0, "Division by zero.");
        let inv = 1.0 / rhs;
        self * inv
    }
}

impl DivAssign<Float> for XYZ {
    fn div_assign(&mut self, rhs: Float) {
        debug_assert_ne!(rhs, 0.0, "Division by zero.");
        let inv = 1.0 / rhs;
        *self *= inv;
    }
}

impl fmt::Display for XYZ {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "[ {}, {}, {} ]", self.x, self.y, self.z)
    }
}

#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct RGB {
    pub r: Float,
    pub g: Float,
    pub b: Float,
}

impl From<[Float; 3]> for RGB {
    fn from(slice: [Float; 3]) -> Self {
        RGB::new(slice[0], slice[1], slice[2])
    }
}

impl From<&[Float; 3]> for RGB {
    fn from(slice: &[Float; 3]) -> Self {
        RGB::new(slice[0], slice[1], slice[2])
    }
}

impl From<(Float, Float, Float)> for RGB {
    fn from(triplet: (Float, Float, Float)) -> Self {
        RGB::new(triplet.0, triplet.1, triplet.2)
    }
}

impl<'a> IntoIterator for &'a RGB {
    type Item = &'a Float;
    type IntoIter = core::array::IntoIter<&'a Float, 3>;

    fn into_iter(self) -> Self::IntoIter {
        [&self.r, &self.g, &self.b].into_iter()
    }
}

impl RGB {
    pub fn new(r: Float, g: Float, b: Float) -> Self {
        Self { r, g, b }
    }

    pub fn average(&self) -> Float {
        (self.r + self.g + self.b) / 3.0
    }

    pub fn max_component_value(&self) -> Float {
        self.r.max(self.g).max(self.b)
    }

    pub fn min_component_value(&self) -> Float {
        self.r.min(self.g).min(self.b)
    }

    pub fn min_component_index(&self) -> u32 {
        if self.r < self.g {
            if self.r < self.b {
                0
            } else {
                2
            }
        } else {
            if self.g < self.b {
                1
            } else {
                2
            }
        }
    }

    pub fn max_component_index(&self) -> u32 {
        if self.r > self.g {
            if self.r > self.b {
                0
            } else {
                2
            }
        } else {
            if self.g > self.b {
                1
            } else {
                2
            }
        }
    }

    pub fn clamp(&self, min: Float, max: Float) -> Self {
        RGB::new(
            clamp(self.r, min, max),
            clamp(self.g, min, max),
            clamp(self.b, min, max),
        )
    }

    pub fn clamp_zero(&self) -> Self {
        RGB::new(self.r.max(0.), self.b.max(0.), self.g.max(0.))
    }
}

impl Index<u32> for RGB {
    type Output = Float;
    fn index(&self, index: u32) -> &Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &self.r,
            1 => &self.g,
            _ => &self.b,
        }
    }
}

impl Index<i32> for RGB {
    type Output = Float;
    fn index(&self, index: i32) -> &Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &self.r,
            1 => &self.g,
            _ => &self.b,
        }
    }
}

impl Index<usize> for RGB {
    type Output = Float;
    fn index(&self, index: usize) -> &Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &self.r,
            1 => &self.g,
            _ => &self.b,
        }
    }
}

impl IndexMut<u32> for RGB {
    fn index_mut(&mut self, index: u32) -> &mut Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &mut self.r,
            1 => &mut self.g,
            _ => &mut self.b,
        }
    }
}

impl IndexMut<i32> for RGB {
    fn index_mut(&mut self, index: i32) -> &mut Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &mut self.r,
            1 => &mut self.g,
            _ => &mut self.b,
        }
    }
}

impl IndexMut<usize> for RGB {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &mut self.r,
            1 => &mut self.g,
            _ => &mut self.b,
        }
    }
}

impl Neg for RGB {
    type Output = Self;
    fn neg(self) -> Self {
        Self {
            r: -self.r,
            g: -self.g,
            b: -self.b,
        }
    }
}

impl Add for RGB {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self {
            r: self.r + rhs.r,
            g: self.g + rhs.g,
            b: self.b + rhs.b,
        }
    }
}

impl AddAssign for RGB {
    fn add_assign(&mut self, rhs: Self) {
        self.r += rhs.r;
        self.g += rhs.g;
        self.b += rhs.b;
    }
}

impl Sub for RGB {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            r: self.r - rhs.r,
            g: self.g - rhs.g,
            b: self.b - rhs.b,
        }
    }
}

impl SubAssign for RGB {
    fn sub_assign(&mut self, rhs: Self) {
        self.r -= rhs.r;
        self.g -= rhs.g;
        self.b -= rhs.b;
    }
}

impl Sub<RGB> for Float {
    type Output = RGB;
    fn sub(self, rhs: RGB) -> RGB {
        RGB::new(self - rhs.r, self - rhs.g, self - rhs.b)
    }
}

impl Mul for RGB {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Self {
            r: self.r * rhs.r,
            g: self.g * rhs.g,
            b: self.b * rhs.b,
        }
    }
}

impl MulAssign for RGB {
    fn mul_assign(&mut self, rhs: Self) {
        self.r *= rhs.r;
        self.g *= rhs.g;
        self.b *= rhs.b;
    }
}

// Scalar multiplication (ryz * float)
impl Mul<Float> for RGB {
    type Output = Self;
    fn mul(self, rhs: Float) -> Self {
        Self {
            r: self.r * rhs,
            g: self.g * rhs,
            b: self.b * rhs,
        }
    }
}

impl MulAssign<Float> for RGB {
    fn mul_assign(&mut self, rhs: Float) {
        self.r *= rhs;
        self.g *= rhs;
        self.b *= rhs;
    }
}

impl Mul<RGB> for Float {
    type Output = RGB;
    fn mul(self, rhs: RGB) -> RGB {
        rhs * self
    }
}

impl Div for RGB {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        Self {
            r: self.r / rhs.r,
            g: self.g / rhs.g,
            b: self.b / rhs.b,
        }
    }
}

impl DivAssign for RGB {
    fn div_assign(&mut self, rhs: Self) {
        self.r /= rhs.r;
        self.g /= rhs.g;
        self.b /= rhs.b;
    }
}

impl Div<Float> for RGB {
    type Output = Self;
    fn div(self, rhs: Float) -> Self {
        debug_assert_ne!(rhs, 0.0, "Division by zero.");
        let inv = 1.0 / rhs;
        self * inv
    }
}

impl DivAssign<Float> for RGB {
    fn div_assign(&mut self, rhs: Float) {
        debug_assert_ne!(rhs, 0.0, "Division by zero.");
        let inv = 1.0 / rhs;
        *self *= inv;
    }
}

impl fmt::Display for RGB {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "[ {}, {}, {} ]", self.r, self.g, self.b)
    }
}

impl Mul<XYZ> for SquareMatrix3f {
    type Output = RGB;

    fn mul(self, v: XYZ) -> RGB {
        let r = self[0][0] * v.x + self[0][1] * v.y + self[0][2] * v.z;
        let g = self[1][0] * v.x + self[1][1] * v.y + self[1][2] * v.z;
        let b = self[2][0] * v.x + self[2][1] * v.y + self[2][2] * v.z;
        RGB::new(r, g, b)
    }
}

impl Mul<RGB> for SquareMatrix3f {
    type Output = XYZ;
    fn mul(self, v: RGB) -> XYZ {
        let x = self[0][0] * v.r + self[0][1] * v.g + self[0][2] * v.b;
        let y = self[1][0] * v.r + self[1][1] * v.g + self[1][2] * v.b;
        let z = self[2][0] * v.r + self[2][1] * v.g + self[2][2] * v.b;
        XYZ::new(x, y, z)
    }
}

pub trait MatrixMulColor {
    fn mul_rgb(&self, v: RGB) -> RGB;
    fn mul_xyz(&self, v: XYZ) -> XYZ;
}

impl MatrixMulColor for SquareMatrix3f {
    fn mul_rgb(&self, v: RGB) -> RGB {
        let m = self;
        RGB::new(
            m[0][0] * v.r + m[0][1] * v.g + m[0][2] * v.b,
            m[1][0] * v.r + m[1][1] * v.g + m[1][2] * v.b,
            m[2][0] * v.r + m[2][1] * v.g + m[2][2] * v.b,
        )
    }

    fn mul_xyz(&self, v: XYZ) -> XYZ {
        let m = self;
        XYZ::new(
            m[0][0] * v.x + m[0][1] * v.y + m[0][2] * v.z,
            m[1][0] * v.x + m[1][1] * v.y + m[1][2] * v.z,
            m[2][0] * v.x + m[2][1] * v.y + m[2][2] * v.z,
        )
    }
}

#[repr(C)]
#[derive(Debug, Default, Copy, Clone)]
pub struct RGBSigmoidPolynomial {
    pub c0: Float,
    pub c1: Float,
    pub c2: Float,
}

impl RGBSigmoidPolynomial {
    pub fn new(c0: Float, c1: Float, c2: Float) -> Self {
        Self { c0, c1, c2 }
    }

    pub fn evaluate(&self, lambda: Float) -> Float {
        let eval = evaluate_polynomial(lambda, &[self.c0, self.c1, self.c2]);
        Self::s(eval)
    }

    pub fn max_value(&self) -> Float {
        let lambda = -self.c1 / (2.0 * self.c0);
        let result = self.evaluate(360.0).max(self.evaluate(830.0));
        if (360.0..830.0).contains(&lambda) {
            return result.max(self.evaluate(lambda));
        }
        result
    }

    fn s(x: Float) -> Float {
        if x.is_infinite() {
            if x > 0.0 {
                return 1.0;
            } else {
                return 0.0;
            }
        }
        0.5 + x / (2.0 * (1.0 + (x * x)).sqrt())
    }
}

#[enum_dispatch]
pub trait ColorEncodingTrait: 'static + Send + Sync {
    fn from_linear(&self, vin: &[Float], vout: &mut [u8]);
    fn to_linear(&self, vin: &[u8], vout: &mut [Float]);
    fn to_float_linear(&self, v: Float) -> Float;

    fn from_linear_scalar(&self, v: Float) -> u8 {
        let mut out = [0u8; 1];
        self.from_linear(&[v], &mut out);
        out[0]
    }

    fn to_linear_scalar(&self, v: u8) -> Float {
        let mut out = [0.0; 1];
        self.to_linear(&[v], &mut out);
        out[0]
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

#[repr(C)]
#[enum_dispatch(ColorEncodingTrait)]
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum ColorEncoding {
    Linear(LinearEncoding),
    SRGB(SRGBEncoding),
}

impl fmt::Display for ColorEncoding {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Encoding")
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct LinearEncoding;
impl ColorEncodingTrait for LinearEncoding {
    fn from_linear(&self, vin: &[Float], vout: &mut [u8]) {
        for (i, &v) in vin.iter().enumerate() {
            vout[i] = (v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8;
        }
    }

    fn to_linear(&self, vin: &[u8], vout: &mut [Float]) {
        for (i, &v) in vin.iter().enumerate() {
            vout[i] = v as Float / 255.0;
        }
    }

    fn to_float_linear(&self, v: Float) -> Float {
        v
    }

    fn from_linear_scalar(&self, v: Float) -> u8 {
        (v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8
    }

    fn to_linear_scalar(&self, v: u8) -> Float {
        v as Float / 255.0
    }
}

impl fmt::Display for LinearEncoding {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Linear")
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct SRGBEncoding;
impl ColorEncodingTrait for SRGBEncoding {
    fn from_linear(&self, vin: &[Float], vout: &mut [u8]) {
        for (i, &v_linear) in vin.iter().enumerate() {
            let v = v_linear.clamp(0.0, 1.0);
            let v_encoded = if v <= 0.0031308 {
                v * 12.92
            } else {
                1.055 * v.powf(1.0 / 2.4) - 0.055
            };
            vout[i] = (v_encoded * 255.0 + 0.5) as u8;
        }
    }

    fn to_linear(&self, vin: &[u8], vout: &mut [Float]) {
        for (i, &v) in vin.iter().enumerate() {
            vout[i] = SRGB_TO_LINEAR_LUT[v as usize];
        }
    }

    fn from_linear_scalar(&self, v: Float) -> u8 {
        let v_clamped = v.clamp(0.0, 1.0);
        let v_encoded = if v_clamped <= 0.0031308 {
            v_clamped * 12.92
        } else {
            1.055 * v_clamped.powf(1.0 / 2.4) - 0.055
        };
        (v_encoded * 255.0 + 0.5) as u8
    }

    fn to_linear_scalar(&self, v: u8) -> Float {
        // Normalize 0-255 to 0.0-1.0 first
        let v_float = v as Float / 255.0;
        // Apply sRGB -> Linear math
        if v_float <= 0.04045 {
            v_float / 12.92
        } else {
            ((v_float + 0.055) / 1.055).powf(2.4)
        }
    }

    fn to_float_linear(&self, v: Float) -> Float {
        let v = v.clamp(0.0, 1.0);
        if v <= 0.04045 {
            v / 12.92
        } else {
            ((v + 0.055) / 1.055).powf(2.4)
        }
    }
}

impl fmt::Display for SRGBEncoding {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "sRGB")
    }
}

pub const LINEAR: ColorEncoding = ColorEncoding::Linear(LinearEncoding);
pub const SRGB: ColorEncoding = ColorEncoding::SRGB(SRGBEncoding);

const SRGB_TO_LINEAR_LUT: [Float; 256] = [
    0.0000000000,
    0.0003035270,
    0.0006070540,
    0.0009105810,
    0.0012141080,
    0.0015176350,
    0.0018211619,
    0.0021246888,
    0.0024282159,
    0.0027317430,
    0.0030352699,
    0.0033465356,
    0.0036765069,
    0.0040247170,
    0.0043914421,
    0.0047769533,
    0.0051815170,
    0.0056053917,
    0.0060488326,
    0.0065120910,
    0.0069954102,
    0.0074990317,
    0.0080231922,
    0.0085681248,
    0.0091340570,
    0.0097212177,
    0.0103298230,
    0.0109600937,
    0.0116122449,
    0.0122864870,
    0.0129830306,
    0.0137020806,
    0.0144438436,
    0.0152085144,
    0.0159962922,
    0.0168073755,
    0.0176419523,
    0.0185002182,
    0.0193823613,
    0.0202885624,
    0.0212190095,
    0.0221738834,
    0.0231533647,
    0.0241576303,
    0.0251868572,
    0.0262412224,
    0.0273208916,
    0.0284260381,
    0.0295568332,
    0.0307134409,
    0.0318960287,
    0.0331047624,
    0.0343398079,
    0.0356013142,
    0.0368894450,
    0.0382043645,
    0.0395462364,
    0.0409151986,
    0.0423114114,
    0.0437350273,
    0.0451862030,
    0.0466650836,
    0.0481718220,
    0.0497065634,
    0.0512694679,
    0.0528606549,
    0.0544802807,
    0.0561284944,
    0.0578054339,
    0.0595112406,
    0.0612460710,
    0.0630100295,
    0.0648032799,
    0.0666259527,
    0.0684781820,
    0.0703601092,
    0.0722718611,
    0.0742135793,
    0.0761853904,
    0.0781874284,
    0.0802198276,
    0.0822827145,
    0.0843762159,
    0.0865004659,
    0.0886556059,
    0.0908417329,
    0.0930589810,
    0.0953074843,
    0.0975873619,
    0.0998987406,
    0.1022417471,
    0.1046164930,
    0.1070231125,
    0.1094617173,
    0.1119324341,
    0.1144353822,
    0.1169706732,
    0.1195384338,
    0.1221387982,
    0.1247718409,
    0.1274376959,
    0.1301364899,
    0.1328683347,
    0.1356333494,
    0.1384316236,
    0.1412633061,
    0.1441284865,
    0.1470272839,
    0.1499598026,
    0.1529261619,
    0.1559264660,
    0.1589608639,
    0.1620294005,
    0.1651322246,
    0.1682693958,
    0.1714410931,
    0.1746473908,
    0.1778884083,
    0.1811642349,
    0.1844749898,
    0.1878207624,
    0.1912016720,
    0.1946178079,
    0.1980693042,
    0.2015562356,
    0.2050787061,
    0.2086368501,
    0.2122307271,
    0.2158605307,
    0.2195262313,
    0.2232279778,
    0.2269658893,
    0.2307400703,
    0.2345506549,
    0.2383976579,
    0.2422811985,
    0.2462013960,
    0.2501583695,
    0.2541521788,
    0.2581829131,
    0.2622507215,
    0.2663556635,
    0.2704978585,
    0.2746773660,
    0.2788943350,
    0.2831487954,
    0.2874408960,
    0.2917706966,
    0.2961383164,
    0.3005438447,
    0.3049873710,
    0.3094689548,
    0.3139887452,
    0.3185468316,
    0.3231432438,
    0.3277781308,
    0.3324515820,
    0.3371636569,
    0.3419144452,
    0.3467040956,
    0.3515326977,
    0.3564002514,
    0.3613068759,
    0.3662526906,
    0.3712377846,
    0.3762622178,
    0.3813261092,
    0.3864295185,
    0.3915725648,
    0.3967553079,
    0.4019778669,
    0.4072403014,
    0.4125427008,
    0.4178851545,
    0.4232677519,
    0.4286905527,
    0.4341537058,
    0.4396572411,
    0.4452012479,
    0.4507858455,
    0.4564110637,
    0.4620770514,
    0.4677838385,
    0.4735315442,
    0.4793202281,
    0.4851499796,
    0.4910208881,
    0.4969330430,
    0.5028865933,
    0.5088814497,
    0.5149177909,
    0.5209956765,
    0.5271152258,
    0.5332764983,
    0.5394796133,
    0.5457245708,
    0.5520114899,
    0.5583404899,
    0.5647116303,
    0.5711249113,
    0.5775805116,
    0.5840784907,
    0.5906189084,
    0.5972018838,
    0.6038274169,
    0.6104956269,
    0.6172066331,
    0.6239604354,
    0.6307572126,
    0.6375969648,
    0.6444797516,
    0.6514056921,
    0.6583748460,
    0.6653873324,
    0.6724432111,
    0.6795425415,
    0.6866854429,
    0.6938719153,
    0.7011020184,
    0.7083759308,
    0.7156936526,
    0.7230552435,
    0.7304608822,
    0.7379105687,
    0.7454043627,
    0.7529423237,
    0.7605246305,
    0.7681512833,
    0.7758223414,
    0.7835379243,
    0.7912980318,
    0.7991028428,
    0.8069523573,
    0.8148466945,
    0.8227858543,
    0.8307699561,
    0.8387991190,
    0.8468732834,
    0.8549926877,
    0.8631572723,
    0.8713672161,
    0.8796223402,
    0.8879231811,
    0.8962693810,
    0.9046613574,
    0.9130986929,
    0.9215820432,
    0.9301108718,
    0.9386858940,
    0.9473065734,
    0.9559735060,
    0.9646862745,
    0.9734454751,
    0.9822505713,
    0.9911022186,
    1.0000000000,
];

pub const RES: u32 = 64;

#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Coeffs {
    pub c0: Float,
    pub c1: Float,
    pub c2: Float,
}

impl From<&[Float; 3]> for Coeffs {
    fn from(slice: &[Float; 3]) -> Coeffs {
        Coeffs {
            c0: slice[0],
            c1: slice[1],
            c2: slice[2],
        }
    }
}

impl Add for Coeffs {
    type Output = Self;
    #[inline(always)]
    fn add(self, rhs: Self) -> Self {
        Self {
            c0: self.c0 + rhs.c0,
            c1: self.c1 + rhs.c1,
            c2: self.c2 + rhs.c2,
        }
    }
}

impl Sub for Coeffs {
    type Output = Self;
    #[inline(always)]
    fn sub(self, rhs: Self) -> Self {
        Self {
            c0: self.c0 - rhs.c0,
            c1: self.c1 - rhs.c1,
            c2: self.c2 - rhs.c2,
        }
    }
}

impl Mul<Float> for Coeffs {
    type Output = Self;
    #[inline(always)]
    fn mul(self, s: Float) -> Self {
        Self {
            c0: self.c0 * s,
            c1: self.c1 * s,
            c2: self.c2 * s,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, PartialEq)]
pub struct RGBToSpectrumTable {
    pub z_nodes: GVec<Float>,
    pub coeffs: GVec<Coeffs>,
    pub n_nodes: u32,
}

unsafe impl Send for RGBToSpectrumTable {}
unsafe impl Sync for RGBToSpectrumTable {}

impl RGBToSpectrumTable {
    #[inline(always)]
    fn get_coeffs(&self, bucket: u32, z: u32, y: u32, x: u32) -> Coeffs {
        let offset = bucket * (RES * RES * RES) + z * (RES * RES) + y * (RES) + x;
        unsafe { *self.coeffs.as_ptr().add(offset as usize) }
    }

    pub fn evaluate(&self, rgb: RGB) -> RGBSigmoidPolynomial {
        let m = rgb.max_component_value();
        let min_val = rgb.min_component_value();
        if m - min_val < 1e-4 {
            let x: Float = clamp(rgb[0], 1e-4, 0.9999);
            let c2 = (0.5 - x) / (x * (1.0 - x)).sqrt();
            return RGBSigmoidPolynomial::new(0.0, 0.0, c2);
        }

        // Identify the primary bucket (c) based on the dominant axis
        let c_idx = if rgb[0] > rgb[1] {
            if rgb[0] > rgb[2] {
                0
            } else {
                2
            }
        } else if rgb[1] > rgb[2] {
            1
        } else {
            2
        };

        let z = m;
        let (x_val, y_val) = match c_idx {
            0 => (rgb[1], rgb[2]), // R is max -> G, B
            1 => (rgb[0], rgb[2]), // G is max -> R, B
            _ => (rgb[0], rgb[1]), // B is max -> R, G
        };

        let (coord_a, coord_b) = if x_val > y_val {
            (y_val, x_val)
        } else {
            (x_val, y_val)
        };
        let x = coord_a / z;
        let y = coord_b / z;

        let z_nodes = &self.z_nodes;
        let zi = find_interval(RES, |i| z_nodes[i as usize] < z) as usize;
        let dz = (z - z_nodes[zi]) / (z_nodes[zi + 1] - z_nodes[zi]);
        let x_float = x * (RES - 1) as Float;
        let xi = (x_float as u32).min(RES - 2);
        let dx = x_float - xi as Float;

        let y_float = y * (RES - 1) as Float;
        let yi = (y_float as u32).min(RES - 2);
        let dy = y_float - yi as Float;

        let c000 = self.get_coeffs(c_idx, zi as u32, yi, xi);
        let c001 = self.get_coeffs(c_idx, zi as u32, yi, xi + 1);
        let c010 = self.get_coeffs(c_idx, zi as u32, yi + 1, xi);
        let c011 = self.get_coeffs(c_idx, zi as u32, yi + 1, xi + 1);
        let c100 = self.get_coeffs(c_idx, zi as u32 + 1, yi, xi);
        let c101 = self.get_coeffs(c_idx, zi as u32 + 1, yi, xi + 1);
        let c110 = self.get_coeffs(c_idx, zi as u32 + 1, yi + 1, xi);
        let c111 = self.get_coeffs(c_idx, zi as u32 + 1, yi + 1, xi + 1);
        let c00 = lerp(dx, c000, c001);
        let c01 = lerp(dx, c010, c011);
        let c10 = lerp(dx, c100, c101);
        let c11 = lerp(dx, c110, c111);
        let c0 = lerp(dy, c00, c01);
        let c1 = lerp(dy, c10, c11);
        let c = lerp(dz, c0, c1);

        RGBSigmoidPolynomial {
            c0: c.c0,
            c1: c.c1,
            c2: c.c2,
        }
    }
}

const LMS_FROM_XYZ: SquareMatrix3f = SquareMatrix::new([
    [0.8951, 0.2664, -0.1614],
    [-0.7502, 1.7135, 0.0367],
    [0.0389, -0.0685, 1.0296],
]);

const XYZ_FROM_LMS: SquareMatrix3f = SquareMatrix::new([
    [0.986993, -0.147054, 0.159963],
    [0.432305, 0.51836, 0.0492912],
    [-0.00852866, 0.0400428, 0.968487],
]);

pub fn white_balance(src_white: Point2f, target_white: Point2f) -> SquareMatrix3f {
    let src_xyz = XYZ::from_xyy(src_white, None);
    let dst_xyz = XYZ::from_xyy(target_white, None);
    let src_lms = LMS_FROM_XYZ * src_xyz;
    let dst_lms = LMS_FROM_XYZ * dst_xyz;
    let lms_correct = SquareMatrix3f::diag(&[
        dst_lms[0] / src_lms[0],
        dst_lms[1] / src_lms[1],
        dst_lms[2] / src_lms[2],
    ]);
    XYZ_FROM_LMS * lms_correct * LMS_FROM_XYZ
}
