/// 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
    }
}

#[derive(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>,
}

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(())
    }
}

const ROOT_NODE_IDX: usize = 0;
impl<M> BVHTriangles<M>
where
    M: Material,
{
    pub fn new(stl: &STL, material: M) -> 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];
                let v1 = t.verts[1];
                let v2 = t.verts[2];
                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!("  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:  {}",
            (bvh.bvh_nodes.capacity() - bvh.bvh_nodes.len()) * std::mem::size_of::<BVHNode>()
        );
        bvh.bvh_nodes.shrink_to_fit();
        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 subdivide(&mut self, idx: usize) {
        let node = &self.bvh_nodes[idx];
        let parent_area = node.aabb.area();
        let parent_cost = node.tri_count as f32 * parent_area;

        let (left_first, tri_count, left_count, i) = {
            let node = &self.bvh_nodes[idx];

            let mut best_axis = usize::MAX;
            let mut best_pos = 0.;
            let mut best_cost = f32::MAX;

            for axis in 0..3 {
                for i in 0..node.tri_count {
                    let triangle =
                        &self.triangles[self.triangle_index[(node.left_first + i) as usize]];
                    let candidate_pos = triangle.centroid[axis];
                    let cost = self.evaluate_sah(node, axis, candidate_pos);
                    if cost <= best_cost {
                        best_pos = candidate_pos;
                        best_axis = axis;
                        best_cost = cost;
                    }
                }
            }
            // Stop subdividing if it isn't getting any better.
            if best_cost >= parent_cost {
                return;
            }

            let axis = best_axis;
            let split_pos = best_pos;

            // 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[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 evaluate_sah(&self, node: &BVHNode, axis: usize, pos: f32) -> f32 {
        // determine triangle counts and bounds for this split candidate
        let mut left_box = AABB::infinite();
        let mut right_box = AABB::infinite();
        let mut left_count = 0;
        let mut right_count = 0;
        for i in 0..node.tri_count {
            let triangle = &self.triangles[self.triangle_index[(node.left_first + i) as usize]];
            if triangle.centroid[axis] < pos {
                left_count += 1;
                left_box.grow(triangle.verts[0]);
                left_box.grow(triangle.verts[1]);
                left_box.grow(triangle.verts[2]);
            } else {
                right_count += 1;
                right_box.grow(triangle.verts[0]);
                right_box.grow(triangle.verts[1]);
                right_box.grow(triangle.verts[2]);
            }
        }
        let cost = left_count as f32 * left_box.area() + right_count as f32 * right_box.area();
        if cost > 0. {
            cost
        } else {
            f32::MAX
        }
    }

    fn intersect_bvh(&self, r: Ray, node_idx: u32, t_min: f32, t_max: f32) -> Option<HitRecord> {
        // TODO(wathiede): visit nodes front to back and remove recursion.
        let node = &self.bvh_nodes[node_idx as usize];
        if !node.aabb.hit(r, t_min, t_max) {
            return None;
        }
        //dbg!(&self);

        if node.is_leaf() {
            return self
                .triangles
                .iter()
                .map(|tri| {
                    if let Some(RayTriangleResult { t, p }) = intersect_tri(r, tri) {
                        // We don't support UV (yet?).
                        let uv = (0.5, 0.5);
                        let v0 = tri.verts[0];
                        let v1 = tri.verts[1];
                        let v2 = tri.verts[2];

                        let v0v1 = v1 - v0;
                        let v0v2 = v2 - v0;
                        let normal = cross(v0v1, v0v2).unit_vector();
                        //println!("hit triangle {tri:?}");
                        Some(HitRecord {
                            t,
                            uv,
                            p,
                            normal,
                            material: &self.material,
                        })
                    } else {
                        None
                    }
                })
                .filter_map(|hr| hr)
                .min_by(|a, b| a.t.partial_cmp(&b.t).unwrap());
        } else {
            let r1 = self.intersect_bvh(r, node.left_first, t_min, t_max);
            let r2 = self.intersect_bvh(r, node.left_first + 1, t_min, t_max);
            // Merge results, if both hit, take the one closest to the ray origin (smallest t
            // value).
            match (&r1, &r2) {
                (Some(_), None) => return r1,
                (None, Some(_)) => return r2,
                (None, None) => (),
                (Some(rp1), Some(rp2)) => return if rp1.t < rp2.t { r1 } else { r2 },
            }
        }
        None
    }
}

/// 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),
        });
    }
    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, 0, t_min, t_max)
    }

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

struct RayTriangleResult {
    t: f32,
    p: Vec3,
}

#[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 std::{
        io::{BufReader, Cursor},
        sync::Arc,
    };
    use stl::STL;

    /*
           #[test]
           fn build_bvh() {
           let stl_triangles: Vec<_> = (0..4)
           .flat_map(|y| {
           (0..2).map(move |x| {
           let x = x as f32;
           let y = y as f32;
           stl::Triangle {
           normal: [1., 0., 0.].into(),
           verts: [
           [2. * x + 0., 2. * y + 0., 0.].into(),
           [2. * x + 1., 2. * y + 0., 0.].into(),
           [2. * x + 1., 2. * y + 1., 0.].into(),
           ],
           attr: 0,
           }
           })
           })
           .collect();
           let stl = STL {
           header: [0; 80],
           triangles: stl_triangles,
           };
        /*
        let mut bvh_triangles: Vec<_> = stl_triangles
        .iter()
        .map(|tri| {
        let div3 = 1. / 3.;
        let v0 = tri.verts[0];
        let v1 = tri.verts[1];
        let v2 = tri.verts[2];
        let centroid = (v0 + v1 + v2) * div3;
        Triangle {
        centroid,
        verts: tri.verts,
        }
        })
        .collect();
        bvh_triangles.sort_by(|a, b| a.centroid.y.partial_cmp(&b.centroid.y).unwrap());
        let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
        let bvh_nodes = Default::default();
        let want = BVHTriangles {
        triangles: bvh_triangles,
        bvh_nodes,
        material,
        };
        */
        let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
        let bvh = BVHTriangles::new(&stl, material);
        dbg!(&bvh);
        assert_eq!(bvh.bvh_nodes.len(), 2 * bvh.triangles.len() - 2);
    }
    */

    #[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();
        // These currently differ between STL and cuboid.
        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.),
            ]);
        }

        // TODO(wathiede): proptest this, it's still not perfectly equal when rendering.

        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:?}"
            );
        }
    }
}