use ray_tracing_core::{camera::Camera, prelude::*, renderer::ClassicalRenderer, scene::Scene};

pub struct MIS {
    width: u32,
    height: u32,
}

impl MIS {
    pub fn new(width: u32, height: u32) -> Self {
        Self { width, height }
    }
}

impl<S, C, R> ClassicalRenderer<R, S, C> for MIS
where
    S: Scene<R>,
    C: Camera<R>,
    R: Rng,
{
    fn render_pixel(&self, scene: &S, camera: &C, x: u32, y: u32, rng: &mut R) -> Color {
        let mut sum = Color::black();
        let mut alpha = Color::white();

        let mut r = camera.forward(x, y, rng);

        let mut count = 0;
        let mut last_bsdf_pdf = 0.0;

        while let Some(i) = scene.intersect(r, 0.001, Float::INFINITY) {
            let frame = i.tangent_frame();

            let intersect_pos = r.at(i.t());

            let w_in = frame.to_frame(-r.dir());

            if let Some(light) = i.light() {
                if count == 0 {
                    sum += alpha * light.emit(w_in, rng);
                } else {
                    let dist = i.t() / r.dir().length();

                    let path_pdf = last_bsdf_pdf;
                    let nee_pdf = i.light_pdf() * dist * dist / w_in.y();

                    let b = path_pdf / (path_pdf + nee_pdf);
                    sum += b * alpha * light.emit(w_in, rng);
                }
            }

            let w_out = if let Some(material) = i.material() {
                let sample_result = material.sample(w_in, rng);

                if let Some(l) = scene.sample_light(w_in, &i, rng) {
                    let light_frame = l.tangent_frame();

                    let light_dir = l.pos() - intersect_pos;

                    let light_ray = Ray::new(intersect_pos, light_dir.normalize(), r.time());

                    if scene
                        .intersect(light_ray, 0.001, light_dir.length() - 0.001)
                        .is_none()
                    {
                        let g = (Dir3::dot(i.normal(), light_dir.normalize())
                            * Dir3::dot(l.normal(), -light_dir.normalize()))
                            / light_dir.length_squared();

                        // calculate MIS wheight
                        let path_pdf = material.pdf(w_in, frame.to_frame(light_dir));
                        let nee_pdf = l.pdf() * light_dir.length_squared() / Float::abs(w_in.y());

                        let b = nee_pdf / (path_pdf + nee_pdf);

                        sum += b
                            * alpha
                            * g
                            * material.eval(w_in, frame.to_frame(light_dir.normalize()), rng)
                            * l.light()
                                .emit(light_frame.to_frame(-light_dir.normalize()), rng)
                            / l.pdf();
                    }
                }

                if sample_result.is_delta() {
                    last_bsdf_pdf = 100000.0;
                } else {
                    last_bsdf_pdf = material.pdf(w_in, sample_result.w_out());
                }
                alpha *= sample_result.color();
                sample_result.w_out()
            } else {
                return sum;
            };

            r = Ray::new(r.at(i.t()), frame.to_world(w_out), r.time());
            count += 1;
            if count > 4 {
                break;
            }
        }

        sum
    }

    fn width(&self) -> u32 {
        self.width
    }

    fn height(&self) -> u32 {
        self.height
    }
}