use core::f64;

use plotters::{
    chart::ChartBuilder,
    coord::Shift,
    prelude::{BitMapBackend, DrawingArea, DrawingBackend, IntoDrawingArea},
    style::WHITE,
};
use rand::{rngs::SmallRng, SeedableRng};
use ray_tracing_core::prelude::*;
use ray_tracing_material::{
    diffuse::DiffuseMaterial,
    iridescent::Iridescent,
    microfacet::{BeckmannDistribution, Microfacet},
    mirror::Mirror,
    oren_nayar::OrenNayar,
};
use rayon::iter::{
    IntoParallelIterator, IntoParallelRefIterator, ParallelExtend, ParallelIterator,
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let w_in = Dir3::new(1.0, 1.0, 0.0).normalize();

    let color = Color::new(1.0, 1.0, 1.0);

    let m = Mirror::new(color);
    generate_chart("mirror.png", &m, 2, w_in)?;

    let m = DiffuseMaterial::new(color);
    generate_chart("diffuse.png", &m, 100, w_in)?;

    let m = Microfacet::new(BeckmannDistribution::new(0.1), color);
    generate_chart("microfacet.png", &m, 100, w_in)?;

    let m = OrenNayar::new(0.5, color);
    generate_chart("oren-nayar.png", &m, 100, w_in)?;

    let m = Iridescent::new(300.0, 300.0, 1.0, 1.5, 10);
    generate_chart("iridescent.png", &m, 2, w_in)?;

    Ok(())
}

fn generate_chart<M: Material<SmallRng>>(
    path: impl AsRef<std::path::Path>,
    m: &M,
    samples_per_pixel: usize,
    w_in: Dir3,
) -> Result<(), Box<dyn std::error::Error>> {
    let inner_size = 400;
    let width = (inner_size + 70) * 3;
    let height = ((inner_size + 97) * 2 + 32) * 2;
    let root = BitMapBackend::new(&path, (width, height)).into_drawing_area();
    root.fill(&WHITE)?;

    let executions = inner_size as usize * inner_size as usize * samples_per_pixel;

    let areas = root.split_evenly((2, 1));
    let eval_histogram = create_histogram_evaled(m, 400, 400, executions, w_in);
    let sample_histogram = create_historgram_sampled(m, 400, 400, executions, w_in);

    let max = Float::max(
        Float::max(eval_histogram.max(), sample_histogram.max()),
        Float::MIN_POSITIVE,
    );

    let eval_energy = eval_histogram.energy();
    let area = areas[0].titled(
        &format!(
            "Evaled ({:.4}, {:.4}, {:.4})",
            eval_energy.r(),
            eval_energy.g(),
            eval_energy.b()
        ),
        ("sans-serif", 30),
    )?;
    plot_material(&area, eval_histogram, max)?;

    let sample_energy = sample_histogram.energy();
    let area = areas[1].titled(
        &format!(
            "Evaled ({:.4}, {:.4}, {:.4})",
            sample_energy.r(),
            sample_energy.g(),
            sample_energy.b()
        ),
        ("sans-serif", 30),
    )?;
    plot_material(&area, sample_histogram, max)?;

    root.present()?;

    Ok(())
}

