use super::{
    Bounds3f, DirectionCone, Float, Interaction, Normal3f, Point2f, Point3f, Point3fi, Ray,
    ShapeIntersection, ShapeSample, ShapeSampleContext, ShapeTrait, SurfaceInteraction,
    TriangleIntersection, TriangleShape, Vector2f, Vector3f,
};
use crate::core::interaction::InteractionTrait;
use crate::core::pbrt::gamma;
use crate::geometry::{Sqrt, Tuple, VectorLike, spherical_triangle_area};
use crate::utils::math::{difference_of_products, square};
use crate::utils::mesh::TriangleMesh;
use crate::utils::sampling::{
    bilinear_pdf, invert_spherical_triangle_sample, sample_bilinear, sample_spherical_triangle,
    sample_uniform_triangle,
};
use std::mem;
use std::sync::{Arc, OnceLock};

pub static TRIANGLE_MESHES: OnceLock<Vec<Arc<TriangleMesh>>> = OnceLock::new();

#[derive(Clone, Copy)]
struct TriangleData {
    vertices: [Point3f; 3],
    uvs: [Point2f; 3],
    normals: Option<[Normal3f; 3]>,
    area: Float,
    normal: Normal3f,
    reverse_orientation: bool,
    transform_swaps_handedness: bool,
}

impl TriangleShape {
    pub const MIN_SPHERICAL_SAMPLE_AREA: Float = 3e-4;
    pub const MAX_SPHERICAL_SAMPLE_AREA: Float = 6.22;

    fn mesh(&self) -> &Arc<TriangleMesh> {
        let meshes = TRIANGLE_MESHES
            .get()
            .expect("Mesh has not been initialized");
        &meshes[self.mesh_ind]
    }

    fn get_data(&self) -> TriangleData {
        let mesh = self.mesh();
        let start = 3 * self.tri_index;
        let indices = &mesh.vertex_indices[start..start + 3];
        let vertices = [mesh.p[indices[0]], mesh.p[indices[1]], mesh.p[indices[2]]];
        let uvs = mesh.uv.as_ref().map_or(
            [
                Point2f::zero(),
                Point2f::new(1.0, 0.0),
                Point2f::new(1.0, 1.0),
            ],
            |uv| [uv[indices[0]], uv[indices[1]], uv[indices[2]]],
        );
        let normals = mesh
            .n
            .as_ref()
            .map(|n| [n[indices[0]], n[indices[1]], n[indices[2]]]);
        let dp1 = vertices[1] - vertices[0];
        let dp2 = vertices[2] - vertices[0];
        let normal = Normal3f::from(dp1.cross(dp2).normalize());
        let area = 0.5 * dp1.cross(dp2).norm();

        TriangleData {
            vertices,
            uvs,
            normals,
            area,
            normal,
            reverse_orientation: mesh.reverse_orientation,
            transform_swaps_handedness: mesh.transform_swaps_handedness,
        }
    }

    fn solid_angle(&self, p: Point3f) -> Float {
        let data = self.get_data();
        let [p0, p1, p2] = data.vertices;
        spherical_triangle_area(
            (p0 - p).normalize(),
            (p1 - p).normalize(),
            (p2 - p).normalize(),
        )
    }

