use crate::core::geometry::{
    spherical_direction, Bounds3f, Frame, Point2f, Point3f, Point3i, Ray, Vector3f, VectorLike,
};
use crate::core::pbrt::{Float, INV_4_PI, PI};
use crate::core::spectrum::{Spectrum, SpectrumTrait};
use crate::spectra::{
    BlackbodySpectrum, DenselySampledSpectrum, RGBIlluminantSpectrum, RGBUnboundedSpectrum,
    SampledSpectrum, SampledWavelengths, LAMBDA_MAX, LAMBDA_MIN,
};
use crate::utils::containers::SampledGrid;
use crate::utils::math::{clamp, square};
use crate::utils::rng::Rng;
use crate::utils::transform::Transform;
use crate::{gvec_with_capacity, leak, GVec, Ptr};
use enum_dispatch::enum_dispatch;
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct PhaseFunctionSample {
    pub p: Float,
    pub wi: Vector3f,
    pub pdf: Float,
}

impl PhaseFunctionSample {
    pub fn new(p: Float, wi: Vector3f, pdf: Float) -> Self {
        Self { p, wi, pdf }
    }
}

#[enum_dispatch]
pub trait PhaseFunctionTrait {
    fn p(&self, wo: Vector3f, wi: Vector3f) -> Float;
    fn sample_p(&self, wo: Vector3f, u: Point2f) -> Option<PhaseFunctionSample>;
    fn pdf(&self, wo: Vector3f, wi: Vector3f) -> Float;
}

#[repr(C)]
#[enum_dispatch(PhaseFunctionTrait)]
#[derive(Debug, Clone, Copy)]
pub enum PhaseFunction {
    HenyeyGreenstein(HGPhaseFunction),
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct HGPhaseFunction {
    g: Float,
}

impl HGPhaseFunction {
    pub fn new(g: Float) -> Self {
        Self { g }
    }

    fn phase_hg(&self, cos_theta: Float) -> Float {
        let denom = 1.0 + square(self.g) + 2.0 * self.g * cos_theta;
        INV_4_PI * (1.0 - square(self.g)) / (denom * denom.sqrt())
    }
}

impl PhaseFunctionTrait for HGPhaseFunction {
    fn p(&self, wo: Vector3f, wi: Vector3f) -> Float {
        self.phase_hg(-(wo.dot(wi)))
    }

    fn sample_p(&self, wo: Vector3f, u: Point2f) -> Option<PhaseFunctionSample> {
        let cos_theta = if self.g.abs() < 1e-3 {
            1.0 - 2.0 * u[0]
        } else {
            let square_term = (1.0 - square(self.g)) / (1.0 + self.g * (1.0 - 2.0 * u[0]));
            -(1.0 + square(self.g) - square(square_term)) / (2.0 * self.g)
        };

        let sin_theta = (1.0 - square(cos_theta)).max(0.0).sqrt();
        let phi = 2.0 * PI * u[1];
        let w_frame = Frame::from_z(wo);
        let wi = w_frame.from_local(spherical_direction(sin_theta, cos_theta, phi));

        let p_val = self.phase_hg(cos_theta);

        Some(PhaseFunctionSample::new(p_val, wi, p_val))
    }

