use rand_distr::{Distribution, Normal};
use ray_tracing_core::{material::Material, math::material::fresnel_real, prelude::*};
use std::fmt::Debug;

#[derive(Debug)]
pub struct Microfacet<D: MicrofacetDistribution + Debug> {
    dist: D,
    color: Color,
}

impl<D: MicrofacetDistribution + Debug> Microfacet<D> {
    pub fn new(dist: D, color: Color) -> Self {
        Self { dist, color }
    }
}

impl<R: Rng, D: MicrofacetDistribution + Debug + Sync + Send> Material<R> for Microfacet<D> {
    fn eval(&self, w_in: Dir3, w_out: Dir3, _rng: &mut R) -> ray_tracing_core::prelude::Color {
        if w_out.y() > 0.0 && w_in.y() > 0.0 {
            let w_h = (w_in + w_out).normalize();

            let g = self.dist.g1(w_in, w_h) * self.dist.g1(w_out, w_h);

            self.color * fresnel_real(Dir3::dot(w_in, w_h), 1.0, 1.5) * g * self.dist.d(w_h)
                / (4.0 * w_in.y() * w_out.y()).max(0.0)
        } else {
            Color::black()
        }
    }

    fn sample(&self, w_in: Dir3, rng: &mut R) -> ray_tracing_core::material::SampleResult {
        let w_h = self.dist.sample_d(rng);

        let w_out = (2.0 * (Dir3::dot(w_in, w_h) * w_h) - w_in).normalize();

        if w_out.y() > 0.0 {
            let g = self.dist.g1(w_in, w_h) * self.dist.g1(w_out, w_h);
            ray_tracing_core::material::SampleResult::new(
                w_out,
                self.color * fresnel_real(Dir3::dot(w_in, w_h), 1.0, 1.5) * g,
                false,
            )
        } else {
            ray_tracing_core::material::SampleResult::new(w_out, Color::black(), false)
        }
    }

    fn pdf(&self, w_in: Dir3, w_out: Dir3) -> Float {
        let w_h = (w_in + w_out).normalize();

        self.dist.d(w_h)
    }
}

#[derive(Debug)]
pub struct BeckmannDistribution {
    alpha: Float,
}

impl BeckmannDistribution {
    pub fn new(alpha: Float) -> Self {
        Self { alpha }
    }
}

impl MicrofacetDistribution for BeckmannDistribution {
    fn g1(&self, w: Dir3, w_h: Dir3) -> Float {
        if w.y() > 0.0 && Dir3::dot(w, w_h) > 0.0 {
            let cos_theta = w.y();
            let a = 1.0 / (self.alpha * (Float::sqrt(1.0 - cos_theta * cos_theta) / cos_theta));
            if a < 1.6 {
                (3.535 * a + 2.181 * a * a) / (1.0 + 2.276 * a + 2.577 * a * a)
            } else {
                1.0
            }
        } else {
            0.0
        }
    }

    fn d(&self, w_h: Dir3) -> Float {
        if w_h.y() > 0.0 {
            let cos_theta_squared = w_h.y() * w_h.y();
            let tan_theta_squared =
                w_h.x() * w_h.x() / cos_theta_squared + w_h.z() * w_h.z() / cos_theta_squared;
            Float::exp(-tan_theta_squared / (self.alpha * self.alpha))
                / (FloatConsts::PI
                    * self.alpha
                    * self.alpha
                    * cos_theta_squared
                    * cos_theta_squared)
        } else {
            0.0
        }
    }

    fn sample_d<R: Rng>(&self, rng: &mut R) -> Dir3 {
        let n = Normal::new(0.0, self.alpha).unwrap();

        let mx = n.sample(rng);
        let mz = n.sample(rng);

        let y = Float::sqrt(1.0 / (mx * mx + 1.0 + mz * mz));

        Dir3::new(mx * y, y, mz * y).normalize()
    }
}

pub trait MicrofacetDistribution {
    fn g1(&self, w: Dir3, w_h: Dir3) -> Float;

    fn d(&self, w_h: Dir3) -> Float;

    fn sample_d<R: Rng>(&self, rng: &mut R) -> Dir3;
}