    fn intersect_triangle(&self, ray: &Ray, t_max: Float) -> Option<TriangleIntersection> {
        let data = self.get_data();
        let [p0, p1, p2] = data.vertices;
        if (p2 - p0).cross(p1 - p0).norm_squared() == 0. {
            return None;
        }
        let mut p0t = p0 - Vector3f::from(ray.o);
        let mut p1t = p1 - Vector3f::from(ray.o);
        let mut p2t = p2 - Vector3f::from(ray.o);

        let kz = ray.d.abs().max_component_index();
        let kx = if kz == 3 { 0 } else { kz + 1 };
        let ky = if kz == 3 { 0 } else { kx + 1 };
        let d = ray.d.permute([kx, ky, kz]);
        p0t = p0t.permute([kx, ky, kz]);
        p1t = p1t.permute([kx, ky, kz]);
        p2t = p2t.permute([kx, ky, kz]);
        // Apply shear transformation to translated vertex positions
        let sx = -d.x() / d.z();
        let sy = -d.y() / d.z();
        let sz = 1. / d.z();
        p0t[0] += sx * p0t.z();
        p0t[1] += sy * p0t.z();
        p1t[0] += sx * p1t.z();
        p1t[1] += sy * p1t.z();
        p2t[0] += sx * p2t.z();
        p2t[0] += sy * p2t.z();

        // Compute edge function coefficients e0, e1, and e2
        let mut e0 = difference_of_products(p1t.x(), p2t.y(), p1t.y(), p2t.x());
        let mut e1 = difference_of_products(p2t.x(), p0t.y(), p2t.y(), p0t.x());
        let mut e2 = difference_of_products(p0t.x(), p1t.y(), p0t.y(), p1t.x());

        // if mem::size_of::<Float>() == mem::size_of::<f32>() && (e0 == 0.0 || e1 == 0.0 || e2 == 0.0)
        if e0 == 0.0 || e1 == 0.0 || e2 == 0.0 {
            let [p0t64, p1t64, p2t64] = [p0t.cast::<f64>(), p1t.cast::<f64>(), p2t.cast::<f64>()];

            e0 = (p2t64.y() * p1t64.x() - p2t64.x() * p1t64.y()) as Float;
            e1 = (p0t64.y() * p2t64.x() - p0t64.x() * p2t64.y()) as Float;
            e2 = (p1t64.y() * p0t64.x() - p1t64.x() * p0t64.y()) as Float;
        }

        if (e0 < 0. || e1 < 0. || e2 < 0.) && (e0 > 0. || e1 > 0. || e2 > 0.) {
            return None;
        }
        let det = e0 + e1 + e2;
        if det == 0. {
            return None;
        }

        // Compute scaled hit distance to triangle and test against ray
        p0t[2] *= sz;
        p1t[2] *= sz;
        p2t[2] *= sz;

        let t_scaled = e0 * p0t.z() + e1 * p1t.z() + e2 * p2t.z();
        if det < 0. && (t_scaled >= 0. || t_scaled < t_max * det)
            || (det > 0. && (t_scaled <= 0. || t_scaled > t_max * det))
        {
            return None;
        }

        // Compute barycentric coordinates and value for triangle intersection
        let inv_det = 1. / det;
        let b0 = e0 * inv_det;
        let b1 = e1 * inv_det;
        let b2 = e2 * inv_det;
        let t = t_scaled * inv_det;

        // Ensure that computed triangle is conservatively greater than zero
        let max_z_t = Vector3f::new(p0t.z(), p1t.z(), p2t.z())
            .abs()
            .max_component_value();
        let delta_z = gamma(3) * max_z_t;

        let max_x_t = Vector3f::new(p0t.x(), p1t.x(), p2t.x())
            .abs()
            .max_component_value();
        let max_y_t = Vector3f::new(p0t.y(), p1t.y(), p2t.y())
            .abs()
            .max_component_value();
        let delta_x = gamma(5) * (max_x_t + max_z_t);
        let delta_y = gamma(5) * (max_y_t + max_z_t);

        let delta_e = 2. * (gamma(2) * max_x_t * max_y_t + delta_y * max_x_t + delta_x * max_y_t);
        let max_e = Vector3f::new(e0, e1, e2).abs().max_component_value();
        let delta_t =
            3. * (gamma(3) * max_e * max_z_t + delta_e * max_z_t + delta_z * max_e) * inv_det.abs();

        if t <= delta_t {
            return None;
        }

        Some(TriangleIntersection::new(b0, b1, b2, t))
    }