    fn pdf(&self, wo: Vector3f, wi: Vector3f) -> Float {
        self.p(wo, wi)
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct MajorantGrid {
    pub bounds: Bounds3f,
    pub res: Point3i,
    pub voxels: GVec<Float>,
    pub n_voxels: u32,
}


impl MajorantGrid {
    #[cfg(not(target_os = "cuda"))]
    pub fn new(bounds: Bounds3f, res: Point3i) -> Self {
        let n_voxels = (res.x() * res.y() * res.z()) as usize;
        let voxels = gvec_with_capacity(n_voxels);
        Self {
            bounds,
            res,
            voxels,
            n_voxels: n_voxels as u32,
        }
    }

    #[inline(always)]
    fn is_valid(&self) -> bool {
        !self.voxels.is_empty()
    }

    #[inline(always)]
    pub fn lookup(&self, x: i32, y: i32, z: i32) -> Float {
        if !self.is_valid() {
            return 0.0;
        }

        let idx = z * self.res.x() * self.res.y() + y * self.res.x() + x;

        if idx >= 0 && (idx as u32) < self.n_voxels {
            unsafe { *self.voxels.as_ptr().add(idx as usize) }
        } else {
            0.0
        }
    }

    #[inline(always)]
    pub fn set(&mut self, x: i32, y: i32, z: i32, v: Float) {
        if !self.is_valid() {
            return;
        }

        let idx = x + self.res.x() * (y + self.res.y() * z);

        unsafe {
            *self.voxels.as_mut_ptr().add(idx as usize) = v;
        }
    }

    pub fn voxel_bounds(&self, x: i32, y: i32, z: i32) -> Bounds3f {
        let p0 = Point3f::new(
            x as Float / self.res.x() as Float,
            y as Float / self.res.y() as Float,
            z as Float / self.res.z() as Float,
        );
        let p1 = Point3f::new(
            (x + 1) as Float / self.res.x() as Float,
            (y + 1) as Float / self.res.y() as Float,
            (z + 1) as Float / self.res.z() as Float,
        );
        Bounds3f::from_points(p0, p1)
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct RayMajorantSegment {
    t_min: Float,
    t_max: Float,
    sigma_maj: SampledSpectrum,
}

#[repr(C)]
#[derive(Clone, Debug)]
pub enum RayMajorantIterator {
    Homogeneous(HomogeneousMajorantIterator),
    DDA(DDAMajorantIterator),
    // Grid(GridMajorantIterator<'a>),
    Void,
}

impl Iterator for RayMajorantIterator {
    type Item = RayMajorantSegment;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            RayMajorantIterator::Homogeneous(iter) => iter.next(),
            RayMajorantIterator::DDA(iter) => iter.next(),
            RayMajorantIterator::Void => None,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug)]
pub struct HomogeneousMajorantIterator {
    called: bool,
    seg: RayMajorantSegment,
}

impl HomogeneousMajorantIterator {
    pub fn new(t_min: Float, t_max: Float, sigma_maj: SampledSpectrum) -> Self {
        Self {
            seg: RayMajorantSegment {
                t_min,
                t_max,
                sigma_maj,
            },
            called: false,
        }
    }
}

impl Iterator for HomogeneousMajorantIterator {
    type Item = RayMajorantSegment;

    fn next(&mut self) -> Option<Self::Item> {
        if self.called {
            return None;
        }

        self.called = true;
        Some(self.seg)
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct DDAMajorantIterator {
    sigma_t: SampledSpectrum,
    t_min: Float,
    t_max: Float,
    grid: Ptr<MajorantGrid>,
    next_crossing_t: [Float; 3],
    delta_t: [Float; 3],
    step: [i32; 3],
    voxel_limit: [i32; 3],
    voxel: [i32; 3],
}

impl DDAMajorantIterator {
    pub fn new(
        ray: &Ray,
        t_min: Float,
        t_max: Float,
        grid: &MajorantGrid,
        sigma_t: &SampledSpectrum,
    ) -> Self {
        let mut iter = Self {
            t_min,
            t_max,
            sigma_t: *sigma_t,
            grid: Ptr::from(&*grid),
            next_crossing_t: [0.0; 3],
            delta_t: [0.0; 3],
            step: [0; 3],
            voxel_limit: [0; 3],
            voxel: [0; 3],
        };

        let diag = grid.bounds.diagonal();

        let mut ray_grid_d = Vector3f::new(
            ray.d.x() / diag.x(),
            ray.d.y() / diag.y(),
            ray.d.z() / diag.z(),
        );

        let p_grid_start = grid.bounds.offset(&ray.at(t_min));
        let grid_intersect = Vector3f::from(p_grid_start);
        let res = [grid.res.x(), grid.res.y(), grid.res.z()];

        for axis in 0..3 {
            iter.voxel[axis] = clamp(
                (grid_intersect[axis] * res[axis] as Float) as i32,
                0,
                res[axis] - 1,
            );

            iter.delta_t[axis] = 1.0 / (ray_grid_d[axis].abs() * res[axis] as Float);

            if ray_grid_d[axis] == -0.0 {
                ray_grid_d[axis] = 0.0;
            }

            if ray_grid_d[axis] >= 0.0 {
                let next_voxel_pos = (iter.voxel[axis] + 1) as Float / res[axis] as Float;
                iter.next_crossing_t[axis] =
                    t_min + (next_voxel_pos - grid_intersect[axis]) / ray_grid_d[axis];
                iter.step[axis] = 1;
                iter.voxel_limit[axis] = res[axis];
            } else {
                let next_voxel_pos = (iter.voxel[axis]) as Float / res[axis] as Float;
                iter.next_crossing_t[axis] =
                    t_min + (next_voxel_pos - grid_intersect[axis]) / ray_grid_d[axis];
                iter.step[axis] = -1;
                iter.voxel_limit[axis] = -1;
            }
        }

        iter
    }
}

impl Iterator for DDAMajorantIterator {
    type Item = RayMajorantSegment;

    fn next(&mut self) -> Option<Self::Item> {
        if self.t_min >= self.t_max {
            return None;
        }

        // Find stepAxis for stepping to next voxel and exit point tVoxelExit
        let d0 = self.next_crossing_t[0];
        let d1 = self.next_crossing_t[1];
        let d2 = self.next_crossing_t[2];

        let bits = ((d0 < d1) as usize) << 2 | ((d0 < d2) as usize) << 1 | ((d1 < d2) as usize);

        const CMP_TO_AXIS: [usize; 8] = [2, 1, 2, 1, 2, 2, 0, 0];
        let step_axis = CMP_TO_AXIS[bits];

        let t_voxel_exit = self.t_max.min(self.next_crossing_t[step_axis]);
        // Get maxDensity for current voxel and initialize RayMajorantSegment, seg

        let density_scale = self
            .grid
            .lookup(self.voxel[0], self.voxel[1], self.voxel[2]);

        let sigma_maj = self.sigma_t * density_scale;

        let seg = RayMajorantSegment {
            t_min: self.t_min,
            t_max: t_voxel_exit,
            sigma_maj,
        };

        // Advance to next voxel in maximum density grid
        self.t_min = t_voxel_exit;

        if self.next_crossing_t[step_axis] > self.t_max {
            self.t_min = self.t_max;
        }

        self.voxel[step_axis] += self.step[step_axis];

        if self.voxel[step_axis] == self.voxel_limit[step_axis] {
            self.t_min = self.t_max;
        }

        // Increment the crossing time for this axis
        self.next_crossing_t[step_axis] += self.delta_t[step_axis];

        Some(seg)
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct MediumProperties {
    pub sigma_a: SampledSpectrum,
    pub sigma_s: SampledSpectrum,
    pub phase: PhaseFunction,
    pub le: SampledSpectrum,
}

impl MediumProperties {
    pub fn sigma_t(&self) -> SampledSpectrum {
        self.sigma_a * self.sigma_s
    }
}

#[enum_dispatch]
pub trait MediumTrait: Send + Sync + core::fmt::Debug {
    fn is_emissive(&self) -> bool;
    fn sample_point(&self, p: Point3f, lambda: &SampledWavelengths) -> MediumProperties;
    fn sample_ray(
        &self,
        ray: &Ray,
        t_max: Float,
        lambda: &SampledWavelengths,
    ) -> RayMajorantIterator;

    fn sample_t_maj<F>(
        &self,
        mut ray: Ray,
        mut t_max: Float,
        u: Float,
        rng: &mut Rng,
        lambda: &SampledWavelengths,
        mut callback: F,
    ) -> SampledSpectrum
    where
        F: FnMut(Point3f, MediumProperties, SampledSpectrum, SampledSpectrum, &mut Rng) -> bool,
    {
        let len = ray.d.norm();
        t_max *= len;
        ray.d /= len;

        let mut iter = self.sample_ray(&ray, t_max, lambda);
        let mut t_maj = SampledSpectrum::new(1.0);

        while let Some(seg) = iter.next() {
            if seg.sigma_maj[0] == 0. {
                let dt = seg.t_max - seg.t_min;
                let dt = if dt.is_infinite() { Float::MAX } else { dt };

                t_maj *= (-dt * seg.sigma_maj).exp();
                continue;
            }

            let mut t_min = seg.t_min;
            loop {
                let dist = -(1. - u.ln()) / seg.sigma_maj[0];
                let t = t_min + dist;

                if t < seg.t_max {
                    t_maj *= (-(t - t_min) * seg.sigma_maj).exp();
                    let p = ray.at(t);
                    let mp = self.sample_point(p, lambda);

                    if !callback(p, mp, seg.sigma_maj, t_maj, rng) {
                        return SampledSpectrum::new(1.);
                    }

                    t_maj = SampledSpectrum::new(1.);
                    t_min = t;
                } else {
                    let dt = seg.t_max - t_min;
                    let dt = if dt.is_infinite() { Float::MAX } else { dt };

                    t_maj *= (-dt * seg.sigma_maj).exp();
                    break;
                }
            }
        }

        SampledSpectrum::new(1.)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
#[enum_dispatch(MediumTrait)]
pub enum Medium {
    Homogeneous(HomogeneousMedium),
    Grid(GridMedium),
    RGBGrid(RGBGridMedium),
    Cloud(CloudMedium),
    NanoVDB(NanoVDBMedium),
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct HomogeneousMedium {
    pub sigma_a_spec: Ptr<DenselySampledSpectrum>,
    pub sigma_s_spec: Ptr<DenselySampledSpectrum>,
    pub le_spec: Ptr<DenselySampledSpectrum>,
    pub phase: HGPhaseFunction,
}

impl MediumTrait for HomogeneousMedium {
    fn is_emissive(&self) -> bool {
        self.le_spec.max_value() > 0.
    }

    fn sample_point(&self, _p: Point3f, lambda: &SampledWavelengths) -> MediumProperties {
        let sigma_a = self.sigma_a_spec.sample(lambda);
        let sigma_s = self.sigma_s_spec.sample(lambda);
        let le = self.le_spec.sample(lambda);
        MediumProperties {
            sigma_a,
            sigma_s,
            phase: self.phase.into(),
            le,
        }
    }

    fn sample_ray(
        &self,
        _ray: &Ray,
        t_max: Float,
        lambda: &SampledWavelengths,
    ) -> RayMajorantIterator {
        let sigma_a = self.sigma_a_spec.sample(lambda);
        let sigma_s = self.sigma_s_spec.sample(lambda);
        let sigma_maj = sigma_a + sigma_s;

        let iter = HomogeneousMajorantIterator::new(0.0, t_max, sigma_maj);

        RayMajorantIterator::Homogeneous(iter)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct GridMedium {
    pub bounds: Bounds3f,
    pub render_from_medium: Transform,
    pub sigma_a_spec: Ptr<DenselySampledSpectrum>,
    pub sigma_s_spec: Ptr<DenselySampledSpectrum>,
    pub density_grid: Ptr<SampledGrid<Float>>,
    pub phase: HGPhaseFunction,
    pub temperature_grid: Ptr<SampledGrid<Float>>,
    pub le_spec: Ptr<DenselySampledSpectrum>,
    pub le_scale: Ptr<SampledGrid<Float>>,
    pub is_emissive: bool,
    pub majorant_grid: Ptr<MajorantGrid>,
}

impl MediumTrait for GridMedium {
    fn is_emissive(&self) -> bool {
        self.is_emissive
    }

    fn sample_point(&self, p: Point3f, lambda: &SampledWavelengths) -> MediumProperties {
        let mut sigma_a = self.sigma_a_spec.sample(lambda);
        let mut sigma_s = self.sigma_s_spec.sample(lambda);
        let p_transform = self.render_from_medium.apply_inverse(p);
        let p = Point3f::from(self.bounds.offset(&p_transform));
        let d = self.density_grid.lookup(p);
        sigma_a *= d;
        sigma_s *= d;
        let scale = if self.is_emissive {
            self.le_scale.lookup(p)
        } else {
            0.0
        };

        let le = if scale > 0.0 {
            let raw_emission = if !self.temperature_grid.is_null() {
                let temp = self.temperature_grid.lookup(p);
                BlackbodySpectrum::new(temp).sample(lambda)
            } else {
                self.le_spec.sample(lambda)
            };

            raw_emission * scale
        } else {
            SampledSpectrum::new(0.0)
        };

        MediumProperties {
            sigma_a,
            sigma_s,
            phase: self.phase.into(),
            le,
        }
    }

    fn sample_ray(
        &self,
        ray: &Ray,
        t_max: Float,
        lambda: &SampledWavelengths,
    ) -> RayMajorantIterator {
        let (local_ray, local_t_max) = self.render_from_medium.apply_inverse_ray(ray, Some(t_max));

        let inv_dir = Vector3f::new(
            1.0 / local_ray.d.x(),
            1.0 / local_ray.d.y(),
            1.0 / local_ray.d.z(),
        );
        let dir_is_neg = [
            if local_ray.d.x() < 0.0 { 1 } else { 0 },
            if local_ray.d.y() < 0.0 { 1 } else { 0 },
            if local_ray.d.z() < 0.0 { 1 } else { 0 },
        ];

        let (t_min, t_max) =
            match self
                .bounds
                .intersect_p(local_ray.o, local_t_max, inv_dir, &dir_is_neg)
            {
                Some((t0, t1)) => (t0, t1),
                None => return RayMajorantIterator::Void, // Missed the medium bounds
            };

        let sigma_t = self.sigma_a_spec.sample(lambda) + self.sigma_s_spec.sample(lambda);

        let iter =
            DDAMajorantIterator::new(&local_ray, t_min, t_max, &self.majorant_grid, &sigma_t);

        RayMajorantIterator::DDA(iter)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct RGBGridMedium {
    pub bounds: Bounds3f,
    pub render_from_medium: Transform,
    pub phase: HGPhaseFunction,
    pub le_scale: Float,
    pub sigma_scale: Float,
    pub sigma_a_grid: Ptr<SampledGrid<RGBUnboundedSpectrum>>,
    pub sigma_s_grid: Ptr<SampledGrid<RGBUnboundedSpectrum>>,
    pub le_grid: Ptr<SampledGrid<RGBIlluminantSpectrum>>,
    pub majorant_grid: Ptr<MajorantGrid>,
}

impl MediumTrait for RGBGridMedium {
    fn is_emissive(&self) -> bool {
        self.le_grid.is_valid() && self.le_scale > 0.
    }

    fn sample_point(&self, p: Point3f, lambda: &SampledWavelengths) -> MediumProperties {
        let p_transform = self.render_from_medium.apply_inverse(p);
        let p = Point3f::from(self.bounds.offset(&p_transform));
        let convert = |s: &RGBUnboundedSpectrum| s.sample(lambda);
        let sigma_a = if self.sigma_a_grid.is_valid() {
            self.sigma_scale * self.sigma_a_grid.lookup_convert(p, convert)
        } else {
            SampledSpectrum::new(1.0) * self.sigma_scale // Or 0.0? Check PBRT logic
        };
        let sigma_s = if self.sigma_s_grid.is_valid() {
            self.sigma_scale * self.sigma_s_grid.lookup_convert(p, convert)
        } else {
            SampledSpectrum::new(1.0) * self.sigma_scale
        };

        let le = if self.le_grid.is_valid() && self.le_scale > 0.0 {
            self.le_grid.lookup_convert(p, |s| s.sample(lambda)) * self.le_scale
        } else {
            SampledSpectrum::new(0.0)
        };

        MediumProperties {
            sigma_a,
            sigma_s,
            phase: self.phase.into(),
            le,
        }
    }

    fn sample_ray(
        &self,
        ray: &Ray,
        t_max: Float,
        _lambda: &SampledWavelengths,
    ) -> RayMajorantIterator {
        let (local_ray, local_t_max) = self.render_from_medium.apply_inverse_ray(ray, Some(t_max));

        let inv_dir = Vector3f::new(
            1.0 / local_ray.d.x(),
            1.0 / local_ray.d.y(),
            1.0 / local_ray.d.z(),
        );
        let dir_is_neg = [
            if local_ray.d.x() < 0.0 { 1 } else { 0 },
            if local_ray.d.y() < 0.0 { 1 } else { 0 },
            if local_ray.d.z() < 0.0 { 1 } else { 0 },
        ];

        let (t_min, t_max) =
            match self
                .bounds
                .intersect_p(local_ray.o, local_t_max, inv_dir, &dir_is_neg)
            {
                Some((t0, t1)) => (t0, t1),
                None => return RayMajorantIterator::Void, // Missed the medium bounds
            };

        let sigma_t = SampledSpectrum::new(1.);

        let iter =
            DDAMajorantIterator::new(&local_ray, t_min, t_max, &self.majorant_grid, &sigma_t);
        RayMajorantIterator::DDA(iter)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct CloudMedium;
impl MediumTrait for CloudMedium {
    fn is_emissive(&self) -> bool {
        todo!()
    }
    fn sample_point(&self, _p: Point3f, _lambda: &SampledWavelengths) -> MediumProperties {
        todo!()
    }
    fn sample_ray(
        &self,
        _ray: &Ray,
        _t_max: Float,
        _lambda: &SampledWavelengths,
    ) -> RayMajorantIterator {
        todo!()
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct NanoVDBMedium;
impl MediumTrait for NanoVDBMedium {
    fn is_emissive(&self) -> bool {
        todo!()
    }
    fn sample_point(&self, _p: Point3f, _lambda: &SampledWavelengths) -> MediumProperties {
        todo!()
    }
    fn sample_ray(
        &self,
        _ray: &Ray,
        _t_max: Float,
        _lambda: &SampledWavelengths,
    ) -> RayMajorantIterator {
        todo!()
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct MediumInterface {
    pub inside: Ptr<Medium>,
    pub outside: Ptr<Medium>,
}


impl Default for MediumInterface {
    fn default() -> Self {
        Self {
            inside: Ptr::null(),
            outside: Ptr::null(),
        }
    }
}

impl MediumInterface {
    pub fn new(inside: Ptr<Medium>, outside: Ptr<Medium>) -> Self {
        Self { inside, outside }
    }
    pub fn empty() -> Self {
        Self::default()
    }
    pub fn is_medium_transition(&self) -> bool {
        self.inside != self.outside
    }
}