use std::fmt;
use std::ops::{Add, AddAssign, Sub, SubAssign, Mul, MulAssign, Div, DivAssign, Neg, Index, IndexMut};
use once_cell::sync::Lazy;
use std::sync::Arc;

use crate::core::pbrt::Float;
use crate::core::cie;
use super::color::{RGB, XYZ, RGBSigmoidPolynomial};
use super::colorspace::RGBColorspace;

pub const CIE_Y_INTEGRAL: Float = 106.856895;
pub const N_SPECTRUM_SAMPLES: usize = 1200;
pub const LAMBDA_MIN: i32 = 360; 
pub const LAMBDA_MAX: i32 = 830;

#[derive(Debug, Copy, Clone)]
pub struct SampledSpectrum {
    values: [Float; N_SPECTRUM_SAMPLES],
}

impl Default for SampledSpectrum {
    fn default() -> Self { Self { values: [0.0; N_SPECTRUM_SAMPLES] } }
}

impl SampledSpectrum {
    pub fn new(c: Float) -> Self {
        Self {
            values: [c; N_SPECTRUM_SAMPLES],
        }
    }

    pub fn from_vector(v: Vec<Float>) -> Self {
        let mut s = Self::new(0.0);
        for i in 0..N_SPECTRUM_SAMPLES {
            s.values[i] = v[i];
        }
        s
    }

    pub fn has_nans(&self) -> bool {
        self.values.iter().any(|&v| v.is_nan())
    }

    pub fn min_component_value(&self) -> Float {
        self.values.iter().fold(Float::INFINITY, |a, &b| a.min(b))
    }

    pub fn max_component_value(&self) -> Float {
        self.values.iter().fold(Float::NEG_INFINITY, |a, &b| a.max(b))
    }

    pub fn average(&self) -> Float {
        self.values.iter().sum::<Float>() / (N_SPECTRUM_SAMPLES as Float)
    }

    pub fn safe_div(&self, rhs: SampledSpectrum) -> Self {
        let mut r = SampledSpectrum::new(0.0);
        for i in 0..N_SPECTRUM_SAMPLES { 
            r.values[i] = if rhs[i] != 0.0 {self.values[i] / rhs.values[i]} else { 0.0 }
        }
        r
    }

    pub fn to_xyz(&self, lambda: &SampledWavelengths) -> XYZ {
        let x = spectra::X.sample(lambda);
        let y = spectra::Y.sample(lambda);
        let z = spectra::Z.sample(lambda);
        let pdf = lambda.pdf();
        XYZ::new(
            (*self * x).safe_div(pdf).average(),
            (*self * y).safe_div(pdf).average(),
            (*self * z).safe_div(pdf).average(),
            ) / CIE_Y_INTEGRAL
    }

    pub fn to_rgb(&self, lambda: &SampledWavelengths, c: &RGBColorspace) -> RGB {
        let xyz = self.to_xyz(lambda);
        c.to_rgb(xyz)
    }
}


impl Index<usize> for SampledSpectrum {
    type Output = Float;
    fn index(&self, i: usize) -> &Self::Output { &self.values[i] }
}

impl IndexMut<usize> for SampledSpectrum {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output { &mut self.values[i] }
}

impl Add for SampledSpectrum {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] += rhs.values[i];
        }
        ret
    }
}

impl AddAssign for SampledSpectrum {
    fn add_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] += rhs.values[i];
        }
    }
}

impl Sub for SampledSpectrum {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] -= rhs.values[i];
        }
        ret
    }
}

impl SubAssign for SampledSpectrum {
    fn sub_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] -= rhs.values[i];
        }
    }
}

impl Sub<SampledSpectrum> for Float {
    type Output = SampledSpectrum;
    fn sub(self, rhs: SampledSpectrum) -> SampledSpectrum {
        let mut ret = SampledSpectrum::new(0.0);
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] = self - rhs.values[i];
        }
        ret
    }
}

impl Mul for SampledSpectrum {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] *= rhs.values[i];
        }
        ret
    }
}

impl MulAssign for SampledSpectrum {
    fn mul_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] *= rhs.values[i];
        }
    }
}

impl Mul<Float> for SampledSpectrum {
    type Output = Self;
    fn mul(self, rhs: Float) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] *= rhs;
        }
        ret
    }
}

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

impl MulAssign<Float> for SampledSpectrum {
    fn mul_assign(&mut self, rhs: Float) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] *= rhs;
        }
    }
}

impl DivAssign for SampledSpectrum {
    fn div_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            debug_assert_ne!(0.0, rhs.values[i]);
            self.values[i] /= rhs.values[i];
        }
    }
}

