use crate::core::bxdf::BSDF;
use crate::core::interaction::{InteractionData, ShadingGeometry, SurfaceInteraction};
use crate::core::pbrt::{Float, PI};
use crate::geometry::{Frame, Normal3f, Point2f, Point3f, Point3fi, Vector3f};
use crate::shapes::Shape;
use crate::spectra::{N_SPECTRUM_SAMPLES, SampledSpectrum};
use crate::utils::math::{catmull_rom_weights, square};
use crate::utils::sampling::sample_catmull_rom_2d;
use enum_dispatch::enum_dispatch;
use std::sync::Arc;

#[derive(Debug)]
pub struct BSSRDFSample<'a> {
    pub sp: SampledSpectrum,
    pub pdf: SampledSpectrum,
    pub sw: BSDF<'a>,
    pub wo: Vector3f,
}

#[derive(Clone, Debug)]
pub struct SubsurfaceInteraction {
    pi: Point3fi,
    n: Normal3f,
    ns: Normal3f,
    dpdu: Vector3f,
    dpdv: Vector3f,
    dpdus: Vector3f,
    dpdvs: Vector3f,
}

impl SubsurfaceInteraction {
    pub fn new(si: &SurfaceInteraction) -> Self {
        Self {
            pi: si.common.pi,
            n: si.common.n,
            dpdu: si.dpdu,
            dpdv: si.dpdv,
            ns: si.shading.n,
            dpdus: si.shading.dpdu,
            dpdvs: si.shading.dpdv,
        }
    }

    pub fn p(&self) -> Point3f {
        self.pi.into()
    }
}

impl From<SurfaceInteraction> for SubsurfaceInteraction {
    fn from(si: SurfaceInteraction) -> SubsurfaceInteraction {
        SubsurfaceInteraction {
            pi: si.common.pi,
            n: si.common.n,
            ns: si.shading.n,
            dpdu: si.dpdu,
            dpdv: si.dpdv,
            dpdus: si.shading.dpdu,
            dpdvs: si.shading.dpdv,
        }
    }
}

