use crate::core::bsdf::{BSDF, BSDFSample};
use crate::core::bxdf::{BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode};
use crate::core::color::RGB;
use crate::core::geometry::{
    Normal3f, Point2f, Vector3f, abs_cos_theta, cos_theta, same_hemisphere,
};
use crate::core::scattering::{
    TrowbridgeReitzDistribution, fr_complex_from_spectrum, fr_dielectric, fresnel_moment1, reflect,
    refract,
};
use crate::spectra::{DeviceStandardColorSpaces, RGBUnboundedSpectrum, SampledSpectrum};
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::sampling::{
    cosine_hemisphere_pdf, sample_cosine_hemisphere, sample_trimmed_logistic,
};
use crate::{Float, INV_2_PI, INV_PI, PI};
use core::any::Any;

static P_MAX: usize = 3;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct HairBxDF {
    pub h: Float,
    pub eta: Float,
    pub sigma_a: SampledSpectrum,
    pub beta_m: Float,
    pub beta_n: Float,
    pub v: [Float; P_MAX + 1],
    pub s: Float,
    pub sin_2k_alpha: [Float; P_MAX],
    pub cos_2k_alpha: [Float; P_MAX],
    pub colorspaces: DeviceStandardColorSpaces,
}

impl HairBxDF {
    pub fn new(
        h: Float,
        eta: Float,
        sigma_a: SampledSpectrum,
        beta_m: Float,
        beta_n: Float,
        alpha: Float,
        colorspaces: DeviceStandardColorSpaces,
    ) -> Self {
        let mut sin_2k_alpha = [0.; P_MAX];
        let mut cos_2k_alpha = [0.; P_MAX];
        sin_2k_alpha[0] = radians(alpha).sin();
        cos_2k_alpha[0] = safe_sqrt(1. - square(sin_2k_alpha[0]));

        for i in 0..P_MAX {
            sin_2k_alpha[i] = 2. * cos_2k_alpha[i - 1] * sin_2k_alpha[i - 1];
            cos_2k_alpha[i] = square(cos_2k_alpha[i - 1]) - square(sin_2k_alpha[i - 1]);
        }

        Self {
            h,
            eta,
            sigma_a,
            beta_m,
            beta_n,
            v: [0.; P_MAX + 1],
            s: 0.,
            sin_2k_alpha,
            cos_2k_alpha,
            colorspaces,
        }
    }

    fn ap(
        cos_theta_o: Float,
        eta: Float,
        h: Float,
        t: SampledSpectrum,
    ) -> [SampledSpectrum; P_MAX + 1] {
        let cos_gamma_o = safe_sqrt(1. - square(h));
        let cos_theta = cos_theta_o * cos_gamma_o;
        let f = fr_dielectric(cos_theta, eta);
        let ap0 = SampledSpectrum::new(f);
        let ap1 = t * (1.0 - f).powi(2);
        let tf = t * f;
        std::array::from_fn(|p| match p {
            0 => ap0,
            1 => ap1,
            _ if p < P_MAX => ap1 * tf.pow_int(p - 1),
            _ => ap1 * tf.pow_int(p - 1) / (SampledSpectrum::new(1.0) - tf),
        })
    }

    fn mp(
        cos_theta_i: Float,
        cos_theta_o: Float,
        sin_theta_i: Float,
        sin_theta_o: Float,
        v: Float,
    ) -> Float {
        let a = cos_theta_i * cos_theta_o / v;
        let b = sin_theta_i * sin_theta_o / v;
        if v <= 0.1 {
            fast_exp(log_i0(a) - b - 1. / v + 0.6931 + (1. / (2. * v).ln()))
        } else {
            fast_exp(-b) * i0(a) / ((1. / v).sinh() * 2. * v)
        }
    }

    fn np(phi: Float, p: i32, s: Float, gamma_o: Float, gamma_t: Float) -> Float {
        let mut dphi = phi - Self::phi(p, gamma_o, gamma_t);
        while dphi > PI {
            dphi -= 2. * PI;
        }
        while dphi < -PI {
            dphi += 2. * PI;
        }

        trimmed_logistic(dphi, s, -PI, PI)
    }

    fn phi(p: i32, gamma_o: Float, gamma_t: Float) -> Float {
        2. * p as Float * gamma_t - 2. * gamma_o + p as Float * PI
    }

