/// Implementation based on blog post @
/// https://jacco.ompf2.com/2022/04/13/how-to-build-a-bvh-part-1-basics/
use std::f32::EPSILON;
use std::fmt;

use log::info;
use stl::STL;

use crate::{
    aabb::AABB,
    hitable::{Hit, HitRecord},
    material::Material,
    ray::Ray,
    vec3::{cross, dot, Vec3},
};

#[derive(Debug, PartialEq)]
struct BVHNode {
    aabb: AABB,
    // When tri_count==0, left_first holds the left child's index in bvh_nodes. When >0 left_first
    // holds the index for the first triangle in triangles.
    left_first: u32,
    tri_count: u32,
}

impl BVHNode {
    fn is_leaf(&self) -> bool {
        self.tri_count > 0
    }

    fn cost(&self) -> f32 {
        let area = self.aabb.area();
        self.tri_count as f32 * area
    }
}

#[derive(Clone, PartialEq)]
pub struct Triangle {
    centroid: Vec3,
    verts: [Vec3; 3],
}
impl fmt::Debug for Triangle {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Tri: <{}, {}, {}> @ {}",
            self.verts[0], self.verts[1], self.verts[2], self.centroid
        )
    }
}

pub struct BVHTriangles<M>
where
    M: Material,
{
    pub triangles: Vec<Triangle>,
    triangle_index: Vec<usize>,
    material: M,
    bvh_nodes: Vec<BVHNode>,
}

#[derive(Debug, Default)]
struct Bin {
    bounds: AABB,
    tri_count: usize,
}

impl<M> fmt::Debug for BVHTriangles<M>
where
    M: Material,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} triangles", self.triangles.len())?;
        if f.alternate() {
            writeln!(f)?;
        }
        for (i, t) in self.triangles.iter().enumerate() {
            if f.alternate() {
                write!(f, "\t")?;
            }
            write!(f, "{i:>3} {:?} ", t)?;
            if f.alternate() {
                writeln!(f)?;
            }
        }
        if f.alternate() {
            writeln!(f)?;
        }
        for (i, n) in self.bvh_nodes.iter().enumerate() {
            write!(f, "N[{i}] {n:?}")?;
            if f.alternate() {
                writeln!(f)?;
            }
            let n = &self.bvh_nodes[i];
            if n.is_leaf() {
                for t_idx in n.left_first..(n.left_first + n.tri_count) {
                    if f.alternate() {
                        write!(f, "\t")?;
                    }
                    write!(f, "{:?} ", self.triangles[t_idx as usize])?;
                    if f.alternate() {
                        writeln!(f)?;
                    }
                }
            }
        }
        Ok(())
    }
}

#[derive(Debug)]
struct SplitCost {
    pos: f32,
    axis: usize,
    cost: f32,
}

