use crate::core::bsdf::BSDFSample;
use crate::core::bxdf::{BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode};
use crate::core::geometry::{Point2f, Vector3f, abs_cos_theta, same_hemisphere};
use crate::spectra::SampledSpectrum;
use crate::utils::sampling::{cosine_hemisphere_pdf, sample_cosine_hemisphere};
use crate::{Float, INV_PI};
use core::any::Any;

#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct DiffuseBxDF {
    pub r: SampledSpectrum,
}

impl DiffuseBxDF {
    pub fn new(r: SampledSpectrum) -> Self {
        Self { r }
    }
}

impl BxDFTrait for DiffuseBxDF {
    fn flags(&self) -> BxDFFlags {
        if !self.r.is_black() {
            BxDFFlags::DIFFUSE_REFLECTION
        } else {
            BxDFFlags::UNSET
        }
    }

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

    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.;
        }
        let pdf = cosine_hemisphere_pdf(abs_cos_theta(wi));
        let bsdf = BSDFSample {
            f: self.r * INV_PI,
            wi,
            pdf,
            flags: BxDFFlags::DIFFUSE_REFLECTION,
            ..Default::default()
        };
        Some(bsdf)
    }

    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) || !same_hemisphere(wo, wi) {
            return 0.;
        }
        cosine_hemisphere_pdf(abs_cos_theta(wi))
    }

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

    fn regularize(&mut self) {}
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct DiffuseTransmissionBxDF {
    pub r: SampledSpectrum,
    pub t: SampledSpectrum,
}

impl DiffuseTransmissionBxDF {
    pub fn new(r: SampledSpectrum, t: SampledSpectrum) -> Self {
        Self { r, t }
    }
}

impl BxDFTrait for DiffuseTransmissionBxDF {
    fn flags(&self) -> BxDFFlags {
        let r_flags = if !self.r.is_black() {
            BxDFFlags::DIFFUSE_REFLECTION
        } else {
            BxDFFlags::UNSET
        };
        let t_flags = if !self.t.is_black() {
            BxDFFlags::DIFFUSE_TRANSMISSION
        } else {
            BxDFFlags::UNSET
        };

        r_flags | t_flags
    }

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

    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());
        let transmission_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::TRANSMISSION.bits());

        let pr = if !f_args.sample_flags.contains(reflection_flags) {
            0.
        } else {
            self.r.max_component_value()
        };

        let pt = if !f_args.sample_flags.contains(transmission_flags) {
            0.
        } else {
            self.t.max_component_value()
        };

        if (pr == 0.) && (pt == 0.) {
            return None;
        }

        let mut wi = sample_cosine_hemisphere(u);
        if wo.z() < 0. {
            wi[2] *= -1.;
        }
        let pdf = cosine_hemisphere_pdf(abs_cos_theta(wi)) * pr / (pr + pt);

        let flags = if uc < pr / (pr + pt) {
            BxDFFlags::DIFFUSE_REFLECTION
        } else {
            BxDFFlags::DIFFUSE_TRANSMISSION
        };

        let bsdf = BSDFSample {
            f: self.r * INV_PI,
            wi,
            pdf,
            flags,
            ..Default::default()
        };
        Some(bsdf)
    }

    fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float {
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        let transmission_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::TRANSMISSION.bits());

        let pr = if !f_args.sample_flags.contains(reflection_flags) {
            0.
        } else {
            self.r.max_component_value()
        };

        let pt = if !f_args.sample_flags.contains(transmission_flags) {
            0.
        } else {
            self.t.max_component_value()
        };

        if (pr == 0.) && (pt == 0.) {
            return 0.;
        }

        let cos_factor = cosine_hemisphere_pdf(abs_cos_theta(wi));

        if same_hemisphere(wo, wi) {
            return pr / (pr + pt) * cos_factor;
        } else {
            return pt / (pr + pt) * cos_factor;
        }
    }

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

    fn regularize(&mut self) {}
}