    fn ap_pdf(&self, cos_theta_o: Float) -> [Float; P_MAX + 1] {
        let sin_theta_o = safe_sqrt(1. - square(cos_theta_o));
        let sin_theta_t = sin_theta_o / self.eta;
        let cos_theta_t = safe_sqrt(1. - square(sin_theta_t));
        let etap = safe_sqrt(square(self.eta) - square(sin_theta_o)) / cos_theta_o;
        let sin_gamma_t = self.h / etap;
        let cos_gamma_t = safe_sqrt(1. - square(sin_gamma_t));
        // let gamma_t = safe_asin(sin_gamma_t);
        let t_value = -self.sigma_a * (2. * cos_gamma_t / cos_theta_t);
        let t = t_value.exp();
        let ap = Self::ap(cos_theta_o, self.eta, self.h, t);
        let sum_y: Float = ap.iter().map(|s| s.average()).sum();
        std::array::from_fn(|i| ap[i].average() / sum_y)
    }

    pub fn sigma_a_from_concentration(
        ce: Float,
        cp: Float,
        stdcs: DeviceStandardColorSpaces,
    ) -> RGBUnboundedSpectrum {
        let eumelanin_sigma_a = RGB::new(0.419, 0.697, 1.37);
        let pheomelanin_sigma_a = RGB::new(0.187, 0.4, 1.05);
        let sigma_a = ce * eumelanin_sigma_a + cp * pheomelanin_sigma_a;
        RGBUnboundedSpectrum::new(&stdcs.srgb, sigma_a)
    }
}

impl BxDFTrait for HairBxDF {
    fn flags(&self) -> BxDFFlags {
        BxDFFlags::GLOSSY_REFLECTION
    }

    fn f(&self, wo: Vector3f, wi: Vector3f, _mode: TransportMode) -> SampledSpectrum {
        // Compute hair coordinate system terms related to wo
        let sin_theta_o = wo.x();
        let cos_theta_o = safe_sqrt(1. - square(sin_theta_o));
        let phi_o = wo.z().atan2(wo.y());
        let gamma_o = safe_asin(self.h);
        // Compute hair coordinate system terms related to wi
        let sin_theta_i = wi.x();
        let cos_theta_i = safe_sqrt(1. - square(sin_theta_i));
        let phi_i = wi.z().atan2(wi.y());

        let sin_theta_t = sin_theta_o / self.eta;
        let cos_theta_t = safe_sqrt(1. - square(sin_theta_t));
        let etap = safe_sqrt(square(self.eta) - square(sin_theta_o)) / cos_theta_o;
        let sin_gamma_t = self.h / etap;
        let cos_gamma_t = safe_sqrt(1. - square(sin_gamma_t));
        let gamma_t = safe_asin(sin_gamma_t);
        let sampled_value = -self.sigma_a * (2. * cos_gamma_t / cos_theta_t);
        let t = sampled_value.exp();
        let phi = phi_i - phi_o;
        let ap_pdf = Self::ap(cos_theta_o, self.eta, self.h, t);
        let mut f_sum = SampledSpectrum::new(0.);
        for (p, &ap) in ap_pdf.iter().enumerate().take(P_MAX) {
            let (sin_thetap_o, cos_thetap_o) = match p {
                0 => (
                    sin_theta_o * self.cos_2k_alpha[1] - cos_theta_o * self.sin_2k_alpha[1],
                    cos_theta_o * self.cos_2k_alpha[1] + sin_theta_o * self.sin_2k_alpha[1],
                ),
                1 => (
                    sin_theta_o * self.cos_2k_alpha[0] + cos_theta_o * self.sin_2k_alpha[0],
                    cos_theta_o * self.cos_2k_alpha[0] - sin_theta_o * self.sin_2k_alpha[0],
                ),
                2 => (
                    sin_theta_o * self.cos_2k_alpha[2] + cos_theta_o * self.sin_2k_alpha[2],
                    cos_theta_o * self.cos_2k_alpha[2] - sin_theta_o * self.sin_2k_alpha[2],
                ),
                _ => (sin_theta_o, cos_theta_o),
            };

            f_sum += Self::mp(
                cos_theta_i,
                cos_thetap_o,
                sin_theta_i,
                sin_thetap_o,
                self.v[p],
            ) * ap
                * Self::np(phi, p as i32, self.s, gamma_o, gamma_t);
        }
        if abs_cos_theta(wi) > 0. {
            f_sum /= abs_cos_theta(wi);
        }
        f_sum
    }