impl Div for SampledSpectrum {
    type Output = Self;
    fn div(self, rhs: Self) -> Self::Output {
        let mut ret = self;
        ret /= rhs;
        ret
    }
}
impl Div<Float> for SampledSpectrum {
    type Output = Self;

    fn div(self, rhs: Float) -> Self::Output {
        debug_assert_ne!(rhs, 0.0);
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] /= rhs;
        }
        ret
    }
}

impl DivAssign<Float> for SampledSpectrum {
    fn div_assign(&mut self, rhs: Float) {
        debug_assert_ne!(rhs, 0.0);
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] /= rhs;
        }
    }
}

impl Neg for SampledSpectrum {
    type Output = Self;

    fn neg(self) -> Self::Output {
        let mut ret = SampledSpectrum::new(0.0);
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] = -self.values[i];
        }
        ret
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct SampledWavelengths {
    pub lambda: [Float; N_SPECTRUM_SAMPLES],
    pub pdf: [Float; N_SPECTRUM_SAMPLES],
}

impl SampledWavelengths {
    pub fn pdf(&self) -> SampledSpectrum {
        SampledSpectrum::from_vector(self.pdf.to_vec())
    }

    pub fn secondary_terminated(&self) -> bool {
        for i in 1..N_SPECTRUM_SAMPLES {
            if self.pdf[i] != 0.0 {
                return false
            }
        }
        true
    }

    pub fn terminate_secondary(&mut self) {
        if !self.secondary_terminated() {
            for i in 1..N_SPECTRUM_SAMPLES {
                self.pdf[i] = 0.0;
            }
            self.pdf[0] /= N_SPECTRUM_SAMPLES as Float;
        }
    }

    pub fn sample_uniform(u: Float, lambda_min: Float, lambda_max: Float) -> Self {
        let mut lambda = [0.0; N_SPECTRUM_SAMPLES];
        lambda[0] = crate::core::pbrt::lerp(u, lambda_min, lambda_min);
        let delta = (lambda_max - lambda_min) / N_SPECTRUM_SAMPLES as Float;
        for i in 1..N_SPECTRUM_SAMPLES {
            lambda[i] = lambda[i - 1] + delta;
            if lambda[i] > lambda_max {
                lambda[i] = lambda_min + (lambda[i] - lambda_max);
            }
        }

        let mut pdf = [0.0; N_SPECTRUM_SAMPLES];
        for i in 0..N_SPECTRUM_SAMPLES {
            pdf[i] = 1.0 / (lambda_max - lambda_min);
        }

        Self { lambda, pdf }
    }

    pub fn sample_visible_wavelengths(u: Float) -> Float {
        538.0 - 138.888889 * Float::atanh(0.85691062 - 1.82750197 * u)
    }

    pub fn visible_wavelengths_pdf(lambda: Float) -> Float {
        if lambda < 360.0 || lambda > 830.0 { return 0.0; }
        0.0039398042 / (Float::cosh(0.0072 * (lambda - 538.0))).sqrt()
    }

    pub fn sample_visible(u: Float) -> Self {
        let mut lambda = [0.0; N_SPECTRUM_SAMPLES];
        let mut pdf = [0.0; N_SPECTRUM_SAMPLES];

        for i in 0..N_SPECTRUM_SAMPLES {
            let mut up = u + i as Float / N_SPECTRUM_SAMPLES as Float;
            if up > 1.0 { up -= 1.0; }
            lambda[i] = Self::sample_visible_wavelengths(up);
            pdf[i] = Self::visible_wavelengths_pdf(lambda[i]);
        }
        Self { lambda, pdf }
    }
}

impl Index<usize> for SampledWavelengths {
    type Output = Float;
    fn index(&self, i: usize) -> &Self::Output { &self.lambda[i] }
}

impl IndexMut<usize> for SampledWavelengths {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output { &mut self.lambda[i] }
}

#[derive(Debug, Clone)]
pub enum Spectrum {
    Constant(ConstantSpectrum),
    DenselySampled(DenselySampledSpectrum),
    PiecewiseLinear(PiecewiseLinearSpectrum),
    Blackbody(BlackbodySpectrum),
    RGBAlbedo(RGBAlbedoSpectrum),
}

impl Spectrum {
    pub fn sample_at(&self, lambda: Float) -> Float {
        match self {
            Spectrum::Constant(s) => s.sample_at(lambda),
            Spectrum::DenselySampled(s) => s.sample_at(lambda),
            Spectrum::PiecewiseLinear(s) => s.sample_at(lambda),
            Spectrum::Blackbody(s) => s.sample_at(lambda),
            Spectrum::RGBAlbedo(s) => s.sample_at(lambda),
        }
    }