    fn interaction_from_intersection(
        &self,
        ti: TriangleIntersection,
        time: Float,
        wo: Vector3f,
    ) -> SurfaceInteraction {
        let data = self.get_data();
        let [p0, p1, p2] = data.vertices;
        let [uv0, uv1, uv2] = data.uvs;
        // Compute triangle partial derivatives
        let (dpdu, dpdv, degenerate_uv, det) = self.compute_partials(data);
        // Interpolate (u, v) parametric coordinates and hit point
        let p_hit_vec =
            ti.b0 * Vector3f::from(p0) + ti.b1 * Vector3f::from(p1) + ti.b2 * Vector3f::from(p2);
        let p_hit = Point3f::from(p_hit_vec);
        let uv_hit_vec =
            ti.b0 * Vector2f::from(uv0) + ti.b1 * Vector2f::from(uv1) + ti.b2 * Vector2f::from(uv2);
        let uv_hit = Point2f::from(uv_hit_vec);

        // Return SurfaceInteraction for triangle hit>
        let flip_normal = data.reverse_orientation ^ data.transform_swaps_handedness;
        let p_abs_sum = (ti.b0 * Vector3f::from(p0)).abs()
            + (ti.b1 * Vector3f::from(p1)).abs()
            + (ti.b2 * Vector3f::from(p2)).abs();
        let p_error = gamma(7) * p_abs_sum;
        let mut isect = SurfaceInteraction::new(
            Point3fi::new_with_error(p_hit, p_error),
            uv_hit,
            wo,
            dpdu,
            dpdv,
            Normal3f::default(),
            Normal3f::default(),
            time,
            flip_normal,
        );

        isect.face_index = self
            .mesh()
            .face_indices
            .as_ref()
            .map_or(0, |fi| fi[self.tri_index]);
        isect.common.n = data.normal;
        isect.shading.n = isect.n();
        if flip_normal {
            isect.common.n = -isect.n();
            isect.shading.n = -isect.shading.n;
        }

        if data.normals.is_some() || self.mesh().s.is_some() {
            self.apply_shading_normals(&mut isect, ti, data, degenerate_uv, det);
        }

        isect
    }

    fn compute_partials(&self, data: TriangleData) -> (Vector3f, Vector3f, bool, Float) {
        let [p0, p1, p2] = data.vertices;
        let [uv0, uv1, uv2] = data.uvs;
        let duv02 = uv0 - uv2;
        let duv12 = uv1 - uv2;
        let dp02 = p0 - p2;
        let dp12 = p1 - p2;
        let det = difference_of_products(duv02[0], duv12[1], duv02[1], duv12[0]);
        let degenerate_uv = det.abs() < 1e-9;
        let (dpdu, dpdv) = if !degenerate_uv {
            let inv_det = 1. / det;
            (
                (dp02 * duv12[1] - dp12 * duv02[1]) * inv_det,
                (dp12 * duv02[0] - dp02 * duv12[0]) * inv_det,
            )
        } else {
            let dp20 = p2 - p0;
            let dp10 = p1 - p0;
            let mut ng = dp20.cross(dp10);
            if ng.norm_squared() == 0. {
                ng = (dp20.cast::<f64>().cross(dp10.cast::<f64>())).cast();
            }
            let n = ng.normalize();
            n.coordinate_system()
        };
        (dpdu, dpdv, degenerate_uv, det)
    }

    fn apply_shading_normals(
        &self,
        isect: &mut SurfaceInteraction,
        ti: TriangleIntersection,
        data: TriangleData,
        degenerate_uv: bool,
        det: Float,
    ) {
        let Some([n0, n1, n2]) = data.normals else {
            return;
        };
        let [uv0, uv1, uv2] = data.uvs;
        let duv02 = uv0 - uv2;
        let duv12 = uv1 - uv2;

        let ns = ti.b0 * n0 + ti.b1 * n1 + ti.b2 * n2;
        let ns = if ns.norm_squared() > 0. {
            ns.normalize()
        } else {
            isect.n()
        };

        let mut ss = self.mesh().s.as_ref().map_or(isect.dpdu, |s| {
            let indices = &self.mesh().vertex_indices[3 * self.tri_index..3 * self.tri_index + 3];
            let interp_s = ti.b0 * s[indices[0]] + ti.b1 * s[indices[1]] + ti.b2 * s[indices[2]];

            if interp_s.norm_squared() > 0. {
                interp_s
            } else {
                isect.dpdu
            }
        });

        let mut ts = Vector3f::from(ns).cross(ss);
        if ts.norm_squared() > 0. {
            ss = ts.cross(Vector3f::from(ns));
        } else {
            (ss, ts) = Vector3f::from(ns).coordinate_system();
        }
        let (dndu, dndv) = if degenerate_uv {
            let dn = (n2 - n0).cross(n1 - n0);
            if dn.norm_squared() == 0. {
                (Normal3f::zero(), Normal3f::zero())
            } else {
                dn.coordinate_system()
            }
        } else {
            let inv_det = 1. / det;
            let dn02 = n0 - n2;
            let dn12 = n1 - n2;
            (
                (dn02 * duv12[1] - dn12 * duv02[1]) * inv_det,
                (dn12 * duv02[0] - dn02 * duv12[0]) * inv_det,
            )
        };

        isect.shading.n = ns;
        isect.shading.dpdu = ss;
        isect.shading.dpdv = ts;
        isect.dndu = dndu;
        isect.dndv = dndv;
    }
}