    fn sample_f(
        &self,
        wo: Vector3f,
        mut uc: Float,
        u: Point2f,
        f_args: FArgs,
    ) -> Option<BSDFSample> {
        let sin_theta_o = wo.x();
        let cos_theta_o = safe_sqrt(1. - square(sin_theta_o));
        let phi_o = wo.z().atan2(wo.y());
        let gamma_o = safe_asin(self.h);
        // Determine which term to sample for hair scattering
        let ap_pdf = self.ap_pdf(cos_theta_o);
        let p = sample_discrete(&ap_pdf, uc, None, Some(&mut uc));
        let (sin_thetap_o, mut cos_thetap_o) = match p {
            0 => (
                sin_theta_o * self.cos_2k_alpha[1] - cos_theta_o * self.sin_2k_alpha[1],
                cos_theta_o * self.cos_2k_alpha[1] + sin_theta_o * self.sin_2k_alpha[1],
            ),
            1 => (
                sin_theta_o * self.cos_2k_alpha[0] + cos_theta_o * self.sin_2k_alpha[0],
                cos_theta_o * self.cos_2k_alpha[0] - sin_theta_o * self.sin_2k_alpha[0],
            ),
            2 => (
                sin_theta_o * self.cos_2k_alpha[2] + cos_theta_o * self.sin_2k_alpha[2],
                cos_theta_o * self.cos_2k_alpha[2] - sin_theta_o * self.sin_2k_alpha[2],
            ),
            _ => (sin_theta_o, cos_theta_o),
        };

        cos_thetap_o = cos_thetap_o.abs();
        let cos_theta =
            1. + self.v[p] * (u[0].max(1e-5) + (1. - u[0]) * fast_exp(-2. / self.v[p])).ln();
        let sin_theta = safe_sqrt(1. - square(cos_theta));
        let cos_phi = (2. * PI * u[1]).cos();
        let sin_theta_i = -cos_theta * sin_thetap_o + sin_theta * cos_phi * cos_thetap_o;
        let cos_theta_i = safe_sqrt(1. - square(sin_theta_i));
        let etap = safe_sqrt(square(self.eta) - square(sin_theta_o)) / cos_theta_o;
        let sin_gamma_t = self.h / etap;
        // let cos_gamma_t = safe_sqrt(1. - square(sin_gamma_t));
        let gamma_t = safe_asin(sin_gamma_t);
        let dphi = if p < P_MAX {
            Self::phi(p as i32, gamma_o, gamma_t) + sample_trimmed_logistic(uc, self.s, -PI, PI)
        } else {
            2. * PI * uc
        };
        let phi_i = phi_o + dphi;
        let wi = Vector3f::new(
            sin_theta_i,
            cos_theta_i * phi_i.cos(),
            cos_theta_i * phi_i.sin(),
        );

        let mut pdf = 0.;
        for (p, &ap) in ap_pdf.iter().enumerate().take(P_MAX) {
            let (sin_thetap_o, cos_thetap_o_raw) = match p {
                0 => (
                    sin_theta_o * self.cos_2k_alpha[1] - cos_theta_o * self.sin_2k_alpha[1],
                    cos_theta_o * self.cos_2k_alpha[1] + sin_theta_o * self.sin_2k_alpha[1],
                ),
                1 => (
                    sin_theta_o * self.cos_2k_alpha[0] + cos_theta_o * self.sin_2k_alpha[0],
                    cos_theta_o * self.cos_2k_alpha[0] - sin_theta_o * self.sin_2k_alpha[0],
                ),
                2 => (
                    sin_theta_o * self.cos_2k_alpha[2] + cos_theta_o * self.sin_2k_alpha[2],
                    cos_theta_o * self.cos_2k_alpha[2] - sin_theta_o * self.sin_2k_alpha[2],
                ),
                _ => (sin_theta_o, cos_theta_o),
            };
            let cos_thetap_o = cos_thetap_o_raw.abs();
            pdf += Self::mp(
                cos_theta_i,
                cos_thetap_o,
                sin_theta_i,
                sin_thetap_o,
                self.v[p],
            ) * ap
                * Self::np(dphi, p as i32, self.s, gamma_o, gamma_t);
        }
        pdf += Self::mp(
            cos_theta_i,
            cos_theta_o,
            sin_theta_i,
            sin_theta_o,
            self.v[P_MAX],
        ) * ap_pdf[P_MAX]
            * INV_2_PI;

        let bsd = BSDFSample {
            f: self.f(wo, wi, f_args.mode),
            wi,
            pdf,
            flags: self.flags(),
            ..Default::default()
        };

        Some(bsd)
    }

