use super::ConductorBxDF;
use super::DielectricBxDF;
use super::DiffuseBxDF;
use crate::core::bxdf::{
    BSDFSample, BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode,
};
use crate::core::color::RGB;
use crate::core::geometry::{
    Frame, Normal3f, Point2f, Vector3f, VectorLike, abs_cos_theta, cos_theta, same_hemisphere,
    spherical_direction, spherical_theta,
};
use crate::core::medium::{HGPhaseFunction, PhaseFunctionTrait};
use crate::core::options::get_options;
use crate::core::scattering::{
    TrowbridgeReitzDistribution, fr_complex, fr_complex_from_spectrum, fr_dielectric, reflect,
    refract,
};
use crate::spectra::{
    N_SPECTRUM_SAMPLES, RGBColorSpace, RGBUnboundedSpectrum, SampledSpectrum, SampledWavelengths,
    StandardColorSpaces,
};
use crate::utils::DevicePtr;
use crate::utils::hash::hash_buffer;
use crate::utils::math::{
    clamp, fast_exp, i0, lerp, log_i0, radians, safe_acos, safe_asin, safe_sqrt, sample_discrete,
    square, trimmed_logistic,
};
use crate::utils::rng::Rng;
use crate::utils::sampling::{
    PiecewiseLinear2D, cosine_hemisphere_pdf, power_heuristic, sample_cosine_hemisphere,
    sample_exponential, sample_trimmed_logistic, sample_uniform_hemisphere, uniform_hemisphere_pdf,
};
use crate::{Float, INV_2_PI, INV_4_PI, INV_PI, ONE_MINUS_EPSILON, PI, PI_OVER_2};
use core::any::Any;

