use crate::core::bxdf::{
    BSDFSample, BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode,
};
use crate::core::geometry::{
    Normal3f, Point2f, Vector3f, VectorLike, abs_cos_theta, cos_theta, same_hemisphere,
};
use crate::core::scattering::{
    TrowbridgeReitzDistribution, fr_complex_from_spectrum, fr_dielectric, reflect, refract,
};
use crate::spectra::SampledSpectrum;
use crate::utils::math::square;
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)]
pub struct DielectricBxDF {
    pub eta: Float,
    pub mf_distrib: TrowbridgeReitzDistribution,
}

impl DielectricBxDF {
    pub fn new(eta: Float, mf_distrib: TrowbridgeReitzDistribution) -> Self {
        Self { eta, mf_distrib }
    }
}

impl BxDFTrait for DielectricBxDF {
    fn flags(&self) -> BxDFFlags {
        let flags = if self.eta == 1. {
            BxDFFlags::TRANSMISSION
        } else {
            BxDFFlags::REFLECTION | BxDFFlags::TRANSMISSION
        };
        flags
            | if self.mf_distrib.effectively_smooth() {
                BxDFFlags::SPECULAR
            } else {
                BxDFFlags::GLOSSY
            }
    }

    fn f(&self, wo: Vector3f, wi: Vector3f, mode: TransportMode) -> SampledSpectrum {
        if self.eta == 1. || self.mf_distrib.effectively_smooth() {
            return SampledSpectrum::new(0.);
        }

        // Generalized half vector wm
        let cos_theta_o = cos_theta(wo);
        let cos_theta_i = cos_theta(wi);
        let reflect = cos_theta_i * cos_theta_o > 0.;

        let mut etap = 1.;
        if !reflect {
            etap = if cos_theta_o > 0. {
                self.eta
            } else {
                1. / self.eta
            };
        }

        let wm_orig = wi * etap + wo;
        if cos_theta_i == 0. || cos_theta_o == 0. || wm_orig.norm_squared() == 0. {
            return SampledSpectrum::new(0.);
        }
        let wm = Normal3f::new(0., 0., 1.).face_forward(wm_orig.normalize());

        if wi.dot(wm.into()) * cos_theta_i < 0. || wo.dot(wm.into()) * cos_theta_o < 0. {
            return SampledSpectrum::new(0.);
        }

        let fr = fr_dielectric(wo.dot(wm.into()), self.eta);

        if reflect {
            SampledSpectrum::new(
                self.mf_distrib.d(wm.into()) * self.mf_distrib.g(wo, wi) * fr
                    / (4. * cos_theta_i * cos_theta_o).abs(),
            )
        } else {
            let denom =
                square(wi.dot(wm.into()) + wo.dot(wm.into()) / etap) * cos_theta_i * cos_theta_o;
            let mut ft = self.mf_distrib.d(wm.into())
                * (1. - fr)
                * self.mf_distrib.g(wo, wi)
                * (wi.dot(wm.into()) * wo.dot(wm.into()) / denom).abs();
            if mode == TransportMode::Radiance {
                ft /= square(etap)
            }

            SampledSpectrum::new(ft)
        }
    }

    fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float {
        if self.eta == 1. || self.mf_distrib.effectively_smooth() {
            return 0.;
        }

        let cos_theta_o = cos_theta(wo);
        let cos_theta_i = cos_theta(wi);

        let reflect = cos_theta_i * cos_theta_o > 0.;
        let mut etap = 1.;
        if !reflect {
            etap = if cos_theta_o > 0. {
                self.eta
            } else {
                1. / self.eta
            };
        }

        let wm_orig = wi * etap + wo;

        if cos_theta_i == 0. || cos_theta_o == 0. || wm_orig.norm_squared() == 0. {
            return 0.;
        }
        let wm = Normal3f::new(0., 0., 1.).face_forward(wm_orig.normalize());

        // Discard backfacing microfacets
        if wi.dot(wm.into()) * cos_theta_i < 0. || wo.dot(wm.into()) * cos_theta_o < 0. {
            return 0.;
        }

        let r = fr_dielectric(wo.dot(wm.into()), self.eta);
        let t = 1. - r;
        let mut pr = r;
        let mut pt = t;
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        let transmission_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::TRANSMISSION.bits());
        if !f_args.sample_flags.contains(reflection_flags) {
            pr = 0.;
        }
        if !f_args.sample_flags.contains(transmission_flags) {
            pt = 0.;
        }
        if pr == 0. && pt == 0. {
            return 0.;
        }