impl ShapeTrait for TriangleShape {
    fn bounds(&self) -> Bounds3f {
        let [p0, p1, p2] = self.get_data().vertices;
        Bounds3f::from_points(p0, p1).union_point(p2)
    }

    fn normal_bounds(&self) -> DirectionCone {
        let data = self.get_data();
        let mut n = data.normal;
        if let Some([n0, n1, n2]) = data.normals {
            n = n.face_forward((n0 + n1 + n2).into());
        } else if data.reverse_orientation ^ data.transform_swaps_handedness {
            n = -n;
        }
        DirectionCone::new_from_vector(Vector3f::from(n))
    }

    fn area(&self) -> Float {
        self.get_data().area
    }

    fn pdf(&self, _interaction: &Interaction) -> Float {
        1. / self.area()
    }

    fn pdf_from_context(&self, ctx: &ShapeSampleContext, wi: Vector3f) -> Float {
        let solid_angle = self.solid_angle(ctx.p());
        if (Self::MIN_SPHERICAL_SAMPLE_AREA..Self::MAX_SPHERICAL_SAMPLE_AREA).contains(&solid_angle)
        {
            let ray = ctx.spawn_ray(wi);
            return self.intersect(&ray, None).map_or(0., |isect| {
                let absdot = Vector3f::from(isect.intr.n()).dot(-wi).abs();
                let d2 = ctx.p().distance_squared(isect.intr.p());
                let pdf = 1. / self.area() * (d2 / absdot);
                if pdf.is_infinite() { 0. } else { pdf }
            });
        }

        let mut pdf = 1. / solid_angle;
        if ctx.ns != Normal3f::zero() {
            let [p0, p1, p2] = self.get_data().vertices;
            let u = invert_spherical_triangle_sample(&[p0, p1, p2], ctx.p(), wi)
                .unwrap_or(Point2f::zero());

            let rp = ctx.p();
            let wi: [Vector3f; 3] = [
                (p0 - rp).normalize(),
                (p1 - rp).normalize(),
                (p2 - rp).normalize(),
            ];
            let w: [Float; 4] = [
                0.01_f32.max(ctx.ns.dot(wi[1].into()).abs()),
                0.01_f32.max(ctx.ns.dot(wi[1].into()).abs()),
                0.01_f32.max(ctx.ns.dot(wi[0].into()).abs()),
                0.01_f32.max(ctx.ns.dot(wi[2].into()).abs()),
            ];
            pdf *= bilinear_pdf(u, &w);
        }
        pdf
    }