const ROOT_NODE_IDX: usize = 0;
impl<M> BVHTriangles<M>
where
    M: Material,
{
    pub fn new(stl: &STL, material: M, scale_factor: f32) -> BVHTriangles<M> {
        let now = std::time::Instant::now();

        assert_eq!(std::mem::size_of::<BVHNode>(), 32);
        let div3 = 1. / 3.;
        let triangles: Vec<_> = stl
            .triangles
            .iter()
            .map(|t| {
                let v0 = t.verts[0] * scale_factor;
                let v1 = t.verts[1] * scale_factor;
                let v2 = t.verts[2] * scale_factor;
                let centroid = (v0 + v1 + v2) * div3;
                Triangle {
                    centroid,
                    verts: [v0, v1, v2],
                }
            })
            .collect();
        let triangle_index = (0..triangles.len()).collect();

        let n = 2 * triangles.len() - 2;
        let bvh_nodes = Vec::with_capacity(n);
        let mut bvh = BVHTriangles {
            triangles,
            triangle_index,
            bvh_nodes,
            material,
        };
        bvh.build_bvh();
        info!(
            "BVHTriangles build time {:0.3}s",
            now.elapsed().as_secs_f32()
        );
        struct Stats {
            nodes: usize,
            leafs: usize,
            min_tris: u32,
            max_tris: u32,
        }
        impl Default for Stats {
            fn default() -> Self {
                Stats {
                    nodes: 0,
                    leafs: 0,
                    min_tris: u32::MAX,
                    max_tris: u32::MIN,
                }
            }
        }
        let stats = bvh.bvh_nodes.iter().fold(Stats::default(), |mut stats, n| {
            stats.nodes += 1;
            if n.is_leaf() {
                stats.leafs += 1;
                stats.min_tris = n.tri_count.min(stats.min_tris);
                stats.max_tris = n.tri_count.max(stats.max_tris);
            }
            stats
        });
        info!("BVHTriangles build stats:");
        info!("  Nodes:   {}", stats.nodes);
        info!("  Leaves:  {}", stats.leafs);
        info!("  Tris:    {}", bvh.triangles.len());
        info!("  Min Tri: {}", stats.min_tris);
        info!("  Max Tri: {}", stats.max_tris);
        info!("  Avg Tri: {}", bvh.triangles.len() / stats.leafs);
        info!("  Predict: {}", bvh.bvh_nodes.capacity());
        info!("  Actual:  {}", bvh.bvh_nodes.len());
        info!(
            "  Savings:  {} bytes",
            (bvh.bvh_nodes.capacity() - bvh.bvh_nodes.len()) * std::mem::size_of::<BVHNode>()
        );
        bvh.bvh_nodes.shrink_to_fit();
        //dbg!(&bvh);
        bvh
    }

    fn build_bvh(&mut self) {
        // assign all triangles to root node
        let root = BVHNode {
            aabb: AABB::default(),
            left_first: 0,
            tri_count: self.triangles.len() as u32,
        };
        self.bvh_nodes.push(root);
        self.update_node_bounds(ROOT_NODE_IDX);
        // subdivide recursively
        self.subdivide(ROOT_NODE_IDX);
    }
    fn update_node_bounds(&mut self, node_idx: usize) {
        let node = &mut self.bvh_nodes[node_idx];
        let mut aabb_min: Vec3 = f32::MAX.into();
        let mut aabb_max: Vec3 = f32::MIN.into();
        for i in node.left_first..(node.left_first + node.tri_count) {
            let leaf_tri_idx = self.triangle_index[i as usize];
            let leaf_tri = &self.triangles[leaf_tri_idx];
            aabb_min = vec3::min(aabb_min, leaf_tri.verts[0]);
            aabb_min = vec3::min(aabb_min, leaf_tri.verts[1]);
            aabb_min = vec3::min(aabb_min, leaf_tri.verts[2]);
            aabb_max = vec3::max(aabb_max, leaf_tri.verts[0]);
            aabb_max = vec3::max(aabb_max, leaf_tri.verts[1]);
            aabb_max = vec3::max(aabb_max, leaf_tri.verts[2]);
        }
        node.aabb = AABB::new(aabb_min, aabb_max);
    }

    fn find_best_split_plane(&self, node: &BVHNode) -> SplitCost {
        let mut best = SplitCost {
            pos: 0.,
            cost: f32::MAX,
            axis: usize::MAX,
        };

        for axis in 0..3 {
            let mut bounds_min = f32::MAX;
            let mut bounds_max = f32::MIN;

            for i in 0..node.tri_count {
                let triangle = &self.triangles[self.triangle_index[(node.left_first + i) as usize]];
                bounds_min = bounds_min.min(triangle.centroid[axis]);
                bounds_max = bounds_max.max(triangle.centroid[axis]);
            }

            if bounds_min == bounds_max {
                continue;
            }

            const NUM_BINS: usize = 8;
            let mut bins: Vec<_> = (0..NUM_BINS)
                .map(|_| Bin {
                    bounds: AABB::infinite(),
                    tri_count: 0,
                })
                .collect();
            let scale = bins.len() as f32 / (bounds_max - bounds_min);
            // populate the bins
            for i in 0..node.tri_count {
                let triangle = &self.triangles[self.triangle_index[(node.left_first + i) as usize]];
                let bin_idx = (triangle.centroid[axis] - bounds_min) * scale;
                let bin_idx = (bins.len() - 1).min(bin_idx as usize);

                bins[bin_idx].tri_count += 1;
                bins[bin_idx].bounds.grow(triangle.verts[0]);
                bins[bin_idx].bounds.grow(triangle.verts[1]);
                bins[bin_idx].bounds.grow(triangle.verts[2]);
            }

            // gather data for the 7 planes between the 8 bins
            let mut left_area: Vec<_> = (0..bins.len() - 1).map(|_| 0.).collect();
            let mut right_area: Vec<_> = (0..bins.len() - 1).map(|_| 0.).collect();
            let mut left_count: Vec<_> = (0..bins.len() - 1).map(|_| 0).collect();
            let mut right_count: Vec<_> = (0..bins.len() - 1).map(|_| 0).collect();
            let mut left_box = AABB::infinite();
            let mut right_box = AABB::infinite();
            let mut left_sum = 0;
            let mut right_sum = 0;
            for i in 0..(bins.len() - 1) {
                left_sum += bins[i].tri_count;
                left_count[i] = left_sum;
                left_box.grow(bins[i].bounds.min());
                left_box.grow(bins[i].bounds.max());
                left_area[i] = left_box.area();

                right_sum += bins[bins.len() - 1 - i].tri_count;
                right_count[bins.len() - 2 - i] = right_sum;
                right_box.grow(bins[bins.len() - 1 - i].bounds.min());
                right_box.grow(bins[bins.len() - 1 - i].bounds.max());
                right_area[bins.len() - 2 - i] = right_box.area();
            }

            let scale = (bounds_max - bounds_min) / bins.len() as f32;
            // calculate SAH cost for the 7 planes
            for i in 0..(bins.len() - 1) {
                let plane_cost =
                    left_count[i] as f32 * left_area[i] + right_count[i] as f32 * right_area[i];
                if plane_cost < best.cost {
                    best.axis = axis;
                    best.pos = bounds_min + scale * (i as f32 + 1.);
                    best.cost = plane_cost;
                }
            }
        }
        best
    }

    fn subdivide(&mut self, idx: usize) {
        let (left_first, tri_count, left_count, i) = {
            let node = &self.bvh_nodes[idx];
            let split = self.find_best_split_plane(&node);
            let no_split_cost = node.cost();
            // Stop subdividing if it isn't getting any better.
            if split.cost >= no_split_cost {
                return;
            }

            // Split the group in two halves.
            let mut i = node.left_first as isize;
            let mut j = i + node.tri_count as isize - 1;
            while i <= j {
                if self.triangles[self.triangle_index[i as usize]].centroid[split.axis] < split.pos
                {
                    i += 1;
                } else {
                    self.triangles.swap(
                        self.triangle_index[i as usize],
                        self.triangle_index[j as usize],
                    );
                    j -= 1;
                }
            }

            // Create child nodes for each half.
            let left_count = i as u32 - node.left_first;
            if left_count == 0 || left_count == node.tri_count {
                return;
            }
            (node.left_first, node.tri_count, left_count, i)
        };

        // create child nodes
        let left_child_idx = self.bvh_nodes.len() as u32;
        let right_child_idx = left_child_idx + 1 as u32;
        let left = BVHNode {
            aabb: AABB::default(),
            left_first,
            tri_count: left_count as u32,
        };
        let right = BVHNode {
            aabb: AABB::default(),
            left_first: i as u32,
            tri_count: (tri_count - left_count) as u32,
        };
        self.bvh_nodes.push(left);
        self.bvh_nodes.push(right);
        let node = &mut self.bvh_nodes[idx];
        node.left_first = left_child_idx;
        node.tri_count = 0;

        // Recurse
        self.update_node_bounds(left_child_idx as usize);
        self.subdivide(left_child_idx as usize);
        self.update_node_bounds(right_child_idx as usize);
        self.subdivide(right_child_idx as usize);
    }

    fn intersect_bvh(&self, r: Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        let mut node = &self.bvh_nodes[ROOT_NODE_IDX];
        let mut stack = Vec::with_capacity(2);
        let mut nearest = None;
        loop {
            if node.is_leaf() {
                let canditate = (node.left_first..(node.left_first + node.tri_count))
                    // Map from idx to Triangle
                    .map(|idx| &self.triangles[self.triangle_index[idx as usize]])
                    // Try to hit all triangles for this node, filtering out misses.
                    .filter_map(|tri| intersect_tri(r, &tri))
                    // Find the nearest hit (if any).
                    .min_by(|a, b| a.t.partial_cmp(&b.t).unwrap());

                // merge candidate with nearest.
                nearest = match (&canditate, &nearest) {
                    (Some(_), None) => canditate,
                    (None, Some(_)) => nearest,
                    (Some(c), Some(n)) => {
                        //info!("merging {c:#?} and {n:#?}");
                        if c.t < n.t {
                            canditate
                        } else {
                            nearest
                        }
                    }
                    (None, None) => None,
                };
                if stack.is_empty() {
                    break;
                }
                node = stack.pop().unwrap();
                continue;
            }

            let child1 = &self.bvh_nodes[node.left_first as usize];
            let child2 = &self.bvh_nodes[node.left_first as usize + 1];
            let dist1 = child1.aabb.hit_distance(r, t_min, t_max);
            let dist2 = child2.aabb.hit_distance(r, t_min, t_max);
            // Swap c1/c2 & d1/d2 based on d1/d2.
            let (child1, child2, dist1, dist2) = match (dist1, dist2) {
                (Some(d1), Some(d2)) if d1 > d2 => (child2, child1, dist2, dist1),
                (None, Some(_)) => (child2, child1, dist2, dist1),
                _ => (child1, child2, dist1, dist2),
            };

            // dist1/child1 should now be the nearest hit.

            // If we missed dist1/child1, then we implicitly missed dist2/child2, so pop a child
            // from the stack or exit the function.
            if dist1.is_none() {
                if stack.is_empty() {
                    break;
                }
                node = stack.pop().unwrap();
            } else {
                // We hit child1, so process it next.
                node = child1;
                // If we also hit child2 save it on the stack so we can process it later.
                if dist2.is_some() {
                    stack.push(child2);
                }
            }
        }
        nearest
            .and_then(|rtr: RayTriangleResult| Some(rtr.hit_record_with_material(&self.material)))
    }
}