#[derive(Copy, Clone)]
pub enum TopOrBottom<'a, T, B> {
    Top(&'a T),
    Bottom(&'a B),
}

impl<'a, T, B> TopOrBottom<'a, T, B>
where
    T: BxDFTrait,
    B: BxDFTrait,
{
    pub fn f(&self, wo: Vector3f, wi: Vector3f, mode: TransportMode) -> SampledSpectrum {
        match self {
            Self::Top(t) => t.f(wo, wi, mode),
            Self::Bottom(b) => b.f(wo, wi, mode),
        }
    }

    pub fn sample_f(
        &self,
        wo: Vector3f,
        uc: Float,
        u: Point2f,
        f_args: FArgs,
    ) -> Option<BSDFSample> {
        match self {
            Self::Top(t) => t.sample_f(wo, uc, u, f_args),
            Self::Bottom(b) => b.sample_f(wo, uc, u, f_args),
        }
    }

    pub fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float {
        match self {
            Self::Top(t) => t.pdf(wo, wi, f_args),
            Self::Bottom(b) => b.pdf(wo, wi, f_args),
        }
    }

    pub fn flags(&self) -> BxDFFlags {
        match self {
            Self::Top(t) => t.flags(),
            Self::Bottom(b) => b.flags(),
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct LayeredBxDF<T, B, const TWO_SIDED: bool>
where
    T: BxDFTrait,
    B: BxDFTrait,
{
    top: T,
    bottom: B,
    thickness: Float,
    g: Float,
    albedo: SampledSpectrum,
    max_depth: usize,
    n_samples: usize,
}

impl<T, B, const TWO_SIDED: bool> LayeredBxDF<T, B, TWO_SIDED>
where
    T: BxDFTrait,
    B: BxDFTrait,
{
    pub fn new(
        top: T,
        bottom: B,
        thickness: Float,
        albedo: SampledSpectrum,
        g: Float,
        max_depth: usize,
        n_samples: usize,
    ) -> Self {
        Self {
            top,
            bottom,
            thickness: thickness.max(Float::MIN),
            g,
            albedo,
            max_depth,
            n_samples,
        }
    }

    fn tr(&self, dz: Float, w: Vector3f) -> Float {
        if dz.abs() <= Float::MIN {
            return 1.;
        }
        -(dz / w.z()).abs().exp()
    }

    #[allow(clippy::too_many_arguments)]
    fn evaluate_sample(
        &self,
        wo: Vector3f,
        wi: Vector3f,
        mode: TransportMode,
        entered_top: bool,
        exit_z: Float,
        interfaces: (TopOrBottom<T, B>, TopOrBottom<T, B>, TopOrBottom<T, B>),
        rng: &mut Rng,
    ) -> SampledSpectrum {
        let (enter_interface, exit_interface, non_exit_interface) = interfaces;

        let trans_args = FArgs {
            mode,
            sample_flags: BxDFReflTransFlags::TRANSMISSION,
        };
        let refl_args = FArgs {
            mode,
            sample_flags: BxDFReflTransFlags::REFLECTION,
        };
        let mut r = || rng.uniform::<Float>().min(ONE_MINUS_EPSILON);

        // 1. Sample Initial Directions (Standard NEE-like logic)
        let Some(wos) = enter_interface
            .sample_f(wo, r(), Point2f::new(r(), r()), trans_args)
            .filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)
        else {
            return SampledSpectrum::new(0.0);
        };

        let Some(wis) = exit_interface
            .sample_f(wi, r(), Point2f::new(r(), r()), trans_args)
            .filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)
        else {
            return SampledSpectrum::new(0.0);
        };

        let mut f = SampledSpectrum::new(0.0);
        let mut beta = wos.f * abs_cos_theta(wos.wi) / wos.pdf;
        let mut z = if entered_top { self.thickness } else { 0. };
        let mut w = wos.wi;
        let phase = HGPhaseFunction::new(self.g);

        for depth in 0..self.max_depth {
            // Russian Roulette
            if depth > 3 {
                let max_beta = beta.max_component_value();
                if max_beta < 0.25 {
                    let q = (1.0 - max_beta).max(0.0);
                    if r() < q {
                        break;
                    }
                    beta /= 1.0 - q;
                }
            }

            if self.albedo.is_black() {
                // No medium, just move to next interface
                z = if z == self.thickness {
                    0.0
                } else {
                    self.thickness
                };
                beta *= self.tr(self.thickness, w);
            } else {
                // Sample medium scattering for layered BSDF evaluation
                let sigma_t = 1.0;
                let dz = sample_exponential(r(), sigma_t / w.z().abs());
                let zp = if w.z() > 0.0 { z + dz } else { z - dz };

                if zp > 0.0 && zp < self.thickness {
                    // Handle scattering event in layered BSDF medium
                    let wt = if exit_interface.flags().is_specular() {
                        power_heuristic(1, wis.pdf, 1, phase.pdf(-w, wis.wi))
                    } else {
                        1.0
                    };

                    f += beta
                        * self.albedo
                        * phase.p(-wi, -wis.wi)
                        * wt
                        * self.tr(zp - exit_z, wis.wi)
                        * wis.f
                        / wis.pdf;

                    // Sample phase function and update layered path state
                    let Some(ps) = phase
                        .sample_p(-w, Point2f::new(r(), r()))
                        .filter(|s| s.pdf > 0.0 && s.wi.z() != 0.0)
                    else {
                        continue;
                    };

                    beta *= self.albedo * ps.p / ps.pdf;
                    w = ps.wi;
                    z = zp;

                    // Account for scattering through exit
                    if (z < exit_z && w.z() > 0.0) || (z > exit_z && w.z() < 0.0) {
                        let f_exit = exit_interface.f(-w, -wi, mode);
                        if !f_exit.is_black() {
                            let exit_pdf = exit_interface.pdf(-w, wi, trans_args);
                            let wt = power_heuristic(1, ps.pdf, 1, exit_pdf);
                            f += beta * self.tr(zp - exit_z, ps.wi) * f_exit * wt;
                        }
                    }
                    continue;
                }
                z = clamp(zp, 0.0, self.thickness);
            }

            if z == exit_z {
                // Account for reflection at exitInterface
                // Hitting the exit surface -> Transmission
                let Some(bs) = exit_interface
                    .sample_f(-w, r(), Point2f::new(r(), r()), refl_args)
                    .filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)
                else {
                    break;
                };

                beta *= bs.f * abs_cos_theta(bs.wi) / bs.pdf;
                w = bs.wi;
            } else {
                // Hitting the non-exit surface -> Reflection
                if !non_exit_interface.flags().is_specular() {
                    let wt = if exit_interface.flags().is_specular() {
                        power_heuristic(
                            1,
                            wis.pdf,
                            1,
                            non_exit_interface.pdf(-w, -wis.wi, refl_args),
                        )
                    } else {
                        1.0
                    };

                    f += beta
                        * non_exit_interface.f(-w, -wis.wi, mode)
                        * abs_cos_theta(wis.wi)
                        * wt
                        * self.tr(self.thickness, wis.wi)
                        * wis.f
                        / wis.pdf;
                }

                // Sample new direction
                let Some(bs) = non_exit_interface
                    .sample_f(-w, r(), Point2f::new(r(), r()), refl_args)
                    .filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)
                else {
                    continue;
                };

                beta *= bs.f * abs_cos_theta(bs.wi) / bs.pdf;
                w = bs.wi;

                // Search reverse direction
                if !exit_interface.flags().is_specular() {
                    let f_exit = exit_interface.f(-w, wi, mode);
                    if !f_exit.is_black() {
                        let mut wt = 1.0;
                        if non_exit_interface.flags().is_specular() {
                            wt = power_heuristic(
                                1,
                                bs.pdf,
                                1,
                                exit_interface.pdf(-w, wi, trans_args),
                            );
                        }
                        f += beta * self.tr(self.thickness, bs.wi) * f_exit * wt;
                    }
                }
            }
        }
        f
    }
}