    fn sample_from_context(&self, ctx: &ShapeSampleContext, u: Point2f) -> Option<ShapeSample> {
        let data = self.get_data();
        let [p0, p1, p2] = data.vertices;
        let solid_angle = self.solid_angle(ctx.p());
        if (Self::MIN_SPHERICAL_SAMPLE_AREA..Self::MAX_SPHERICAL_SAMPLE_AREA).contains(&solid_angle)
        {
            // Sample shape by area and compute incident direction wi
            return self.sample(u).and_then(|mut ss| {
                let mut intr_clone = (*ss.intr).clone();
                intr_clone.common.time = ctx.time;
                ss.intr = Arc::new(intr_clone);

                let wi = (ss.intr.p() - ctx.p()).normalize();
                if wi.norm_squared() == 0. {
                    return None;
                }
                let absdot = Vector3f::from(ss.intr.n()).abs_dot(-wi);
                let d2 = ctx.p().distance_squared(ss.intr.p());
                ss.pdf /= absdot / d2;
                if ss.pdf.is_infinite() { None } else { Some(ss) }
            });
        }

        // Sample spherical triangle from reference point
        let mut pdf = 1.;
        if ctx.ns != Normal3f::zero() {
            let rp = ctx.p();
            let wi: [Vector3f; 3] = [
                (p0 - rp).normalize(),
                (p1 - rp).normalize(),
                (p2 - rp).normalize(),
            ];
            let w: [Float; 4] = [
                0.01_f32.max(ctx.ns.dot(wi[1].into()).abs()),
                0.01_f32.max(ctx.ns.dot(wi[1].into()).abs()),
                0.01_f32.max(ctx.ns.dot(wi[0].into()).abs()),
                0.01_f32.max(ctx.ns.dot(wi[2].into()).abs()),
            ];

            let u = sample_bilinear(u, &w);
            pdf = bilinear_pdf(u, &w);
        }

        let (b, tri_pdf) = sample_spherical_triangle(&[p0, p1, p2], ctx.p(), u)?;
        if tri_pdf == 0. {
            return None;
        }
        pdf *= tri_pdf;
        let b2 = 1. - b[0] - b[1];

        let p_abs_sum = b[0] * Vector3f::from(p0)
            + b[1] * Vector3f::from(p1)
            + (1. - b[0] - b[1]) * Vector3f::from(p2);
        let p_error = gamma(6) * p_abs_sum;
        // Return ShapeSample for solid angle sampled point on triangle
        let p_vec =
            b[0] * Vector3f::from(p0) + b[1] * Vector3f::from(p1) + b[2] * Vector3f::from(p2);
        let p = Point3f::from(p_vec);
        let mut n = Normal3f::from((p1 - p0).cross(p2 - p0).normalize());
        if let Some([n0, n1, n2]) = data.normals {
            let ns = b[0] * n0 + b[1] * n1 + b2 * n2;
            n = n.face_forward(ns.into());
        } else if data.reverse_orientation ^ data.transform_swaps_handedness {
            n = -n;
        }

        let [uv0, uv1, uv2] = data.uvs;
        let uv_sample_vec =
            b[0] * Vector2f::from(uv0) + b[1] * Vector2f::from(uv1) + b[2] * Vector2f::from(uv2);
        let uv_sample = Point2f::from(uv_sample_vec);
        let pi = Point3fi::new_with_error(p, p_error);
        let mut si = SurfaceInteraction::new_simple(pi, n, uv_sample);
        si.common.time = ctx.time;
        Some(ShapeSample {
            intr: Arc::new(si),
            pdf,
        })
    }

    fn sample(&self, u: Point2f) -> Option<ShapeSample> {
        let data = self.get_data();
        let [p0, p1, p2] = data.vertices;
        let [uv0, uv1, uv2] = data.uvs;
        let b = sample_uniform_triangle(u);
        let p_vec =
            b[0] * Vector3f::from(p0) + b[1] * Vector3f::from(p1) + b[2] * Vector3f::from(p2);
        let b2 = 1. - b[0] - b[1];
        let p = Point3f::from(p_vec);
        let mut n = data.normal;
        if let Some([n0, n1, n2]) = data.normals {
            let interp_n = b[0] * n0 + b[1] * n1 + b2 * n2;
            n = n.face_forward(interp_n.into());
        } else if data.reverse_orientation ^ data.transform_swaps_handedness {
            n = -n;
        }

        let uv_sample_vec =
            b[0] * Vector2f::from(uv0) + b[1] * Vector2f::from(uv1) + b[2] * Vector2f::from(uv2);
        let uv_sample = Point2f::from(uv_sample_vec);
        let p_abs_sum = (b[0] * Vector3f::from(p0)).abs()
            + (b[1] * Vector3f::from(p1)).abs()
            + ((1. - b[0] - b[1]) * Vector3f::from(p2)).abs();
        let p_error = gamma(6) * p_abs_sum;
        let pi = Point3fi::new_with_error(p, p_error);
        Some(ShapeSample {
            intr: Arc::new(SurfaceInteraction::new_simple(pi, n, uv_sample)),
            pdf: 1. / self.area(),
        })
    }

    fn intersect(&self, ray: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        self.intersect_triangle(ray, t_max.unwrap_or(Float::INFINITY))
            .map(|ti| {
                let intr = self.interaction_from_intersection(ti, ray.time, -ray.d);
                ShapeIntersection { intr, t_hit: ti.t }
            })
    }

    fn intersect_p(&self, ray: &Ray, t_max: Option<Float>) -> bool {
        self.intersect_triangle(ray, t_max.unwrap_or(Float::INFINITY))
            .is_some()
    }
}