    pub fn max_value(&self) -> Float {
        match self {
            Spectrum::Constant(_s) => 0.0,
            Spectrum::DenselySampled(_s) => 0.0,
            Spectrum::PiecewiseLinear(_s) => 0.0,
            Spectrum::Blackbody(_s) => 0.0,
            Spectrum::RGBAlbedo(s) => s.max_value(),
        }
    }

    pub fn sample(&self, wavelengths: &SampledWavelengths) -> SampledSpectrum {
        let mut s = SampledSpectrum::default();
        for i in 0..N_SPECTRUM_SAMPLES {
            s[i] = self.sample_at(wavelengths[i]);
        }
        s
    }

    pub fn to_xyz(&self) -> XYZ {
        XYZ::new(inner_product(&spectra::X, self), 
            inner_product(&spectra::Y, self),
            inner_product(&spectra::Z, self))/CIE_Y_INTEGRAL
    } 
}

pub fn inner_product(f: &Spectrum, g: &Spectrum) -> Float {
    let mut integral = 0.0;
    for lambda in LAMBDA_MIN..=LAMBDA_MAX {
        integral += f.sample_at(lambda as f32) * g.sample_at(lambda as f32);
    }
    integral
}


#[derive(Debug, Clone, Copy)]
pub struct ConstantSpectrum {
    c: Float,
}

impl ConstantSpectrum {
    pub fn new(c: Float) -> Self { Self { c } }

    pub fn sample_at(&self, _lambda: Float) -> Float {
        self.c
    }

    fn max_value(&self) -> Float { self.c }
}

#[derive(Debug, Clone, Copy)]
pub struct RGBAlbedoSpectrum {
    rsp: RGBSigmoidPolynomial,
}

impl RGBAlbedoSpectrum {
    pub fn new(cs: &RGBColorspace, rgb: RGB) -> Self {
        Self { rsp: cs.to_rgb_coeffs(rgb) }
    }

    pub fn sample_at(&self, lambda: Float) -> Float {
        self.rsp.evaluate(lambda)
    }

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

    fn sample(&self, lambda: &SampledWavelengths) -> SampledSpectrum {
        let mut s = SampledSpectrum::default();
        for i in 0..N_SPECTRUM_SAMPLES {
            s[i] = self.rsp.evaluate(lambda[i]);
        }
        s
    }
}

#[derive(Debug, Clone, Copy)]
pub struct UnboundedRGBSpectrum {
    scale: Float,
    rsp: RGBSigmoidPolynomial,
}

impl UnboundedRGBSpectrum {
    pub fn new(cs: RGBColorspace, rgb: RGB) -> Self {
        let m = rgb.r.max(rgb.g).max(rgb.b);
        let scale = 2.0 * m;
        let scaled_rgb = if scale != 0.0 { rgb / scale } else { RGB::new(0.0, 0.0, 0.0) };
        Self { scale, rsp: cs.to_rgb_coeffs(scaled_rgb) }
    }
    pub fn sample_at(&self, lambda: Float) -> Float {
        self.scale * self.rsp.evaluate(lambda)
    }

    pub fn max_value(&self) -> Float {
        self.scale * self.rsp.max_value()
    }
}

#[derive(Debug, Clone, Default)]
pub struct RGBIlluminantSpectrum {
    scale: Float,
    rsp: RGBSigmoidPolynomial,
    illuminant: Option<Arc<DenselySampledSpectrum>>,
}

impl RGBIlluminantSpectrum {
    pub fn sample_at(&self, lambda: Float) -> Float {
        match &self.illuminant {
            Some(illuminant) => self.scale * self.rsp.evaluate(lambda) * illuminant.sample_at(lambda),
            None => 0.0,
        }
    }

    pub fn max_value(&self) -> Float {
        match &self.illuminant {
            Some(illuminant) => self.scale * self.rsp.max_value() * illuminant.max_value(),
            None => 0.0,
        }
    }
}


#[derive(Debug, Clone)]
pub struct DenselySampledSpectrum {
    lambda_min: i32,
    lambda_max: i32,
    values: Vec<Float>,
}

impl DenselySampledSpectrum {
    pub fn new(lambda_min: i32, lambda_max: i32) -> Self {
        let n_values = (lambda_max - lambda_min + 1).max(0) as usize;
        Self { lambda_min, lambda_max, values: vec![0.0; n_values] }
    }

    pub fn from_spectrum(spec: &Spectrum, lambda_min: i32, lambda_max: i32) -> Self {
        let mut s = Self::new(lambda_min, lambda_max);
        if s.values.is_empty() { return s; }
        for lambda in lambda_min..=lambda_max {
            let index = (lambda - lambda_min) as usize;
            s.values[index] = spec.sample_at(lambda as Float);
        }
        s
    }