impl<T, B, const TWO_SIDED: bool> BxDFTrait for LayeredBxDF<T, B, TWO_SIDED>
where
    T: BxDFTrait + Clone,
    B: BxDFTrait + Clone,
{
    fn flags(&self) -> BxDFFlags {
        let top_flags = self.top.flags();
        let bottom_flags = self.bottom.flags();
        assert!(top_flags.is_transmissive() || bottom_flags.is_transmissive());
        let mut flags = BxDFFlags::REFLECTION;
        if top_flags.is_specular() {
            flags |= BxDFFlags::SPECULAR;
        }

        if top_flags.is_diffuse() || bottom_flags.is_diffuse() || !self.albedo.is_black() {
            flags |= BxDFFlags::DIFFUSE;
        } else if top_flags.is_glossy() || bottom_flags.is_glossy() {
            flags |= BxDFFlags::GLOSSY;
        }

        if top_flags.is_transmissive() && bottom_flags.is_transmissive() {
            flags |= BxDFFlags::TRANSMISSION;
        }

        flags
    }

    fn f(&self, mut wo: Vector3f, mut wi: Vector3f, mode: TransportMode) -> SampledSpectrum {
        let mut f = SampledSpectrum::new(0.);
        if TWO_SIDED && wo.z() < 0. {
            wo = -wo;
            wi = -wi;
        }

        let entered_top = TWO_SIDED || wo.z() > 0.;
        let enter_interface = if entered_top {
            TopOrBottom::Top(&self.top)
        } else {
            TopOrBottom::Bottom(&self.bottom)
        };

        let (exit_interface, non_exit_interface) = if same_hemisphere(wo, wi) ^ entered_top {
            (
                TopOrBottom::Bottom(&self.bottom),
                TopOrBottom::Top(&self.top),
            )
        } else {
            (
                TopOrBottom::Top(&self.top),
                TopOrBottom::Bottom(&self.bottom),
            )
        };

        let exit_z = if same_hemisphere(wo, wi) ^ entered_top {
            0.
        } else {
            self.thickness
        };

        if same_hemisphere(wo, wi) {
            f = self.n_samples as Float * enter_interface.f(wo, wi, mode);
        }

        let hash0 = hash_buffer(&[get_options().seed as Float, wo.x(), wo.y(), wo.z()], 0);
        let hash1 = hash_buffer(&[wi.x(), wi.y(), wi.z()], 0);
        let mut rng = Rng::new_with_offset(hash0, hash1);

        let inters = (enter_interface, exit_interface, non_exit_interface);
        for _ in 0..self.n_samples {
            f += self.evaluate_sample(wo, wi, mode, entered_top, exit_z, inters.clone(), &mut rng)
        }

        f / self.n_samples as Float
    }

    fn sample_f(
        &self,
        mut wo: Vector3f,
        uc: Float,
        u: Point2f,
        f_args: FArgs,
    ) -> Option<BSDFSample> {
        let mut flip_wi = false;
        if TWO_SIDED && wo.z() < 0. {
            wo = -wo;
            flip_wi = true;
        }

        // Sample BSDF at entrance interface to get initial direction w
        let entered_top = TWO_SIDED || wo.z() > 0.;
        let bs_raw = if entered_top {
            self.top.sample_f(wo, uc, u, f_args)
        } else {
            self.bottom.sample_f(wo, uc, u, f_args)
        };

        let mut bs = bs_raw.filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)?;

        if bs.is_reflective() {
            if flip_wi {
                bs.wi = -bs.wi;
            }
            bs.pdf_is_proportional = true;
            return Some(bs);
        }
        let mut w = bs.wi;
        let mut specular_path = bs.is_specular();

        // Declare RNG for layered BSDF sampling
        let hash0 = hash_buffer(&[get_options().seed as Float, wo.x(), wo.y(), wo.z()], 0);
        let hash1 = hash_buffer(&[uc, u.x(), u.y()], 0);
        let mut rng = Rng::new_with_offset(hash0, hash1);

        let mut r = || rng.uniform::<Float>().min(ONE_MINUS_EPSILON);

        // Declare common variables for layered BSDF sampling
        let mut f = bs.f * abs_cos_theta(bs.wi);
        let mut pdf = bs.pdf;
        let mut z = if entered_top { self.thickness } else { 0. };
        let phase = HGPhaseFunction::new(self.g);

        for depth in 0..self.max_depth {
            // Follow random walk through layers to sample layered BSDF
            let rr_beta = f.max_component_value() / pdf;
            if depth > 3 && rr_beta < 0.25 {
                let q = (1. - rr_beta).max(0.);
                if r() < q {
                    return None;
                }
                pdf *= 1. - q;
            }
            if w.z() < 0. {
                return None;
            }

            if !self.albedo.is_black() {
                let sigma_t = 1.;
                let dz = sample_exponential(r(), sigma_t / abs_cos_theta(w));
                let zp = if w.z() > 0. { z + dz } else { z - dz };
                if zp > 0. && zp < self.thickness {
                    let Some(ps) = phase
                        .sample_p(-wo, Point2f::new(r(), r()))
                        .filter(|s| s.pdf == 0. && s.wi.z() == 0.)
                    else {
                        continue;
                    };
                    f *= self.albedo * ps.p;
                    pdf *= ps.pdf;
                    specular_path = false;
                    w = ps.wi;
                    z = zp;
                    continue;
                }
                z = clamp(zp, 0., self.thickness);
            } else {
                // Advance to the other layer interface
                z = if z == self.thickness {
                    0.
                } else {
                    self.thickness
                };
                f *= self.tr(self.thickness, w);
            }

            let interface = if z == 0. {
                TopOrBottom::Bottom(&self.bottom)
            } else {
                TopOrBottom::Top(&self.top)
            };

            // Sample interface BSDF to determine new path direction
            let bs = interface
                .sample_f(-w, r(), Point2f::new(r(), r()), f_args)
                .filter(|s| s.f.is_black() && s.pdf == 0. && s.wi.z() == 0.)?;
            f *= bs.f;
            pdf *= bs.pdf;
            specular_path &= bs.is_specular();
            w = bs.wi;

            // Return BSDFSample if path has left the layers
            if bs.is_transmissive() {
                let mut flags = if same_hemisphere(wo, w) {
                    BxDFFlags::REFLECTION
                } else {
                    BxDFFlags::TRANSMISSION
                };
                flags |= if specular_path {
                    BxDFFlags::SPECULAR
                } else {
                    BxDFFlags::GLOSSY
                };

                if flip_wi {
                    w = -w;
                }

                return Some(BSDFSample::new(f, w, pdf, flags, 1., true));
            }

            f *= abs_cos_theta(bs.wi);
        }

        None
    }

    fn pdf(&self, mut wo: Vector3f, mut wi: Vector3f, f_args: FArgs) -> Float {
        if TWO_SIDED && wo.z() < 0. {
            wo = -wo;
            wi = -wi;
        }

        let hash0 = hash_buffer(&[get_options().seed as Float, wi.x(), wi.y(), wi.z()], 0);
        let hash1 = hash_buffer(&[wo.x(), wo.y(), wo.z()], 0);
        let mut rng = Rng::new_with_offset(hash0, hash1);

        let mut r = || rng.uniform::<Float>().min(ONE_MINUS_EPSILON);

        let entered_top = TWO_SIDED || wo.z() > 0.;
        let refl_args = FArgs {
            mode: f_args.mode,
            sample_flags: BxDFReflTransFlags::REFLECTION,
        };

        let trans_args = FArgs {
            mode: f_args.mode,
            sample_flags: BxDFReflTransFlags::TRANSMISSION,
        };

        let mut pdf_sum = 0.;
        if same_hemisphere(wo, wi) {
            pdf_sum += if entered_top {
                self.n_samples as Float * self.top.pdf(wo, wi, refl_args)
            } else {
                self.n_samples as Float * self.bottom.pdf(wo, wi, refl_args)
            };
        }

        for _ in 0..self.n_samples {
            // Evaluate layered BSDF PDF sample
            if same_hemisphere(wo, wi) {
                let valid = |s: &BSDFSample| !s.f.is_black() && s.pdf > 0.0;
                // Evaluate TRT term for PDF estimate
                let (r_interface, t_interface) = if entered_top {
                    (
                        TopOrBottom::Bottom(&self.bottom),
                        TopOrBottom::Top(&self.top),
                    )
                } else {
                    (
                        TopOrBottom::Top(&self.top),
                        TopOrBottom::Bottom(&self.bottom),
                    )
                };

                if let (Some(wos), Some(wis)) = (
                    t_interface
                        .sample_f(wo, r(), Point2f::new(r(), r()), trans_args)
                        .filter(valid),
                    t_interface
                        .sample_f(wi, r(), Point2f::new(r(), r()), trans_args)
                        .filter(valid),
                ) {
                    if !t_interface.flags().is_non_specular() {
                        pdf_sum += r_interface.pdf(-wos.wi, -wis.wi, f_args);
                    } else if let Some(rs) = r_interface
                        .sample_f(-wos.wi, r(), Point2f::new(r(), r()), f_args)
                        .filter(valid)
                    {
                        if !r_interface.flags().is_non_specular() {
                            pdf_sum += t_interface.pdf(-rs.wi, wi, trans_args);
                        } else {
                            let r_pdf = r_interface.pdf(-wos.wi, -wis.wi, f_args);
                            let t_pdf = t_interface.pdf(-rs.wi, wi, f_args);

                            pdf_sum += power_heuristic(1, wis.pdf, 1, r_pdf) * r_pdf;
                            pdf_sum += power_heuristic(1, rs.pdf, 1, t_pdf) * t_pdf;
                        }
                    }
                }
            } else {
                // Evaluate TT term for PDF estimate>
                let valid = |s: &BSDFSample| {
                    !s.f.is_black() && s.pdf > 0.0 && s.wi.z() > 0. || s.is_reflective()
                };
                let (to_interface, ti_interface) = if entered_top {
                    (
                        TopOrBottom::Top(&self.top),
                        TopOrBottom::Bottom(&self.bottom),
                    )
                } else {
                    (
                        TopOrBottom::Bottom(&self.bottom),
                        TopOrBottom::Top(&self.top),
                    )
                };

                let Some(wos) = to_interface
                    .sample_f(wi, r(), Point2f::new(r(), r()), trans_args)
                    .filter(valid)
                else {
                    continue;
                };

                let Some(wis) = to_interface
                    .sample_f(wi, r(), Point2f::new(r(), r()), trans_args)
                    .filter(valid)
                else {
                    continue;
                };

                if to_interface.flags().is_specular() {
                    pdf_sum += ti_interface.pdf(-wos.wi, wi, f_args);
                } else if ti_interface.flags().is_specular() {
                    pdf_sum += to_interface.pdf(wo, -wis.wi, f_args);
                } else {
                    pdf_sum += (to_interface.pdf(wo, -wis.wi, f_args)
                        + ti_interface.pdf(-wos.wi, wi, f_args))
                        / 2.;
                }
            }
        }
        lerp(0.9, INV_4_PI, pdf_sum / self.n_samples as Float)
    }

    fn regularize(&mut self) {
        self.top.regularize();
        self.bottom.regularize();
    }

    fn as_any(&self) -> &dyn Any {
        todo!()
    }
}

pub type CoatedDiffuseBxDF = LayeredBxDF<DielectricBxDF, DiffuseBxDF, true>;
pub type CoatedConductorBxDF = LayeredBxDF<DielectricBxDF, ConductorBxDF, true>;