        if reflect {
            self.mf_distrib.pdf(
                wo,
                Vector3f::from(wm) / (4. * wo.dot(wm.into()).abs()) * pr / (pt + pr),
            )
        } else {
            let denom = square(wi.dot(wm.into()) + wo.dot(wm.into()) / etap);
            let dwm_dwi = wi.dot(wm.into()).abs() / denom;
            self.mf_distrib.pdf(wo, wm.into()) * dwm_dwi * pr / (pr + pt)
        }
    }

    fn sample_f(&self, wo: Vector3f, uc: Float, u: Point2f, f_args: FArgs) -> Option<BSDFSample> {
        if self.eta == 1. || self.mf_distrib.effectively_smooth() {
            let r = fr_dielectric(cos_theta(wo), self.eta);
            let t = 1. - r;
            let mut pr = r;
            let mut pt = t;
            let reflection_flags =
                BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
            let transmission_flags =
                BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::TRANSMISSION.bits());
            if !f_args.sample_flags.contains(reflection_flags) {
                pr = 0.;
            }
            if !f_args.sample_flags.contains(transmission_flags) {
                pt = 0.;
            }
            // If probabilities are null, doesnt contribute
            if pr == 0. && pt == 0. {
                return None;
            }

            if uc < pr / (pr + pt) {
                let wi = Vector3f::new(-wo.x(), -wo.y(), wo.z());
                let fr = SampledSpectrum::new(r / abs_cos_theta(wi));
                let bsdf = BSDFSample {
                    f: fr,
                    wi,
                    pdf: pr / (pr + pt),
                    flags: BxDFFlags::SPECULAR_REFLECTION,
                    ..Default::default()
                };
                Some(bsdf)
            } else {
                // Compute ray direction for specular transmission
                if let Some((wi, etap)) = refract(wo, Normal3f::new(0., 0., 1.), self.eta) {
                    let mut ft = SampledSpectrum::new(t / abs_cos_theta(wi));
                    if f_args.mode == TransportMode::Radiance {
                        ft /= square(etap);
                    }
                    let bsdf = BSDFSample {
                        f: ft,
                        wi,
                        pdf: pt / (pr + pt),
                        flags: BxDFFlags::SPECULAR_TRANSMISSION,
                        eta: etap,
                        ..Default::default()
                    };
                    Some(bsdf)
                } else {
                    None
                }
            }
        } else {
            // Sample rough dielectric BSDF
            let wm = self.mf_distrib.sample_wm(wo, u);
            let r = fr_dielectric(wo.dot(wm), self.eta);
            let t = 1. - r;
            let mut pr = r;
            let mut pt = t;
            let reflection_flags =
                BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
            let transmission_flags =
                BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::TRANSMISSION.bits());
            if !f_args.sample_flags.contains(reflection_flags) {
                pr = 0.;
            }
            if !f_args.sample_flags.contains(transmission_flags) {
                pt = 0.;
            }
            if pr == 0. && pt == 0. {
                return None;
            }
            let pdf: Float;
            if uc < pr / (pr + pt) {
                // Sample reflection at rough dielectric interface
                let wi = reflect(wo, wm.into());
                if !same_hemisphere(wo, wi) {
                    return None;
                }

                pdf = self.mf_distrib.pdf(wo, wm) / (4. * wo.dot(wm).abs()) * pr / (pr + pt);
                let f = SampledSpectrum::new(
                    self.mf_distrib.d(wm) * self.mf_distrib.g(wo, wi) * r
                        / (4. * cos_theta(wi) * cos_theta(wo)),
                );
                let bsdf = BSDFSample {
                    f,
                    wi,
                    pdf,
                    flags: BxDFFlags::GLOSSY_REFLECTION,
                    ..Default::default()
                };
                Some(bsdf)
            } else {
                // Sample transmission at rough dielectric interface
                if let Some((wi, etap)) = refract(wo, wm.into(), self.eta) {
                    if same_hemisphere(wo, wi) || wi.z() == 0. {
                        None
                    } else {
                        let denom = square(wi.dot(wm) + wo.dot(wm) / etap);
                        let dwm_mi = wi.dot(wm).abs() / denom;
                        pdf = self.mf_distrib.pdf(wo, wm) * dwm_mi * pt / (pr + pt);
                        let mut ft = SampledSpectrum::new(
                            t * self.mf_distrib.d(wm)
                                * self.mf_distrib.g(wo, wi)
                                * (wi.dot(wm) * wo.dot(wm)).abs()
                                / (cos_theta(wi) * cos_theta(wo) * denom),
                        );
                        if f_args.mode == TransportMode::Radiance {
                            ft /= square(etap);
                        }
                        let bsdf = BSDFSample {
                            f: ft,
                            wi,
                            pdf,
                            flags: BxDFFlags::GLOSSY_TRANSMISSION,
                            eta: etap,
                            ..Default::default()
                        };
                        Some(bsdf)
                    }
                } else {
                    None
                }
            }
        }
    }

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

    fn regularize(&mut self) {
        self.mf_distrib.regularize();
    }
}
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct ThinDielectricBxDF {
    pub eta: Float,
}