/// Based on https://www.scratchapixel.com/lessons/3d-basic-rendering/ray-tracing-rendering-a-triangle/moller-trumbore-ray-triangle-intersection.html
fn intersect_tri(r: Ray, tri: &Triangle) -> Option<RayTriangleResult> {
    // #ifdef MOLLER_TRUMBORE
    //     Vec3f v0v1 = v1 - v0;
    //     Vec3f v0v2 = v2 - v0;
    //     Vec3f pvec = dir.crossProduct(v0v2);
    //     float det = v0v1.dotProduct(pvec);
    // #ifdef CULLING
    //     // if the determinant is negative, the triangle is 'back facing'
    //     // if the determinant is close to 0, the ray misses the triangle
    //     if (det < kEpsilon) return false;
    // #else
    //     // ray and triangle are parallel if det is close to 0
    //     if (fabs(det) < kEpsilon) return false;
    // #endif
    //     float invDet = 1 / det;
    //
    //     Vec3f tvec = orig - v0;
    //     u = tvec.dotProduct(pvec) * invDet;
    //     if (u < 0 || u > 1) return false;
    //
    //     Vec3f qvec = tvec.crossProduct(v0v1);
    //     v = dir.dotProduct(qvec) * invDet;
    //     if (v < 0 || u + v > 1) return false;
    //
    //     t = v0v2.dotProduct(qvec) * invDet;
    //
    let v0 = tri.verts[0];
    let v1 = tri.verts[1];
    let v2 = tri.verts[2];

    let v0v1 = v1 - v0;
    let v0v2 = v2 - v0;
    let p = cross(r.direction, v0v2);
    let det = dot(v0v1, p);
    if det.abs() < EPSILON {
        return None;
    }

    let inv_det = 1. / det;

    let t = r.origin - v0;
    let u = dot(t, p) * inv_det;
    if u < 0. || u > 1. {
        return None;
    }

    let q = cross(t, v0v1);
    let v = dot(r.direction, q) * inv_det;
    if v < 0. || u + v > 1. {
        return None;
    }
    let t = dot(v0v2, q) * inv_det;

    if t > EPSILON {
        return Some(RayTriangleResult {
            t,
            p: r.point_at_parameter(t),
            tri: tri.clone(),
        });
    }
    None
}

