use crate::core::bxdf::{
    BSDFSample, BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode,
};
use crate::core::geometry::{
    Point2f, Vector3f, VectorLike, abs_cos_theta, cos_theta, same_hemisphere, spherical_direction,
    spherical_theta,
};
use crate::core::scattering::reflect;
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::math::square;
use crate::utils::ptr::{Ptr, Slice};
use crate::utils::sampling::{PiecewiseLinear2D, cosine_hemisphere_pdf, sample_cosine_hemisphere};
use crate::{Float, INV_PI, PI, PI_OVER_2};
use core::any::Any;

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct MeasuredBxDFData {
    pub wavelengths: Slice<Float>,
    pub spectra: PiecewiseLinear2D<3>,
    pub ndf: PiecewiseLinear2D<0>,
    pub vndf: PiecewiseLinear2D<2>,
    pub sigma: PiecewiseLinear2D<0>,
    pub isotropic: bool,
    pub luminance: PiecewiseLinear2D<2>,
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct MeasuredBxDF {
    pub brdf: Ptr<MeasuredBxDFData>,
    pub lambda: SampledWavelengths,
}

unsafe impl Send for MeasuredBxDF {}
unsafe impl Sync for MeasuredBxDF {}

impl MeasuredBxDF {
    pub fn new(brdf: &MeasuredBxDFData, lambda: &SampledWavelengths) -> Self {
        Self {
            brdf: Ptr::from(brdf),
            lambda: *lambda,
        }
    }

    pub fn theta2u(theta: Float) -> Float {
        (theta * (2. / PI)).sqrt()
    }

    pub fn phi2u(phi: Float) -> Float {
        phi * 1. / (2. * PI) + 0.5
    }

    pub fn u2theta(u: Float) -> Float {
        square(u) * PI_OVER_2
    }

    pub fn u2phi(u: Float) -> Float {
        (2. * u - 1.) * PI
    }
}

impl BxDFTrait for MeasuredBxDF {
    fn flags(&self) -> BxDFFlags {
        BxDFFlags::REFLECTION | BxDFFlags::GLOSSY
    }

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

        let mut wo_curr = wo;
        let mut wi_curr = wi;
        if wo.z() < 0. {
            wo_curr = -wo_curr;
            wi_curr = -wi_curr;
        }

        // Get half direction vector
        let wm_curr = wi_curr + wo_curr;
        if wm_curr.norm_squared() == 0. {
            return SampledSpectrum::new(0.);
        }
        let wm = wm_curr.normalize();

        // Map vectors to unit square
        let theta_o = spherical_theta(wo_curr);
        let phi_o = wo_curr.y().atan2(wo_curr.x());
        let theta_m = spherical_theta(wm);
        let phi_m = wm.y().atan2(wm.x());
        let u_wo = Point2f::new(MeasuredBxDF::theta2u(theta_o), MeasuredBxDF::phi2u(phi_o));
        let u_wm_phi = if self.brdf.isotropic {
            phi_m - phi_o
        } else {
            phi_m
        };
        let mut u_wm = Point2f::new(
            MeasuredBxDF::theta2u(theta_m),
            MeasuredBxDF::phi2u(u_wm_phi),
        );
        u_wm[1] -= u_wm[1].floor();

        // Inverse parametrization
        let ui = self.brdf.vndf.invert(u_wm, [phi_o, theta_o]);
        let fr = SampledSpectrum::from_fn(|i| {
            self.brdf
                .spectra
                .evaluate(ui.p, [phi_o, theta_o, self.lambda[i]])
                .max(0.0)
        });

        fr * self.brdf.ndf.evaluate(u_wm, [])
            / (4. * self.brdf.sigma.evaluate(u_wo, []) * cos_theta(wi))
    }

    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.;
        }

        let mut wo_curr = wo;
        let mut wi_curr = wi;
        if wo.z() < 0. {
            wo_curr = -wo_curr;
            wi_curr = -wi_curr;
        }

        let wm_curr = wi_curr + wo_curr;
        if wm_curr.norm_squared() == 0. {
            return 0.;
        }
        let wm = wm_curr.normalize();
        let theta_o = spherical_theta(wo_curr);
        let phi_o = wo_curr.y().atan2(wo_curr.x());
        let theta_m = spherical_theta(wm);
        let phi_m = wm.y().atan2(wm.x());

        let u_wm_phi = if self.brdf.isotropic {
            phi_m - phi_o
        } else {
            phi_m
        };

        let mut u_wm = Point2f::new(
            MeasuredBxDF::theta2u(theta_m),
            MeasuredBxDF::phi2u(u_wm_phi),
        );
        u_wm[1] = u_wm[1] - u_wm[1].floor();

        let ui = self.brdf.vndf.invert(u_wm, [phi_o, theta_o]);
        let sample = ui.p;
        let vndf_pdf = ui.pdf;

        let pdf = self.brdf.luminance.evaluate(sample, [phi_o, theta_o]);
        let sin_theta_m = (square(wm.x()) + square(wm.y())).sqrt();
        let jacobian = 4. * wm.dot(wo) * f32::max(2. * square(PI) * u_wm.x() * sin_theta_m, 1e-6);
        vndf_pdf * pdf / jacobian
    }

    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 flip_w = false;
        let mut wo_curr = wo;
        if wo.z() <= 0. {
            wo_curr = -wo_curr;
            flip_w = true;
        }

        let theta_o = spherical_theta(wo_curr);
        let phi_o = wo_curr.y().atan2(wo_curr.x());
        // Warp sample using luminance distribution
        let mut s = self.brdf.luminance.sample(u, [phi_o, theta_o]);
        let u = s.p;
        let lum_pdf = s.pdf;

        // Sample visible normal distribution of measured BRDF
        s = self.brdf.vndf.sample(u, [phi_o, theta_o]);
        let u_wm = s.p;
        let mut pdf = s.pdf;

        // Map from microfacet normal to incident direction
        let mut phi_m = MeasuredBxDF::u2phi(u_wm.y());
        let theta_m = MeasuredBxDF::u2theta(u_wm.x());
        if self.brdf.isotropic {
            phi_m += phi_o;
        }
        let sin_theta_m = theta_m.sin();
        let cos_theta_m = theta_m.cos();
        let wm = spherical_direction(sin_theta_m, cos_theta_m, phi_m);
        let mut wi = reflect(wo_curr, wm.into());
        if wi.z() <= 0. {
            return None;
        }

        // Interpolate spectral BRDF
        let mut f = SampledSpectrum::from_fn(|i| {
            self.brdf
                .spectra
                .evaluate(u, [phi_o, theta_o, self.lambda[i]])
                .max(0.0)
        });

        let u_wo = Point2f::new(MeasuredBxDF::theta2u(theta_o), MeasuredBxDF::phi2u(phi_o));
        f *= self.brdf.ndf.evaluate(u_wm, [])
            / (4. * self.brdf.sigma.evaluate(u_wo, []) * abs_cos_theta(wi));
        pdf /= 4. * wm.dot(wo_curr) * f32::max(2. * square(PI) * u_wm.x(), 1e-6);

        if flip_w {
            wi = -wi;
        }

        let bsdf = BSDFSample {
            f,
            wi,
            pdf: pdf * lum_pdf,
            flags: BxDFFlags::GLOSSY_REFLECTION,
            ..Default::default()
        };

        Some(bsdf)
    }

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

    fn regularize(&mut self) {
        return;
    }
}