impl ThinDielectricBxDF {
    pub fn new(eta: Float) -> Self {
        Self { eta }
    }
}

impl BxDFTrait for ThinDielectricBxDF {
    fn flags(&self) -> BxDFFlags {
        BxDFFlags::REFLECTION | BxDFFlags::TRANSMISSION | BxDFFlags::SPECULAR
    }

    fn f(&self, _wo: Vector3f, _wi: Vector3f, _mode: TransportMode) -> SampledSpectrum {
        SampledSpectrum::new(0.)
    }

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

    fn sample_f(&self, wo: Vector3f, uc: Float, _u: Point2f, f_args: FArgs) -> Option<BSDFSample> {
        let mut r = fr_dielectric(abs_cos_theta(wo), self.eta);
        let mut t = 1. - r;
        if r < 1. {
            r += square(t) * r / (1. - square(r));
            t = 1. - r;
        }
        let mut pr = r;
        let mut pt = t;
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        let transmission_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::TRANSMISSION.bits());
        if !f_args.sample_flags.contains(reflection_flags) {
            pr = 0.;
        }
        if !f_args.sample_flags.contains(transmission_flags) {
            pt = 0.;
        }
        if pr == 0. && pt == 0. {
            return None;
        }

        if uc < pr / (pr + pt) {
            let wi = Vector3f::new(-wo.x(), -wo.y(), wo.z());
            let f = SampledSpectrum::new(r / abs_cos_theta(wi));
            let bsdf = BSDFSample {
                f,
                wi,
                pdf: pr / (pr + pt),
                flags: BxDFFlags::SPECULAR_REFLECTION,
                ..Default::default()
            };
            Some(bsdf)
        } else {
            // Perfect specular transmission
            let wi = -wo;
            let f = SampledSpectrum::new(t / abs_cos_theta(wi));
            let bsdf = BSDFSample {
                f,
                wi,
                pdf: pr / (pr + pt),
                flags: BxDFFlags::SPECULAR_TRANSMISSION,
                ..Default::default()
            };
            Some(bsdf)
        }
    }

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