    pub fn from_function<F>(f: F, lambda_min: i32, lambda_max: i32) -> Self where F: Fn(Float) -> Float {
        let mut s = Self::new(lambda_min, lambda_max);
        if s.values.is_empty() { return s; }
        for lambda in lambda_min..=lambda_max {
            let index = (lambda - lambda_min) as usize;
            s.values[index] = f(lambda as Float);
        }
        s
    }

    pub fn sample_at(&self, lambda: Float) -> Float {
        let offset = (lambda.round() as i32) - self.lambda_min;
        if offset < 0 || offset as usize >= self.values.len() {
            0.0
        } else {
            self.values[offset as usize]
        }
    }

    pub fn max_value(&self) -> Float {
        self.values.iter().fold(Float::MIN, |a,b| a.max(*b))
    }
}

#[derive(Debug, Clone)]
pub struct PiecewiseLinearSpectrum {
    samples: Vec<(Float, Float)>,
}

impl PiecewiseLinearSpectrum {
    pub fn from_interleaved(data: &[Float]) -> Self {
        assert!(data.len() % 2 == 0, "Interleaved data must have an even number of elements");
        let mut samples = Vec::new();
        for pair in data.chunks(2) {
            samples.push((pair[0], pair[1]));
        }
        // PBRT requires the samples to be sorted by wavelength for interpolation.
        samples.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        Self { samples }
    }

    fn sample_at(&self, lambda: Float) -> Float {
        if self.samples.is_empty() {
            return 0.0;
        }
        // Handle boundary conditions
        if lambda <= self.samples[0].0 {
            return self.samples[0].1;
        }
        if lambda >= self.samples.last().unwrap().0 {
            return self.samples.last().unwrap().1;
        }

        let i = self.samples.partition_point(|s| s.0 < lambda);
        let s1 = self.samples[i - 1];
        let s2 = self.samples[i];
        
        let t = (lambda - s1.0) / (s2.0 - s1.0);
        (1.0 - t) * s1.1 + t * s2.1
    }
}

#[derive(Debug, Clone, Copy)]
pub struct BlackbodySpectrum {
    temperature: Float,
    normalization_factor: Float,
}

// Planck's Law
impl BlackbodySpectrum {
    const C: Float = 299792458.0;
    const H: Float = 6.62606957e-34;
    const KB: Float = 1.3806488e-23;

    pub fn new(temperature: Float) -> Self {
        let lambda_max = 2.8977721e-3 / temperature * 1e9;
        let max_val = Self::planck_law(lambda_max, temperature);
        Self {
            temperature,
            normalization_factor: if max_val > 0.0 { 1.0 / max_val } else { 0.0 },
        }
    }
    
    fn planck_law(lambda_nm: Float, temp: Float) -> Float {
        if temp <= 0.0 { return 0.0; }
        let lambda_m = lambda_nm * 1e-9; // Convert nm to meters
        let c1 = 2.0 * Self::H * Self::C * Self::C;
        let c2 = (Self::H * Self::C) / Self::KB;
        
        let numerator = c1 / lambda_m.powi(5);
        let denominator = (c2 / (lambda_m * temp)).exp() - 1.0;
        
        if denominator.is_infinite() { 0.0 } else { numerator / denominator }
    }

    fn sample_at(&self, lambda: Float) -> Float {
        Self::planck_law(lambda, self.temperature) * self.normalization_factor
    }
}

#[derive(Debug, Clone, Copy)]
pub struct RGBSpectrum { pub c: [Float; 3] }

#[derive(Debug, Clone, Copy)]
pub struct RGBUnboundedSpectrum(pub RGBSpectrum);

pub mod spectra {
    use super::*;
    pub static X: Lazy<Spectrum> = Lazy::new(|| {
        let pls = PiecewiseLinearSpectrum::from_interleaved(&cie::CIE_X);
        let dss = DenselySampledSpectrum::from_spectrum(
            &Spectrum::PiecewiseLinear(pls),
            LAMBDA_MIN,
            LAMBDA_MAX,
        );
        Spectrum::DenselySampled(dss)
    });
    pub static Y: Lazy<Spectrum> = Lazy::new(|| {
        let pls = PiecewiseLinearSpectrum::from_interleaved(&cie::CIE_Y);
        let dss = DenselySampledSpectrum::from_spectrum(
            &Spectrum::PiecewiseLinear(pls),
            LAMBDA_MIN,
            LAMBDA_MAX,
        );
        Spectrum::DenselySampled(dss)
    });
    pub static Z: Lazy<Spectrum> = Lazy::new(|| {
        let pls = PiecewiseLinearSpectrum::from_interleaved(&cie::CIE_Z);
        let dss = DenselySampledSpectrum::from_spectrum(
            &Spectrum::PiecewiseLinear(pls),
            LAMBDA_MIN,
            LAMBDA_MAX,
        );
        Spectrum::DenselySampled(dss)
    });

 }