    fn pdf(&self, wo: Vector3f, wi: Vector3f, _f_args: FArgs) -> Float {
        let sin_theta_o = wo.x();
        let cos_theta_o = safe_sqrt(1. - square(sin_theta_o));
        let phi_o = wo.z().atan2(wo.y());
        let gamma_o = safe_asin(self.h);
        // Determine which term to sample for hair scattering
        let sin_theta_i = wi.x();
        let cos_theta_i = safe_sqrt(1. - square(sin_theta_i));
        let phi_i = wi.z().atan2(wi.y());
        // Compute $\gammat$ for refracted ray
        let etap = safe_sqrt(self.eta * self.eta - square(sin_theta_o)) / cos_theta_o;
        let sin_gamma_t = self.h / etap;
        let gamma_t = safe_asin(sin_gamma_t);
        // Compute PDF for $A_p$ terms
        let ap_pdf = self.ap_pdf(cos_theta_o);
        let phi = phi_i - phi_o;

        let mut pdf = 0.;
        for (p, &ap) in ap_pdf.iter().enumerate().take(P_MAX) {
            let (sin_thetap_o, raw_cos_thetap_o) = match p {
                0 => (
                    sin_theta_o * self.cos_2k_alpha[1] - cos_theta_o * self.sin_2k_alpha[1],
                    cos_theta_o * self.cos_2k_alpha[1] + sin_theta_o * self.sin_2k_alpha[1],
                ),
                1 => (
                    sin_theta_o * self.cos_2k_alpha[0] + cos_theta_o * self.sin_2k_alpha[0],
                    cos_theta_o * self.cos_2k_alpha[0] - sin_theta_o * self.sin_2k_alpha[0],
                ),
                2 => (
                    sin_theta_o * self.cos_2k_alpha[2] + cos_theta_o * self.sin_2k_alpha[2],
                    cos_theta_o * self.cos_2k_alpha[2] - sin_theta_o * self.sin_2k_alpha[2],
                ),
                _ => (sin_theta_o, cos_theta_o),
            };

            let cos_thetap_o = raw_cos_thetap_o.abs();

            pdf += Self::mp(
                cos_theta_i,
                cos_thetap_o,
                sin_theta_i,
                sin_thetap_o,
                self.v[p],
            ) * ap
                * Self::np(phi, p as i32, self.s, gamma_o, gamma_t);
        }

        pdf += Self::mp(
            cos_theta_i,
            cos_theta_o,
            sin_theta_i,
            sin_theta_o,
            self.v[P_MAX],
        ) * ap_pdf[P_MAX]
            * INV_2_PI;
        pdf
    }

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

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

impl BxDFTrait for NormalizedFresnelBxDF {
    fn f(&self, wo: Vector3f, wi: Vector3f, mode: TransportMode) -> SampledSpectrum {
        if !same_hemisphere(wo, wi) {
            return SampledSpectrum::new(0.);
        }

        let c = 1. - 2. * fresnel_moment1(1. / self.eta);
        let mut f = SampledSpectrum::new((1. - fr_dielectric(cos_theta(wi), self.eta)) / (c * PI));
        if mode == TransportMode::Radiance {
            f /= square(self.eta)
        }
        f
    }

    fn sample_f(&self, wo: Vector3f, _uc: Float, u: Point2f, f_args: FArgs) -> Option<BSDFSample> {
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        if !f_args.sample_flags.contains(reflection_flags) {
            return None;
        }
        let mut wi = sample_cosine_hemisphere(u);
        if wo.z() < 0. {
            wi[2] *= -1.;
        }

        Some(BSDFSample {
            f: self.f(wo, wi, f_args.mode),
            wi,
            pdf: self.pdf(wo, wi, f_args),
            flags: BxDFFlags::DIFFUSE_REFLECTION,
            ..Default::default()
        })
    }

    fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float {
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        if !f_args.sample_flags.contains(reflection_flags) {
            return 0.;
        }

        if !same_hemisphere(wo, wi) {
            return 0.;
        }

        abs_cos_theta(wi) * INV_PI
    }

    fn flags(&self) -> BxDFFlags {
        BxDFFlags::REFLECTION | BxDFFlags::DIFFUSE
    }

    fn regularize(&mut self) {
        return;
    }

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

#[derive(Debug)]
pub struct EmptyBxDF;
impl BxDFTrait for EmptyBxDF {
    fn f(&self, _wo: Vector3f, _wi: Vector3f, _mode: TransportMode) -> SampledSpectrum {
        SampledSpectrum::default()
    }

    fn sample_f(
        &self,
        _wo: Vector3f,
        _u: Float,
        _u2: Point2f,
        _f_args: FArgs,
    ) -> Option<BSDFSample> {
        None
    }

    fn pdf(&self, _wo: Vector3f, _wi: Vector3f, _f_args: FArgs) -> Float {
        0.0
    }

    fn flags(&self) -> BxDFFlags {
        BxDFFlags::UNSET
    }

    fn regularize(&mut self) {
        return;
    }

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