fn plot_material<DB: DrawingBackend>(
    root: &DrawingArea<DB, Shift>,
    histogram: Histogram,
    max: Float,
) -> Result<(), Box<dyn std::error::Error>>
where
    DB::ErrorType: 'static,
{
    let areas = root.split_evenly((2, 3));
    for (channel, upper, lower) in (0..3).map(|i| (i, &areas[i], &areas[i + 3])) {
        let channel_name = match channel {
            0 => "red",
            1 => "green",
            2 => "blue",
            _ => unreachable!(),
        };

        let mut upper_chart = ChartBuilder::on(upper)
            .margin(5)
            .x_label_area_size(30)
            .top_x_label_area_size(30)
            .y_label_area_size(30)
            .right_y_label_area_size(30)
            .caption(
                format!("Upper Hemisphere {channel_name}"),
                ("sans-serif", 20),
            )
            .build_cartesian_2d(-1.0..1.0, -1.0..1.0)?;

        upper_chart
            .configure_mesh()
            .disable_x_mesh()
            .disable_y_mesh()
            .draw()?;

        let upper_area = upper_chart.plotting_area();

        let mut lower_chart = ChartBuilder::on(lower)
            .margin(5)
            .x_label_area_size(30)
            .top_x_label_area_size(30)
            .y_label_area_size(30)
            .right_y_label_area_size(30)
            .caption(
                format!("Lower Hemisphere {channel_name}"),
                ("sans-serif", 20),
            )
            .build_cartesian_2d(-1.0..1.0, -1.0..1.0)?;

        lower_chart
            .configure_mesh()
            .disable_x_mesh()
            .disable_y_mesh()
            .draw()?;

        let lower_area = lower_chart.plotting_area();

        let (xrange, yrange) = upper_area.get_pixel_range();
        let width = (xrange.end - xrange.start) as usize;
        let height = (yrange.end - yrange.start) as usize;
        assert_eq!(width, histogram.width);
        assert_eq!(height, histogram.height);

        let (xrange, yrange) = lower_area.get_pixel_range();
        let width = (xrange.end - xrange.start) as usize;
        let height = (yrange.end - yrange.start) as usize;
        assert_eq!(width, histogram.width);
        assert_eq!(height, histogram.height);

        for (dir, f) in histogram.iter() {
            let f = match channel {
                0 => f.r(),
                1 => f.g(),
                2 => f.b(),
                _ => unreachable!(),
            };

            assert!(f >= 0.0 && f <= max && f.is_finite() && max != 0.0, "{f}");
            let color = plotters::prelude::ViridisRGB::get_color_normalized(f, 0.0, max);
            if dir.y() < 0.0 {
                lower_area.draw_pixel((dir.x() as f64, dir.z() as f64), &color)?;
            } else {
                upper_area.draw_pixel((dir.x() as f64, dir.z() as f64), &color)?;
            }
        }
    }

    Ok(())
}

struct Histogram {
    data: Vec<Color>,
    inserted: usize,
    width: usize,
    height: usize,
    max: f64,
}

impl Histogram {
    fn new(width: usize, height: usize) -> Self {
        Self {
            data: vec![Color::black(); 2 * width * height],
            inserted: 0,
            width,
            height,
            max: 0.0,
        }
    }

    fn discretize(&self, p: Dir3) -> usize {
        assert!(p.x() >= -1.0);
        assert!(p.x() <= 1.0);
        assert!(p.y() >= -1.0);
        assert!(p.y() <= 1.0);
        assert!(p.z() >= -1.0);
        assert!(p.z() <= 1.0);
        (((p.x() + 1.0) * 0.5 * self.width as f32).floor() as usize).min(self.width - 1)
            + (((p.z() + 1.0) * 0.5 * self.height as f32).floor() as usize).min(self.height - 1)
                * self.width
            + if p.y() < 0.0 {
                self.width * self.height
            } else {
                0
            }
    }

    fn insert(&mut self, pos: Dir3, value: Color) {
        let index = self.discretize(pos);
        self.data[index] += value;
        self.max = f64::max(self.max, self.data[index].r() as f64);
        self.max = f64::max(self.max, self.data[index].g() as f64);
        self.max = f64::max(self.max, self.data[index].b() as f64);
        self.inserted += 1;
    }

    fn max(&self) -> f32 {
        dbg!(self.max, self.inserted, self.inserted as f64);
        (self.max / (self.inserted as f64)) as f32
    }

