use std::any::TypeId;
use std::fmt;
use std::ops::{
    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};

use super::spectrum::Spectrum;
use crate::core::pbrt::{Float, lerp};
use crate::geometry::Point2f;
use crate::utils::math::SquareMatrix;
use once_cell::sync::Lazy;

pub trait Triplet {
    fn from_triplet(c1: Float, c2: Float, c3: Float) -> Self;
}

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

impl Triplet for XYZ {
    fn from_triplet(c1: Float, c2: Float, c3: Float) -> Self {
        XYZ::new(c1, c2, c3)
    }
}

impl<'a> IntoIterator for &'a XYZ {
    type Item = &'a Float;
    type IntoIter = std::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<usize> for XYZ {
    type Output = Float;
    fn index(&self, index: usize) -> &Self::Output {
        debug_assert!(index < 3);
        match index {
            0 => &self.x,
            1 => &self.y,
            _ => &self.z,
        }
    }
}

impl IndexMut<usize> for XYZ {
    fn index_mut(&mut self, index: usize) -> &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)
    }
}

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

impl Triplet for RGB {
    fn from_triplet(c1: Float, c2: Float, c3: Float) -> Self {
        RGB::new(c1, c2, c3)
    }
}

impl<'a> IntoIterator for &'a RGB {
    type Item = &'a Float;
    type IntoIter = std::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(&self) -> Float {
        self.r.max(self.g).max(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<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)
    }
}

pub const RES: usize = 64;
pub type CoefficientArray = [[[[[Float; 3]; RES]; RES]; RES]; 3];

#[derive(Debug)]
pub struct RGBToSpectrumTable {
    z_nodes: &'static [f32],
    coeffs: &'static CoefficientArray,
}

impl RGBToSpectrumTable {
    pub fn srgb() -> Self {
        // use crate::core::constants::{RGB_TO_SPECTRUM_Z_NODES, RGB_TO_SPECTRUM_COEFFS};
        // Self::new(&RGB_TO_SPECTRUM_Z_NODES, &RGB_TO_SPECTRUM_COEFFS)
        todo!("Link the static constant arrays for sRGB coefficients here")
    }
}

#[derive(Debug, Default, Copy, Clone)]
pub struct RGBSigmoidPolynomial {
    c0: Float,
    c1: Float,
    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 =
            match crate::core::pbrt::evaluate_polynomial(lambda, &[self.c0, self.c1, self.c2]) {
                Some(value) => value,
                None => {
                    panic!("evaluate_polynomial returned None with non-empty coefficients")
                }
            };

        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 lambda >= 360.0 && lambda <= 830.0 {
            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())
    }
}

impl RGBToSpectrumTable {
    pub fn new(z_nodes: &'static [f32], coeffs: &'static CoefficientArray) -> Self {
        Self { z_nodes, coeffs }
    }

    pub fn to_polynomial(&self, rgb: RGB) -> RGBSigmoidPolynomial {
        if rgb[0] == rgb[1] && rgb[1] == rgb[2] {
            return RGBSigmoidPolynomial::new(
                0.0,
                0.0,
                (rgb[0] - 0.5) / (rgb[0] * (1.0 - rgb[0])).sqrt(),
            );
        }
        let maxc;
        if rgb[0] > rgb[1] {
            if rgb[0] > rgb[2] {
                maxc = 0;
            } else {
                maxc = 2;
            }
        } else {
            if rgb[1] > rgb[2] {
                maxc = 1;
            } else {
                maxc = 2;
            }
        }

        let z = rgb[maxc];
        let x = rgb[(maxc + 1) % 3] * (RES - 1) as Float / z;
        let y = rgb[(maxc + 2) % 3] * (RES - 1) as Float / z;

        let xi = x.min(RES as Float - 2.0);
        let yi = y.min(RES as Float - 2.0);
        let zi = crate::core::pbrt::find_interval(RES, |i: usize| self.z_nodes[i] < z);
        let dx = (x - xi) as usize;
        let dy = (y - yi) as usize;
        let dz = (z - self.z_nodes[zi]) / (self.z_nodes[zi + 1] - self.z_nodes[zi]);
        let mut c = [0.0; 3];
        for i in 0..3 {
            let co = |dx: usize, dy: usize, dz: usize| {
                self.coeffs[maxc][zi as usize + dz][yi as usize + dy][xi as usize + dx][i]
            };
            c[i] = lerp(
                dz,
                lerp(
                    dy as Float,
                    lerp(dx as Float, co(0, 0, 0) as Float, co(1, 0, 0)) as Float,
                    lerp(dx as Float, co(0, 1, 0) as Float, co(1, 1, 0) as Float),
                ),
                lerp(
                    dy as Float,
                    lerp(dx as Float, co(0, 0, 1) as Float, co(1, 0, 1)) as Float,
                    lerp(dx as Float, co(0, 1, 1) as Float, co(1, 1, 1) as Float),
                ),
            );
        }
        RGBSigmoidPolynomial {
            c0: c[0],
            c1: c[1],
            c2: c[2],
        }
    }
}

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

const XYZ_FROM_LMS: SquareMatrix<Float, 3> = 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) -> SquareMatrix<Float, 3> {
    // Find LMS coefficients for source and target white
    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;

    // Return white balancing matrix for source and target white
    let lms_correct = SquareMatrix::<Float, 3>::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
}

pub trait ColorEncoding: 'static + Send + Sync + fmt::Debug + fmt::Display {
    fn from_linear_slice(&self, vin: &[Float], vout: &mut [u8]);
    fn to_linear_slice(&self, vin: &[u8], vout: &mut [Float]);
    fn to_float_linear(&self, v: Float) -> Float;
    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

#[derive(Debug)]
pub struct LinearEncoding;
impl ColorEncoding for LinearEncoding {
    fn from_linear_slice(&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_slice(&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
    }
}

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

#[derive(Debug)]
pub struct SRGBEncoding;
impl ColorEncoding for SRGBEncoding {
    fn from_linear_slice(&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_slice(&self, vin: &[u8], vout: &mut [Float]) {
        for (i, &v) in vin.iter().enumerate() {
            vout[i] = SRGB_TO_LINEAR_LUT[v as usize];
        }
    }

    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 static LINEAR: Lazy<&'static dyn ColorEncoding> = Lazy::new(|| &LinearEncoding);
pub static SRGB: Lazy<&'static dyn ColorEncoding> = Lazy::new(|| &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,
];