use crate::bxdfs::NormalizedFresnelBxDF;
use crate::core::bsdf::BSDF;
use crate::core::geometry::{Frame, Normal3f, Point2f, Point3f, Point3fi, Vector3f};
use crate::core::interaction::{InteractionBase, ShadingGeom, SurfaceInteraction};
use crate::core::shape::Shape;
use crate::spectra::{N_SPECTRUM_SAMPLES, SampledSpectrum};
use crate::utils::Ptr;
use crate::utils::math::{catmull_rom_weights, square};
use crate::utils::sampling::sample_catmull_rom_2d;
use crate::{Float, PI};
use enum_dispatch::enum_dispatch;
use std::sync::Arc;

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

#[derive(Clone, Debug)]
pub struct SubsurfaceInteraction {
    pub pi: Point3fi,
    pub n: Normal3f,
    pub ns: Normal3f,
    pub dpdu: Vector3f,
    pub dpdv: Vector3f,
    pub dpdus: Vector3f,
    pub 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: InteractionBase::new_minimal(ssi.pi, ssi.n),
            dpdu: ssi.dpdu,
            dpdv: ssi.dpdv,
            dndu: Normal3f::zero(),
            dndv: Normal3f::zero(),
            shading: ShadingGeom {
                n: ssi.ns,
                dpdu: ssi.dpdus,
                dpdv: ssi.dpdvs,
                dndu: Normal3f::zero(),
                dndv: Normal3f::zero(),
            },
            face_index: 0,
            area_light: Ptr::null(),
            material: Ptr::null(),
            dpdx: Vector3f::zero(),
            dpdy: Vector3f::zero(),
            dudx: 0.,
            dvdx: 0.,
            dudy: 0.,
            dvdy: 0.,
            shape: Ptr::from(&Shape::default()),
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct BSSRDFTable {
    pub n_rho: u32,
    pub n_radius: u32,
    pub rho_samples: Ptr<Float>,
    pub radius_samples: Ptr<Float>,
    pub profile: Ptr<Float>,
    pub rho_eff: Ptr<Float>,
    pub profile_cdf: Ptr<Float>,
}

impl BSSRDFTable {
    pub fn get_rho(&self) -> &[Float] {
        unsafe { core::slice::from_raw_parts(self.rho_samples.as_ref(), self.n_rho as usize) }
    }

    pub fn get_radius(&self) -> &[Float] {
        unsafe { core::slice::from_raw_parts(self.radius_samples.as_ref(), self.n_radius as usize) }
    }

    pub fn get_profile(&self) -> &[Float] {
        let n_profile = (self.n_rho * self.n_radius) as usize;
        unsafe { core::slice::from_raw_parts(self.profile.as_ref(), n_profile) }
    }

    pub fn get_cdf(&self) -> &[Float] {
        let n_profile = (self.n_rho * self.n_radius) as usize;
        unsafe { core::slice::from_raw_parts(self.profile_cdf.as_ref(), n_profile) }
    }

    pub fn eval_profile(&self, rho_index: u32, radius_index: u32) -> Float {
        debug_assert!(rho_index < self.n_rho);
        debug_assert!(radius_index < self.n_radius);
        let idx = (rho_index * self.n_radius + radius_index) as usize;
        unsafe { *self.profile.add(idx) }
    }
}

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

#[enum_dispatch]
pub trait BSSRDFTrait {
    fn sample_sp(&self, u1: Float, u2: Point2f) -> Option<BSSRDFProbeSegment>;
    fn probe_intersection_to_sample(
        &self,
        si: &SubsurfaceInteraction,
        bxdf: NormalizedFresnelBxDF,
    ) -> BSSRDFSample;
}

#[repr(C)]
#[enum_dispatch(BSSRDFTrait)]
#[derive(Debug, Copy, Clone)]
pub enum BSSRDF {
    Tabulated(TabulatedBSSRDF),
}

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

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

    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.);
        let rho_samples = self.table.get_rho();
        let radius_samples = self.table.get_radius();
        for i in 0..N_SPECTRUM_SAMPLES {
            let r_optical = r * self.sigma_t[i];
            let (rho_offset, rho_weights) = match catmull_rom_weights(rho_samples, self.rho[i]) {
                Some(res) => res,
                None => continue,
            };

            let (radius_offset, radius_weights) =
                match catmull_rom_weights(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 as u32, radius_offset + k as u32);
                    }
                }
            }
            if r_optical != 0. {
                sr /= 2. * PI * r_optical;
            }
            sr_spectrum[i] = sr;
        }

        sr_spectrum *= square(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 rho_samples = self.table.get_rho();
        let radius_samples = self.table.get_radius();
        let profile = self.table.get_profile();
        let cdf = self.table.get_cdf();

        let (ret, _, _) =
            sample_catmull_rom_2d(rho_samples, radius_samples, 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.);
        let rhoeff_samples = self.table.get_rho();
        let radius_samples = self.table.get_radius();
        for i in 0..N_SPECTRUM_SAMPLES {
            let r_optical = r * self.sigma_t[i];
            let (rho_offset, rho_weights) = match catmull_rom_weights(rhoeff_samples, self.rho[i]) {
                Some(res) => res,
                None => continue,
            };

            let (radius_offset, radius_weights) =
                match catmull_rom_weights(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 += rhoeff_samples[rho_offset as usize + 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 as u32, radius_offset + k as u32)
                                * 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 BSSRDFTrait for TabulatedBSSRDF {
    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,
        _bxdf: NormalizedFresnelBxDF,
    ) -> BSSRDFSample {
        todo!()
    }
}