    fn iter(&self) -> impl Iterator<Item = (Dir3, Color)> + use<'_> {
        let m = 1f32 / self.inserted as f32;
        let width = self.width;
        let height = self.height;
        self.data
            .iter()
            .enumerate()
            .map(move |(i, &f)| {
                if i < width * height {
                    let x = i % width;
                    let z = i / width;

                    let x = 2.0 * (x as f32 + 0.5) / (width as f32) - 1.0;
                    let z = 2.0 * (z as f32 + 0.5) / (height as f32) - 1.0;
                    let y = f32::sqrt(1.0 - x * x - z * z);

                    (Dir3::new(x, y, z), f * m)
                } else {
                    let x = i % width;
                    let z = (i - width * height) / width;

                    let x = 2.0 * (x as f32 + 0.5) / (width as f32) - 1.0;
                    let z = 2.0 * (z as f32 + 0.5) / (height as f32) - 1.0;
                    let y = -f32::sqrt(1.0 - x * x - z * z);

                    (Dir3::new(x, y, z), f * m)
                }
            })
            .filter(|(dir, _)| dir.x() * dir.x() + dir.z() * dir.z() < 1.0)
    }

    fn energy(&self) -> Color {
        let total = self.data.par_iter().copied().sum::<Color>();

        total / (self.inserted as Float)
    }
}

impl ParallelExtend<(Dir3, Color)> for Histogram {
    fn par_extend<I>(&mut self, par_iter: I)
    where
        I: IntoParallelIterator<Item = (Dir3, Color)>,
    {
        let (tx, rx) = std::sync::mpsc::sync_channel(100);

        std::thread::scope(move |s| {
            s.spawn(move || {
                while let Ok((i, f)) = rx.recv() {
                    self.insert(i, f)
                }
            });

            {
                let tx = tx;
                par_iter.into_par_iter().for_each(|t| {
                    let _ = tx.send(t);
                });
            }
        });
    }
}

fn create_historgram_sampled<M: Material<SmallRng>>(
    m: &M,
    width: usize,
    height: usize,
    executions: usize,
    w_in: Dir3,
) -> Histogram {
    let mut h = Histogram::new(width, height);
    h.par_extend(
        (0..executions)
            .into_par_iter()
            .map_init(SmallRng::from_entropy, |rng, _| {
                let sample = m.sample(w_in, rng);

                assert!(
                    sample.color().r() >= 0.0
                        && sample.color().r().is_finite()
                        && sample.color().g() >= 0.0
                        && sample.color().g().is_finite()
                        && sample.color().b() >= 0.0
                        && sample.color().b().is_finite()
                        && (-1.0..=1.0).contains(&sample.w_out().x())
                        && (-1.0..=1.0).contains(&sample.w_out().y())
                        && (-1.0..=1.0).contains(&sample.w_out().z()),
                    "invalid_sampling: w_in: {w_in:?}; w_out: {:?}; material: {m:?}; color: {:?}",
                    sample.w_out(),
                    sample.color()
                );

                (sample.w_out(), sample.color())
            }),
    );

    h
}

fn create_histogram_evaled<M: Material<SmallRng>>(
    m: &M,
    width: usize,
    height: usize,
    executions: usize,
    w_in: Dir3,
) -> Histogram {
    let mut h = Histogram::new(width, height);

    h.par_extend(
        (0..executions)
            .into_par_iter()
            .map_init(SmallRng::from_entropy, |rng, _| {
                let w_out = Dir3::sample_uniform_sphere(rng);

                let color = m.eval(w_in, w_out, rng) * w_out.y().abs() * 4.0 * FloatConsts::PI;

                assert!(
                    color.r() >= 0.0
                        && color.r().is_finite()
                        && color.g() >= 0.0
                        && color.g().is_finite()
                        && color.b() >= 0.0
                        && color.b().is_finite(),
                    "invalid_eval: w_in: {w_in:?}; w_out: {w_out:?}; material: {m:?}; color: {color:?}"
                );

                (w_out, color)
            }),
    );

    h
}