impl From<&SubsurfaceInteraction> for SurfaceInteraction {
    fn from(ssi: &SubsurfaceInteraction) -> SurfaceInteraction {
        SurfaceInteraction {
            common: InteractionData {
                pi: ssi.pi,
                n: ssi.n,
                wo: Vector3f::zero(),
                time: 0.,
                medium_interface: None,
                medium: None,
            },
            uv: Point2f::zero(),
            dpdu: ssi.dpdu,
            dpdv: ssi.dpdv,
            dndu: Normal3f::zero(),
            dndv: Normal3f::zero(),
            shading: ShadingGeometry {
                n: ssi.ns,
                dpdu: ssi.dpdus,
                dpdv: ssi.dpdvs,
                dndu: Normal3f::zero(),
                dndv: Normal3f::zero(),
            },
            face_index: 0,
            area_light: None,
            material: None,
            dpdx: Vector3f::zero(),
            dpdy: Vector3f::zero(),
            dudx: 0.,
            dvdx: 0.,
            dudy: 0.,
            dvdy: 0.,
            shape: Shape::default().into(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct BSSRDFTable {
    rho_samples: Vec<Float>,
    radius_samples: Vec<Float>,
    profile: Vec<Float>,
    rho_eff: Vec<Float>,
    profile_cdf: Vec<Float>,
}

impl BSSRDFTable {
    pub fn new(n_rho_samples: usize, n_radius_samples: usize) -> Self {
        let rho_samples: Vec<Float> = Vec::with_capacity(n_rho_samples);
        let radius_samples: Vec<Float> = Vec::with_capacity(n_radius_samples);
        let profile: Vec<Float> = Vec::with_capacity(n_radius_samples * n_rho_samples);
        let rho_eff: Vec<Float> = Vec::with_capacity(n_rho_samples);
        let profile_cdf: Vec<Float> = Vec::with_capacity(n_radius_samples * n_rho_samples);
        Self {
            rho_samples,
            radius_samples,
            profile,
            rho_eff,
            profile_cdf,
        }
    }

    pub fn eval_profile(&self, rho_index: usize, radius_index: usize) -> Float {
        assert!(rho_index < self.rho_samples.len());
        assert!(radius_index < self.radius_samples.len());
        self.profile[rho_index * self.radius_samples.len() + radius_index]
    }
}

#[derive(Clone, Default, Debug)]
pub struct BSSRDFProbeSegment {
    pub p0: Point3f,
    pub p1: Point3f,
}

#[enum_dispatch]
pub trait BSSRDFTrait: Send + Sync + std::fmt::Debug {
    fn sample_sp(&self, u1: Float, u2: Point2f) -> Option<BSSRDFProbeSegment>;
    fn probe_intersection_to_sample(&self, si: &SubsurfaceInteraction) -> BSSRDFSample<'_>;
}

#[enum_dispatch(BSSRDFTrait)]
#[derive(Debug, Clone)]
pub enum BSSRDF<'a> {
    Tabulated(TabulatedBSSRDF<'a>),
}

#[derive(Clone, Debug)]
pub struct TabulatedBSSRDF<'a> {
    po: Point3f,
    wo: Vector3f,
    ns: Normal3f,
    eta: Float,
    sigma_t: SampledSpectrum,
    rho: SampledSpectrum,
    table: &'a BSSRDFTable,
}

impl<'a> TabulatedBSSRDF<'a> {
    pub fn new(
        po: Point3f,
        wo: Vector3f,
        ns: Normal3f,
        eta: Float,
        sigma_a: &SampledSpectrum,
        sigma_s: &SampledSpectrum,
        table: &'a BSSRDFTable,
    ) -> Self {
        let sigma_t = *sigma_a + *sigma_s;
        let rho = SampledSpectrum::safe_div(sigma_s, &sigma_t);
        Self {
            po,
            wo,
            ns,
            eta,
            table,
            sigma_t,
            rho,
        }
    }

    pub fn sp(&self, pi: Point3f) -> SampledSpectrum {
        self.sr(self.po.distance(pi))
    }

    pub fn sr(&self, r: Float) -> SampledSpectrum {
        let mut sr_spectrum = SampledSpectrum::new(0.);
        for i in 0..N_SPECTRUM_SAMPLES {
            let r_optical = r * self.sigma_t[i];
            let (rho_offset, rho_weights) =
                match catmull_rom_weights(&self.table.rho_samples, self.rho[i]) {
                    Some(res) => res,
                    None => continue,
                };

            let (radius_offset, radius_weights) =
                match catmull_rom_weights(&self.table.radius_samples, r_optical) {
                    Some(res) => res,
                    None => continue,
                };

            let mut sr = 0.;
            for (j, rho_weight) in rho_weights.iter().enumerate() {
                for (k, radius_weight) in radius_weights.iter().enumerate() {
                    let weight = rho_weight * radius_weight;
                    if weight != 0. {
                        sr += weight * self.table.eval_profile(rho_offset + j, radius_offset + k);
                    }
                }
            }
            if r_optical != 0. {
                sr /= 2. * PI * r_optical;
            }
            sr_spectrum[i] = sr;
        }

        sr_spectrum *= self.sigma_t * self.sigma_t;
        SampledSpectrum::clamp_zero(&sr_spectrum)
    }

    pub fn sample_sr(&self, u: Float) -> Option<Float> {
        if self.sigma_t[0] == 0. {
            return None;
        }
        let (ret, _, _) = sample_catmull_rom_2d(
            &self.table.rho_samples,
            &self.table.radius_samples,
            &self.table.profile,
            &self.table.profile_cdf,
            self.rho[0],
            u,
        );
        Some(ret / self.sigma_t[0])
    }

    pub fn pdf_sr(&self, r: Float) -> SampledSpectrum {
        let mut pdf = SampledSpectrum::new(0.);
        for i in 0..N_SPECTRUM_SAMPLES {
            let r_optical = r * self.sigma_t[i];
            let (rho_offset, rho_weights) =
                match catmull_rom_weights(&self.table.rho_samples, self.rho[i]) {
                    Some(res) => res,
                    None => continue,
                };

            let (radius_offset, radius_weights) =
                match catmull_rom_weights(&self.table.radius_samples, r_optical) {
                    Some(res) => res,
                    None => continue,
                };

            let mut sr = 0.;
            let mut rho_eff = 0.;
            for (j, rho_weight) in rho_weights.iter().enumerate() {
                if *rho_weight != 0. {
                    // Update _rhoEff_ and _sr_ for wavelength
                    rho_eff += self.table.rho_eff[rho_offset + j] * rho_weight;

                    // Fix: Use .iter().enumerate() for 'k'
                    for (k, radius_weight) in radius_weights.iter().enumerate() {
                        if *radius_weight != 0. {
                            sr += self.table.eval_profile(rho_offset + j, radius_offset + k)
                                * rho_weight
                                * radius_weight;
                        }
                    }
                }
            }

            // Cancel marginal PDF factor from tabulated BSSRDF profile
            if r_optical != 0. {
                sr /= 2. * PI * r_optical;
            }
            pdf[i] = sr * square(self.sigma_t[i]) / rho_eff;
        }

        SampledSpectrum::clamp_zero(&pdf)
    }

    pub fn pdf_sp(&self, pi: Point3f, ni: Normal3f) -> SampledSpectrum {
        let d = pi - self.po;
        let f = Frame::from_z(self.ns.into());
        let d_local = f.to_local(d);
        let n_local = f.to_local(ni.into());
        let r_proj = [
            (square(d_local.y() + square(d_local.z()))).sqrt(),
            (square(d_local.z() + square(d_local.x()))).sqrt(),
            (square(d_local.x() + square(d_local.y()))).sqrt(),
        ];

        let axis_prob = [0.25, 0.25, 0.25];
        let mut pdf = SampledSpectrum::new(0.);
        for axis in 0..3 {
            pdf += self.pdf_sr(r_proj[axis] * n_local[axis].abs() * axis_prob[axis]);
        }
        pdf
    }
}

impl<'a> BSSRDFTrait for TabulatedBSSRDF<'a> {
    fn sample_sp(&self, u1: Float, u2: Point2f) -> Option<BSSRDFProbeSegment> {
        let f = if u1 < 0.25 {
            Frame::from_x(self.ns.into())
        } else if u1 < 0.5 {
            Frame::from_y(self.ns.into())
        } else {
            Frame::from_z(self.ns.into())
        };

        let r = self.sample_sr(u2[0])?;
        let phi = 2. * PI * u2[1];
        let r_max = self.sample_sr(0.999)?;
        let l = 2. * (square(r_max) - square(r)).sqrt();
        let p_start = self.po + r * (f.x * phi.cos() + f.y * phi.sin()) - l * f.z / 2.;
        let p_target = p_start + l * f.z;
        Some(BSSRDFProbeSegment {
            p0: p_start,
            p1: p_target,
        })
    }

    fn probe_intersection_to_sample(&self, _si: &SubsurfaceInteraction) -> BSSRDFSample<'_> {
        todo!()
    }
}