impl<M> Hit for BVHTriangles<M>
where
    M: Material,
{
    fn hit(&self, r: Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        self.intersect_bvh(r, t_min, t_max)
    }

    fn bounding_box(&self, _t_min: f32, _t_max: f32) -> Option<AABB> {
        Some(self.bvh_nodes[ROOT_NODE_IDX].aabb)
    }
}

#[derive(Debug)]
struct RayTriangleResult {
    t: f32,
    p: Vec3,
    tri: Triangle,
}

impl RayTriangleResult {
    fn hit_record_with_material<'m>(self, material: &'m dyn Material) -> HitRecord<'m> {
        // We don't support UV (yet?).
        let uv = (0.5, 0.5);
        let v0 = self.tri.verts[0];
        let v1 = self.tri.verts[1];
        let v2 = self.tri.verts[2];

        let v0v1 = v1 - v0;
        let v0v2 = v2 - v0;
        let normal = cross(v0v1, v0v2).unit_vector();
        //println!("hit triangle {tri:?}");
        HitRecord {
            t: self.t,
            uv,
            p: self.p,
            normal,
            material,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        bvh_triangles::BVHTriangles,
        cuboid::Cuboid,
        hitable::Hit,
        material::{Dielectric, Lambertian},
        ray::Ray,
        texture::ConstantTexture,
    };
    use pretty_assertions::assert_eq;
    use proptest::prelude::*;
    use std::{
        io::{BufReader, Cursor},
        sync::Arc,
    };
    use stl::STL;

    #[test]
    fn compare_cuboid() {
        let c = Cuboid::new(
            [0., 0., 0.].into(),
            [20., 20., 20.].into(),
            Arc::new(Dielectric::new(1.5)),
        );
        let stl_cube = STL::parse(
            BufReader::new(Cursor::new(include_bytes!("../stls/cube.stl"))),
            false,
        )
        .expect("failed to parse cube");

        let s = BVHTriangles::new(
            &stl_cube,
            Dielectric::new(1.5),
            //Metal::new(Vec3::new(0.8, 0.8, 0.8), 0.2),
            //Lambertian::new(ConstantTexture::new(Vec3::new(1.0, 0.2, 0.2))),
        );
        //dbg!(&s);
        let mut rays: Vec<_> = (1..20)
            .flat_map(|y| {
                (1..20).flat_map(move |x| {
                    let x = x as f32;
                    let y = y as f32;
                    vec![
                        // Outward in angle
                        Ray::new([-1., x, y], [1., 0., 0.], 0.),
                        Ray::new([21., x, y], [-1., 0., 0.], 0.),
                        Ray::new([x, -1., y], [0., 1., 0.], 0.),
                        Ray::new([x, 21., y], [0., -1., 0.], 0.),
                        Ray::new([x, y, -1.], [0., 0., 1.], 0.),
                        Ray::new([x, y, 21.], [0., 0., -1.], 0.),
                        // Inward out (
                        Ray::new([x, y, 10.], [1., 0., 0.], 0.),
                        Ray::new([x, y, 10.], [-1., 0., 0.], 0.),
                        Ray::new([x, 10., y], [1., 0., 0.], 0.),
                        Ray::new([x, 10., y], [-1., 0., 0.], 0.),
                        Ray::new([10., x, y], [1., 0., 0.], 0.),
                        Ray::new([10., x, y], [-1., 0., 0.], 0.),
                        Ray::new([x, y, 10.], [0., 1., 0.], 0.),
                        Ray::new([x, y, 10.], [0., -1., 0.], 0.),
                        Ray::new([x, 10., y], [0., 1., 0.], 0.),
                        Ray::new([x, 10., y], [0., -1., 0.], 0.),
                        Ray::new([10., x, y], [0., 1., 0.], 0.),
                        Ray::new([10., x, y], [0., -1., 0.], 0.),
                        Ray::new([x, y, 10.], [0., 0., 1.], 0.),
                        Ray::new([x, y, 10.], [0., 0., -1.], 0.),
                        Ray::new([x, 10., y], [0., 0., 1.], 0.),
                        Ray::new([x, 10., y], [0., 0., -1.], 0.),
                        Ray::new([10., x, y], [0., 0., 1.], 0.),
                        Ray::new([10., x, y], [0., 0., -1.], 0.),
                    ]
                    .into_iter()
                })
            })
            .collect();
        if false {
            // Outward in at an angle.
            let sqrt2 = 2f32.sqrt();
            rays.extend(vec![
                Ray::new([-1., 10., 10.], [sqrt2, sqrt2, 0.], 0.),
                Ray::new([-1., 10., 10.], [sqrt2, -sqrt2, 0.], 0.),
                Ray::new([-1., 10., 10.], [sqrt2, 0., sqrt2], 0.),
                Ray::new([-1., 10., 10.], [sqrt2, 0., -sqrt2], 0.),
            ]);
        }

        for r in rays.into_iter() {
            let c_hit = c
                .hit(r, 0., f32::MAX)
                .expect(&format!("c_hit missed {r:#?}"));
            let s_hit = s
                .hit(r, 0., f32::MAX)
                .expect(&format!("s_hit missed {r:#?}"));
            assert!(
                (c_hit.t - s_hit.t).abs() < EPSILON,
                "{r:?} [t] c_hit: {c_hit:#?}, s_hit: {s_hit:#?}"
            );
            // uv isn't valid for BVHTriangles.
            // assert_eq!( c_hit.uv, s_hit.uv, "{i}: [uv] c_hit: {c_hit:?}, s_hit: {s_hit:?}");
            assert_eq!(
                c_hit.p, s_hit.p,
                "{r:?}: [p] c_hit: {c_hit:#?}, s_hit: {s_hit:#?}"
            );
            assert_eq!(
                c_hit.normal, s_hit.normal,
                "{r:?}: [normal] c_hit: {c_hit:?}, s_hit: {s_hit:?}"
            );
        }
    }

    proptest! {
        // TODO(wathiede): make this work.
        //#[test]
        fn compare_cuboid_proptest(
            ox in -20.0f32..40.0,
            oy in -20.0f32..40.0,
            oz in -20.0f32..40.0,
            dx in -1.0f32..1.0,
            dy in -1.0f32..1.0,
            dz in -1.0f32..1.0,
        ) {
            let r = Ray::new([ox,oy,oz].into(), Vec3::new(dx, dy, dz).unit_vector(), 0.5);
            let c = Cuboid::new(
                [0., 0., 0.].into(),
                [20., 20., 20.].into(),
                Arc::new(Dielectric::new(1.5)),
            );
            let stl_cube = STL::parse(
                BufReader::new(Cursor::new(include_bytes!("../stls/cube.stl"))),
                false,
            )
                .expect("failed to parse cube");

            let s = BVHTriangles::new(
                &stl_cube,
                Dielectric::new(1.5),
                //Metal::new(Vec3::new(0.8, 0.8, 0.8), 0.2),
                //Lambertian::new(ConstantTexture::new(Vec3::new(1.0, 0.2, 0.2))),
            );

            let c_hit = c .hit(r, 0., f32::MAX);
            let s_hit = s .hit(r, 0., f32::MAX);

            match (c_hit, s_hit) {
                (Some(_), None)=>assert!(false, "hit cuboid but not STL"),
                (None, Some(_))=>assert!(false, "hit STL but not cuboid"),
                (Some(c_hit), Some(s_hit))=> {
                    assert!(
                        (c_hit.t - s_hit.t).abs() < 0.00001,
                        "{r:?} [t] c_hit: {c_hit:#?}, s_hit: {s_hit:#?}"
                    );
                    // uv isn't valid for BVHTriangles.
                    // assert_eq!( c_hit.uv, s_hit.uv, "{i}: [uv] c_hit: {c_hit:?}, s_hit: {s_hit:?}");
                    assert_eq!(
                        c_hit.p, s_hit.p,
                        "{r:?}: [p] c_hit: {c_hit:#?}, s_hit: {s_hit:#?}"
                    );
                    assert_eq!(
                        c_hit.normal, s_hit.normal,
                        "{r:?}: [normal] c_hit: {c_hit:?}, s_hit: {s_hit:?}"
                    );
                },
                // It's okay if they both miss.
                (None,None)=>(),
            }
        }
    }
}