#include <core/fwd.h>
#include <render/shape.h>
#include <render/intersection.h>
#include <core/math_func.h>
#include <core/ray.h>
#include <iomanip>
#include <map>
#include <tiny_obj_loader.h>
#include <core/logger.h>
#include <core/timer.h>
#include <unordered_set>
#include <random>
#include <algorithm>
#include <iostream>
#include <fstream>

namespace
{
    [[maybe_unused]] void print(std::unordered_set<int> const &s)
    {
        std::copy(s.begin(),
                  s.end(),
                  std::ostream_iterator<int>(std::cout, " "));
    }
    

    std::vector<Edge> constructEdges(const std::vector<Vector> &vertices,
                                     const std::vector<Vector3i> &indices,
                                     const std::vector<Vector> &faceNormals)
    {
        /* pair of vertex indices, vector (face i, third vertex i) */
        std::map<std::pair<int, int>, std::vector<Vector2i>> edge_map;
        for (int itri = 0; itri < indices.size(); itri++)
        {
            Vector3i ind = indices[itri];
            for (int iedge = 0; iedge < 3; iedge++)
            {
                int k1 = iedge,
                    k2 = (iedge + 1) % 3;
                std::pair<int, int> key = (ind[k1] < ind[k2])
                                              ? std::make_pair(ind[k1], ind[k2])
                                              : std::make_pair(ind[k2], ind[k1]);
                if (edge_map.find(key) == edge_map.end())
                    edge_map[key] = std::vector<Vector2i>();
                edge_map[key].push_back(Vector2i(itri,
                                                 ind[(iedge + 2) % 3]));
            }
        }

        std::vector<Edge> edges;
        for (const auto &it : edge_map)
        {
            Float length = (vertices[it.first.first] -
                            vertices[it.first.second])
                               .norm();

            // check if good mesh
            if (it.second.size() > 2)
            {
                std::cerr << "Every edge can be shared by at most 2 faces!"
                          << std::endl;
                assert(false);
            }
            else if (it.second.size() == 2) // an edge is shared by two faces
            {
                // face id
                const int ind0 = it.second[0][0],
                          ind1 = it.second[1][0];
                if (ind0 == ind1)
                {
                    std::cerr << "Duplicated faces!" << std::endl;
                    assert(false);
                }

                const Vector &n0 = faceNormals[ind0],
                             &n1 = faceNormals[ind1];
                Float val = n0.dot(n1);
                if (val < -1.0f + EdgeEpsilon)
                {
                    std::cerr << "Inconsistent normal orientation! n0.dot(n1) = "
                              << val << std::endl;
                    // assert(false);
                }
                else if (val < 1.0f - EdgeEpsilon) // not coplanar
                {
                    //
                    [[maybe_unused]] Float tmp0 = n0.dot(vertices[it.second[1][1]] -
                                                         vertices[it.first.first]),
                                           tmp1 = n1.dot(vertices[it.second[0][1]] -
                                                         vertices[it.first.first]);
                    assert(math::signum(tmp0) * math::signum(tmp1) > 0.5f);
                    edges.push_back(Edge(it.first.first, it.first.second,
                                         ind0, ind1, length, it.second[0][1],
                                         tmp0 > Epsilon ? -1 : 1));
                }
            }
            else
            {
                assert(it.second.size() == 1);
                edges.push_back(Edge(it.first.first, it.first.second,
                                     it.second[0][0], -1, length, it.second[0][1], 0));
            }
        }
        return edges;
    }
    
    void constructBoundaryTest(const std::vector<Vector> &vertices,
                                const std::vector<Vector3i> &indices,
                                const std::vector<Vector> &faceNormals,
                                std::vector<std::vector<int>> &boundaryEdge,
                                std::vector<Float> &vertDist2Boundary
                                )
    {
        int vertex_size = vertices.size();
        int indice_size = indices.size();
        assert(boundaryEdge.size() == indice_size);
        assert(vertDist2Boundary.size() == vertex_size);
        /* pair of vertex indices, vector (face i, third vertex i) */
        std::map<std::pair<int, int>, std::vector<Vector2i>> edge_map;
        for (int itri = 0; itri < indices.size(); itri++)
        {
            Vector3i ind = indices[itri];
            for (int iedge = 0; iedge < 3; iedge++)
            {
                int k1 = iedge,
                    k2 = (iedge + 1) % 3;
                std::pair<int, int> key = (ind[k1] < ind[k2])
                                              ? std::make_pair(ind[k1], ind[k2])
                                              : std::make_pair(ind[k2], ind[k1]);
                if (edge_map.find(key) == edge_map.end())
                    edge_map[key] = std::vector<Vector2i>();
                edge_map[key].push_back(Vector2i(itri,
                                                 ind[(iedge + 2) % 3]));
            }
        }

        std::vector<Edge> edges;
        for (const auto &it : edge_map)
        {
            Float length = (vertices[it.first.first] -
                            vertices[it.first.second])
                               .norm();

            // check if good mesh
            if (it.second.size() > 2)
            {
                std::cerr << "Every edge can be shared by at most 2 faces!"
                          << std::endl;
                assert(false);
            }
            else if (it.second.size() == 2) // an edge is shared by two faces
            {
                // face id
                const int ind0 = it.second[0][0],
                          ind1 = it.second[1][0];
                if (ind0 == ind1)
                {
                    std::cerr << "Duplicated faces!" << std::endl;
                    assert(false);
                }

                const Vector &n0 = faceNormals[ind0],
                             &n1 = faceNormals[ind1];
                Float val = n0.dot(n1);
                if (val < -1.0f + EdgeEpsilon)
                {
                    std::cerr << "Inconsistent normal orientation! n0.dot(n1) = "
                              << val << std::endl;
                    // assert(false);
                }
                else if (val < 1.0f - EdgeEpsilon) // not coplanar
                {
                    //
                    [[maybe_unused]] Float tmp0 = n0.dot(vertices[it.second[1][1]] -
                                                         vertices[it.first.first]),
                                           tmp1 = n1.dot(vertices[it.second[0][1]] -
                                                         vertices[it.first.first]);
                    assert(math::signum(tmp0) * math::signum(tmp1) > 0.5f);
                    edges.push_back(Edge(it.first.first, it.first.second,
                                         ind0, ind1, length, it.second[0][1],
                                         tmp0 > Epsilon ? -1 : 1));
                }
            }
            else
            {
                assert(it.second.size() == 1);
                boundaryEdge[it.second[0][0]].push_back(it.second[0][1]);
                Vector3i ind = indices[it.second[0][0]];
                

                vertDist2Boundary[it.first.first] = 0.0;
                vertDist2Boundary[it.first.second] = 0.0;
                Vector p0 = vertices[it.first.first];
                Vector p1 = vertices[it.first.second];
                Vector p2 = vertices[it.second[0][1]];
                Vector v = p2 - p0;
                Vector u = p1 - p0;
                Float c1 = v.dot(u);
                Float c2 = u.dot(u);
                Float b = c1 / c2;
                Vector Pb = p0 + b * u;
                Float len = (p2 - Pb).norm();
                vertDist2Boundary[it.second[0][1]] = len < vertDist2Boundary[it.second[0][1]] ? len : vertDist2Boundary[it.second[0][1]];
                edges.push_back(Edge(it.first.first, it.first.second,
                                     it.second[0][0], -1, length, it.second[0][1], 0));
            }
        }
    }


    std::vector<Edge> constructSortEdges(const std::vector<Vector> &vertices,
                                         const std::vector<Vector3i> &indices,
                                         const std::vector<Vector> &faceNormals, const Sort_config &sort_config)
    {
        std::map<int, std::unordered_set<int>> vertices_neighbor;
        [[maybe_unused]] int vertex_size = vertices.size();
        [[maybe_unused]] int indice_size = indices.size();

        /* pair of vertex indices, vector (face i, third vertex i) */
        std::map<std::pair<int, int>, std::vector<Vector2i>> edge_map;
        for (int itri = 0; itri < indices.size(); itri++)
        {
            Vector3i ind = indices[itri];
            for (int iedge = 0; iedge < 3; iedge++)
            {
                int k1 = iedge,
                    k2 = (iedge + 1) % 3;
                std::pair<int, int> key = (ind[k1] < ind[k2])
                                              ? std::make_pair(ind[k1], ind[k2])
                                              : std::make_pair(ind[k2], ind[k1]);
                if (edge_map.find(key) == edge_map.end())
                    edge_map[key] = std::vector<Vector2i>();
                edge_map[key].push_back(Vector2i(itri,
                                                 ind[(iedge + 2) % 3]));
            }
        }
        std::vector<Edge> edges;
        for (const auto &it : edge_map)
        {
            Float length = (vertices[it.first.first] -
                            vertices[it.first.second])
                               .norm();
            // check if good mesh
            if (it.second.size() > 2)
            {
                std::cerr << "Every edge can be shared by at most 2 faces!"
                          << std::endl;
                assert(false);
            }
            else if (it.second.size() == 2) // an edge is shared by two faces
            {
                // face id
                const int ind0 = it.second[0][0],
                          ind1 = it.second[1][0];
                if (ind0 == ind1)
                {
                    std::cerr << "Duplicated faces!" << std::endl;
                    assert(false);
                }

                const Vector &n0 = faceNormals[ind0],
                             &n1 = faceNormals[ind1];
                Float val = n0.dot(n1);
                if (val < -1.0f + EdgeEpsilon)
                {
                    std::cerr << "Inconsistent normal orientation! n0.dot(n1) = "
                              << val << std::endl;
                }
                else if (val < 1.0f - EdgeEpsilon) // not coplanar
                {
                    [[maybe_unused]] Float tmp0 = n0.dot(vertices[it.second[1][1]] -
                                                         vertices[it.first.first]),
                                           tmp1 = n1.dot(vertices[it.second[0][1]] -
                                                         vertices[it.first.first]);
                    assert(math::signum(tmp0) * math::signum(tmp1) > 0.5f);
                    edges.push_back(Edge(it.first.first, it.first.second,
                                         ind0, ind1, length, it.second[0][1],
                                         tmp0 > Epsilon ? -1 : 1));
                }
            }
            else
            {
                assert(it.second.size() == 1);
                edges.push_back(Edge(it.first.first, it.first.second,
                                     it.second[0][0], -1, length, it.second[0][1], 0));
            }
        }

        if (edges[0].length < 0.00001)
        {
            return edges;
        }
        auto rd = std::random_device{};
        auto rng = std::default_random_engine{rd()};
        rng.seed(10);
        std::shuffle(std::begin(edges), std::end(edges), rng);

        if (sort_config.max_length == -1)
        {
            // std::cout << "using unsorted edge" << std::endl;
            return edges;
        }

        for (int i = 0; i < edges.size(); ++i)
        {

            if (vertices_neighbor.find(edges[i].v0) == vertices_neighbor.end())
                vertices_neighbor[edges[i].v0] = std::unordered_set<int>();
            vertices_neighbor[edges[i].v0].insert(edges[i].v1);

            if (vertices_neighbor.find(edges[i].v1) == vertices_neighbor.end())
                vertices_neighbor[edges[i].v1] = std::unordered_set<int>();
            vertices_neighbor[edges[i].v1].insert(edges[i].v0);
        }

        std::map<std::pair<int, int>, bool> edge_walk_map;
        for (int i = 0; i < edges.size(); ++i)
        {
            std::pair<int, int> walk_key = std::make_pair(edges[i].v0, edges[i].v1);
            edge_walk_map[walk_key] = false;
        }

        // std::vector<int> jump;

        std::vector<Vector2i> draw;
        for (int i = 0; i < vertices_neighbor.size(); ++i)
        {
            // Greedy Search Here
            for (int neig : vertices_neighbor[i])
            {
                int prev = i;
                int curr = neig;
                Vector P_start = vertices[curr] - vertices[prev];

                std::pair<int, int> walk_key = (prev < curr)
                                                   ? std::make_pair(prev, curr)
                                                   : std::make_pair(curr, prev);
                if (edge_walk_map[walk_key] == false)
                {
                    // Greedy Search
                    for (int depth = 0; depth < sort_config.max_length; ++depth)
                    {
                        std::pair<int, int> walk_key = (prev < curr)
                                                           ? std::make_pair(prev, curr)
                                                           : std::make_pair(curr, prev);
                        if (edge_walk_map[walk_key] == true)
                        {
                            break;
                        }
                        edge_walk_map[walk_key] = true;
                        draw.push_back(Vector2i(prev, curr));
                        int best_id = -1;
                        Float best_cos = 1.1;

                        for (int curr_id : vertices_neighbor[curr])
                        {
                            std::pair<int, int> walk_best_key = (curr < curr_id)
                                                                    ? std::make_pair(curr, curr_id)
                                                                    : std::make_pair(curr_id, curr);
                            if (edge_walk_map[walk_best_key] == false)
                            {
                                Vector P0 = vertices[curr] - vertices[prev];
                                Vector P1 = vertices[curr] - vertices[curr_id];

                                Float cos_val = P0.dot(P1) / (P0.norm() * P1.norm());
                                [[maybe_unused]] Float start_cos_val = P_start.dot(P1) / (P_start.norm() * P1.norm());
                                if (cos_val < best_cos)
                                {
                                    best_id = curr_id;
                                    best_cos = cos_val;
                                }
                            }
                        }
                        Vector P_end = vertices[curr] - vertices[best_id];
                        Float global_cos = P_start.dot(P_end) / (P_start.norm() * P_end.norm());
                        // std::cout << prev << " -> " << curr << " :-> " << best_id << std::endl;

                        if (best_cos > sort_config.cos_val_local || global_cos > sort_config.cos_val_global || best_id == -1)
                        {
                            // std::cout << "break" << std::endl;
                            break;
                        }

                        // std::cout << prev << " -> " << curr << " :-> " << best_id << std::endl;

                        prev = curr;
                        curr = best_id;
                    }
                    // jump.push_back(draw.size()-1);
                }
            }
        }
        // std::cout << jump.size() << std::endl;
        // std::cout << draw.size() << std::endl;

        // for (int j : jump) {
        //     std::cout << j << " ";
        // }
        // std::cout << std::endl;

        std::map<std::pair<int, int>, int> sort_edge_map;
        for (int i = 0; i < edges.size(); ++i)
        {
            const auto &e = edges[i];
            std::pair<int, int> e_pair(e.v0, e.v1);
            sort_edge_map[e_pair] = i;
        }

        std::vector<Edge> sort_edges;

        for (Vector2i draw_edge : draw)
        {
            std::pair<int, int> key = (draw_edge[0] < draw_edge[1])
                                          ? std::make_pair(draw_edge[0], draw_edge[1])
                                          : std::make_pair(draw_edge[1], draw_edge[0]);

            int id_edge = sort_edge_map[key];
            Edge data = edges[id_edge];
            data.v0 = draw_edge[0];
            data.v1 = draw_edge[1];
            sort_edges.push_back(data);
        }

        assert(sort_edges.size() == edges.size());
        std::cout << "using sorted edge" << std::endl;
        return sort_edges;
    }
}

Shape::Shape(const std::vector<Vector> &vertices, const std::vector<Vector3i> &indices,
             const std::vector<Vector2> &uvs, const std::vector<Vector> &normals,
             int num_vertices, int num_triangles, int light_id, int bsdf_id,
             int med_int_id, int med_ext_id)
    : vertices(vertices), indices(indices), uvs(uvs),
      num_vertices(num_vertices), num_triangles(num_triangles),
      light_id(light_id), bsdf_id(bsdf_id), med_int_id(med_int_id),
      med_ext_id(med_ext_id)
{
    assert(vertices.size() == num_vertices);
    assert(indices.size() == num_triangles);

    if (m_use_face_normals)
        assert(normals.size() == 0);
    else
    {
        if (normals.size() > 0)
        {
            assert(light_id < 0);
            assert(normals.size() == num_vertices);
            this->normals.resize(num_vertices);
            for (int i = 0; i < num_vertices; i++)
                this->normals[i] = normals[i].normalized();
        }
    }

    configure();
}

Shape::Shape(const std::vector<Vector> &vertices, const std::vector<Vector3i> &indices,
             const std::vector<Vector2> &uvs, const std::vector<Vector> &normals,
             int num_vertices, int num_triangles, int light_id, int bsdf_id,
             int med_int_id, int med_ext_id, const Sort_config &config)
    : vertices(vertices), indices(indices), uvs(uvs),
      num_vertices(num_vertices), num_triangles(num_triangles),
      light_id(light_id), bsdf_id(bsdf_id), med_int_id(med_int_id),
      med_ext_id(med_ext_id)
{
    assert(vertices.size() == num_vertices);
    assert(indices.size() == num_triangles);

#ifdef SHAPE_COMPUTE_VTX_NORMAL
    assert(normals.size() == 0);
#else
    if (normals.size() > 0)
    {
        assert(light_id < 0);
        assert(normals.size() == num_vertices);
        this->normals.resize(num_vertices);
        for (int i = 0; i < num_vertices; i++)
            this->normals[i] = normals[i].normalized();
    }
#endif
    sort_config = config;
    configure();
}

// load vertices, indices, normals, uvs from obj file and then configure
void Shape::load(const std::string &filename)
{
    vertices.clear();
    indices.clear();
    normals.clear();
    uvs.clear();

    tinyobj::attrib_t attrib;
    std::vector<tinyobj::shape_t> shapes;
    std::vector<tinyobj::material_t> materials;
    std::string warn, err;
    PSDR_ASSERT_MSG(tinyobj::LoadObj(&attrib, &shapes, &materials,
                                     &warn, &err, filename.c_str()),
                    "Failed to load obj file: " + filename + "\n" + err);
    assert(attrib.vertices.size() % 3 == 0);

    // Load vertices
    vertices.resize(attrib.vertices.size() / 3);
    for (int i = 0; i < vertices.size(); i++)
        vertices[i] = Vector(attrib.vertices[i * 3],
                             attrib.vertices[i * 3 + 1],
                             attrib.vertices[i * 3 + 2]);
    // Load uvs
    if (attrib.texcoords.size() > 0)
    {
        uvs.resize(attrib.texcoords.size() / 2);
        for (int i = 0; i < uvs.size(); i++)
            uvs[i] = Vector2(attrib.texcoords[i * 2],
                             attrib.texcoords[i * 2 + 1]);
    }

    // Load faces
    indices.clear();
    uv_indices.clear();
    for (const auto &shape : shapes)
    {
        auto &idx = shape.mesh.indices;
        for (int f = 0; f < shape.mesh.num_face_vertices.size(); f++)
        {
            int fv = shape.mesh.num_face_vertices[f];
            PSDR_ASSERT_MSG(fv == 3, "Only triangular faces are supported!");
            // Load vertex indices
            indices.push_back(Vector3i(idx[f * 3].vertex_index,
                                       idx[f * 3 + 1].vertex_index,
                                       idx[f * 3 + 2].vertex_index));
            // Load uv indices
            if (uvs.size() > 0)
            {
                PSDR_ASSERT_MSG(idx[f * 3].texcoord_index >= 0,
                                "Texture coordinates are not provided!");
                uv_indices.push_back(Vector3i(idx[f * 3].texcoord_index,
                                              idx[f * 3 + 1].texcoord_index,
                                              idx[f * 3 + 2].texcoord_index));
            }
        }
    }
    configure();
}

__attribute__((optnone)) double Shape::B(const Ray& ray, const Intersection& its) const {
    Vector3i ind = getIndices(its.triangle_id);
    // if (this->boundaryEdge[its.indices[1]].size() == 0) {
    //     // Float propagation = 0.01;
    //     // Float Bv0 = B_v(ray.dir, normals[int(0)].val, propagation);
    //     // Float Bv1 = B_v(ray.dir, normals[int(1)].val, propagation);
    //     // Float Bv2 = B_v(ray.dir, normals[int(2)].val, propagation);
    //     // Vector2 uv = its.barycentric;
    //     // return (1.0f - uv(0) - uv(1))*Bv0 + uv(0)*Bv1 + uv(1)*Bv2;
    //     return abs(ray.dir.dot(its.shFrame.n));
    // }
    // // if (this->boundaryEdge[its.indices[1]].size() == 3)
    // //     return min3(its.barycentric[0], its.barycentric[1], 1 - its.barycentric[0] - its.barycentric[1]);
    if (boundaryEdge[its.triangle_id].size() >= 2) {
        Float min_dist = 1e10;
        for (int i = 0; i < boundaryEdge[its.triangle_id].size(); i++) {
            int nonB = 0;
            // std::cout << this->boundaryEdge[its.indices[0]][i] << std::endl;
            for (int j = 0; j < 3; j++) {
                if (ind[j] == boundaryEdge[its.triangle_id][i]) 
                    nonB = j;
            }
            Vector edgeDir = (vertices[ind[(nonB + 1) % 3]] - vertices[ind[(nonB + 2) % 3]]).normalized();
            Vector ptDir = its.p - vertices[ind[(nonB + 2) % 3]];
            Float dist = (ptDir - edgeDir.dot(ptDir) * edgeDir).norm();
            if (dist < min_dist) min_dist = dist;
        }
        return min_dist;
    }
    Float B_attached = abs(ray.dir.dot(faceNormals[its.triangle_id]));
    // B_attached = B_attached / (1.0 - (1.0 - 0.1) * (1.0 - B_attached));
    const Float &b0 = vertDist2Boundary[ind[0]], &b1 = vertDist2Boundary[ind[1]], &b2 = vertDist2Boundary[ind[2]];
    Float B_topology = (1.0f - its.barycentric(0) - its.barycentric(1))*b0 + its.barycentric(0)*b1 + its.barycentric(1)*b2;
    // B_topology = B_topology / (1.0 - (1.0 - 0.1) * (1.0 - B_topology));
    // B_topology = 1000;
    Float B = B_attached < B_topology ? B_attached : B_topology;
    return B;
}


void Shape::configure()
{
    num_triangles = indices.size();
    num_vertices = vertices.size();
    boundaryEdge.resize(num_triangles);
    vertDist2Boundary.resize(num_vertices);
    std::fill(vertDist2Boundary.begin(), vertDist2Boundary.end(), 1.0);

    // ================= Normal preprocessing =================
    // Compute adjacent faces
    computeAdjacentFaces();

    // Compute face distribution
    computeFaceDistribution();

    // Compute face normals
    faceNormals.clear();
    faceNormals.resize(num_triangles);
    computeFaceNormals();

    // Compute vertex normals
    if (!m_use_face_normals)
    {
        normals.resize(num_vertices);
        computeVertexNormals();
    }
    // ========================================================
    // ================= Edge preprocessing ===================
    // Compute edges
    if (enable_draw)
    {
        edges = constructSortEdges(vertices, indices, faceNormals, sort_config);
    }
    else
    {
        edges = constructEdges(vertices, indices, faceNormals);
    }
    constructBoundaryTest(vertices, indices, faceNormals, boundaryEdge, vertDist2Boundary);

    // Compute edge distribution
    edge_distrb.clear();
    for (const Edge &edge : edges)
        edge_distrb.append(edge.length);
    edge_distrb.normalize();
    // ========================================================
}

void Shape::computeAdjacentFaces()
{
    adjacentFaces.clear();
    adjacentFaces.resize(vertices.size());
    for (int i = 0; i < indices.size(); i++)
    {
        assert(i < indices.size());
        auto &face = indices.at(i);
        for (int j = 0; j < 3; j++)
        {
            adjacentFaces.at(face[j]).push_back(i);
        }
    }
}

void Shape::computeFaceDistribution()
{
    face_distrb.clear();
    face_distrb.reserve(num_triangles);
    for (int i = 0; i < num_triangles; ++i)
    {
        const auto &ind = getIndices(i);
        const auto &v0 = getVertex(ind(0));
        const auto &v1 = getVertex(ind(1));
        const auto &v2 = getVertex(ind(2));
        auto cur = (v1 - v0).cross(v2 - v0);
        Float area = cur.norm();
        face_distrb.append(0.5f * area);
    }
    face_distrb.normalize();
}
INACTIVE_FN(Shape_computeFaceDistribution, &Shape::computeFaceDistribution);

void Shape::computeFaceNormals()
{
    assert(faceNormals.size() == num_triangles);
    for (int i = 0; i < num_triangles; ++i)
    {
        const auto &ind = getIndices(i);
        const auto &v0 = getVertex(ind(0));
        const auto &v1 = getVertex(ind(1));
        const auto &v2 = getVertex(ind(2));
        [[maybe_unused]] auto n = (v1 - v0).cross(v2 - v0);
        [[maybe_unused]] Float area = n.norm();
        // ! assert(area > Epsilon);
        faceNormals[i] = n.normalized();
    }
}

Float unitAngle(const Vector &u, const Vector &v)
{
    if (u.dot(v) < 0.0)
        return M_PI - asin((0.5 * (v + u).norm()));
    else
        return 2.0 * asin((0.5 * (v - u).norm()));
}

void Shape::computeVertexNormals()
{
    assert(normals.size() == num_vertices);
    for (int i = 0; i < num_vertices; ++i)
    {
        normals[i] = Vector::Zero();
        for (int j : adjacentFaces[i])
            normals[i] += getFaceNormal(j);
        normals[i].normalize();
    }
}

Float Shape::getArea(int index) const
{
    auto &ind = getIndices(index);
    auto &v0 = getVertex(ind(0));
    auto &v1 = getVertex(ind(1));
    auto &v2 = getVertex(ind(2));
    return 0.5f * (v1 - v0).cross(v2 - v0).norm();
}

// void Shape::samplePosition(int index, const Vector2 &rnd2, Vector &pos,
// Vector &norm) const {
//     const Vector3i& ind = getIndices(index);
//     const Vector& v0 = getVertex(ind(0));
//     const Vector& v1 = getVertex(ind(1));
//     const Vector& v2 = getVertex(ind(2));
//     Float a = std::sqrt(rnd2[0]);
//     pos = v0 + (v1 - v0)*(1.0f - a) + (v2 - v0)*(a*rnd2[1]);
//     norm = faceNormals[index].val;
// }

int Shape::samplePosition(const Vector2 &_rnd2,
                          PositionSamplingRecord &pRec) const
{
    Vector2 rnd2(_rnd2);
    int tri_idx = static_cast<int>(face_distrb.sampleReuse(rnd2[0]));
    pRec.tri_id = tri_idx;
    Float a = std::sqrt(rnd2[0]);
    Vector2 uv = Vector2{1.0f - a, a * rnd2[1]};
    Float u = detach(uv.x());
    Float v = detach(uv.y());
    pRec.barycentric = Vector2{u, v};
    const Vector3i ind = getIndices(tri_idx);
    const Vector &v0 = getVertex(ind(0));
    const Vector &v1 = getVertex(ind(1));
    const Vector &v2 = getVertex(ind(2));
    pRec.p = v0 + (v1 - v0) * u + (v2 - v0) * v;
    pRec.n = getFaceNormal(tri_idx);
    pRec.uv = Vector2(u, v);
    Float area = getArea();
    pRec.pdf = detach(1. / area);
    pRec.measure = EMArea;
    area = getArea(tri_idx);
    pRec.J = area / detach(area);
    return tri_idx;
}

void Shape::rayIntersect(int tri_index, const Ray &ray,
                         Intersection &its,
                         IntersectionMode mode) const
{
    its.type = EVSurface;
    if (its.isEmitter())
        its.type = EVEmitter;
    const Vector3i &ind = getIndices(tri_index);
    const Vector &v0 = getVertex(ind(0)),
                 &v1 = getVertex(ind(1)),
                 &v2 = getVertex(ind(2));

    const Array uvt = rayIntersectTriangle(v0, v1, v2, ray);
    const Float u = uvt(0), v = uvt(1), t = uvt(2);

    Vector2 barycentric = Vector2(u, v);
    if (mode == EMaterial)
        barycentric = Vector2(detach(u), detach(v));

    Vector geom_normal = getFaceNormal(tri_index);
    Vector shading_normal = getShadingNormal(tri_index, barycentric);
    // Flip geometric normal to the same side of shading normal
    if (geom_normal.dot(shading_normal) < 0.f)
    {
        geom_normal = -geom_normal;
    }
    its.geoFrame = Frame(geom_normal);
    its.shFrame = Frame(shading_normal);

    if (mode == ESpatial)
        its.p = ray.org + ray.dir * t;
    else if (mode == EMaterial)
        its.p = (1 - barycentric[0] - barycentric[1]) * v0 +
                barycentric[0] * v1 + barycentric[1] * v2;
    else if (mode == EReference)
        its.p = (1 - u - v) * detach(v0) +
                u * detach(v1) +
                v * detach(v2);

    if (mode == ESpatial)
    {
        its.t = t;
        its.uv = getUV(tri_index, barycentric);
        its.wi = its.toLocal(-ray.dir);
        its.barycentric = barycentric;
    }
    else
    {
        Vector dir = its.p - ray.org;
        Float dist = dir.norm();
        dir /= dist;
        its.t = dist;
        its.barycentric = barycentric;
        its.uv = getUV(tri_index, barycentric);
        its.wi = its.toLocal(-dir);
    }
    its.J = getArea(tri_index);
    its.J /= detach(its.J);
}


std::tuple<Vector, Vector, Float> Shape::getPoint(int tri_index, const Vector2 &_barycentric) const
{
    Vector2 barycentric = detach(_barycentric);
    assert(tri_index >= 0 && tri_index < num_triangles);
    const Vector3i &ind = indices[tri_index];
    const Vector &v0 = getVertex(ind[0]),
                 &v1 = getVertex(ind[1]),
                 &v2 = getVertex(ind[2]);
    Vector x = (1. - barycentric.x() - barycentric.y()) * v0 +
               barycentric.x() * v1 +
               barycentric.y() * v2;
    Vector n = getFaceNormal(tri_index);
    Float J = getArea(tri_index);
    return std::make_tuple(x, n, J / detach(J));
}

void Shape::getPoint(const Intersection &its, Intersection &itsAD) const
{
    Vector2 barycentric = detach(its.barycentric);
    int tri_index = its.indices[1];
    assert(tri_index >= 0 && tri_index < num_triangles);
    const Vector3i &ind = indices[tri_index];
    const Vector &v0 = getVertex(ind[0]),
                 &v1 = getVertex(ind[1]),
                 &v2 = getVertex(ind[2]);
    itsAD.p = (1. - barycentric.x() - barycentric.y()) * v0 +
              barycentric.x() * v1 +
              barycentric.y() * v2;
    Vector geom_normal = getFaceNormal(tri_index);
    Vector shading_normal = getShadingNormal(tri_index, barycentric);
    itsAD.geoFrame = Frame(geom_normal);
    itsAD.shFrame = Frame(shading_normal);
    itsAD.barycentric = barycentric;
    itsAD.J = getArea(tri_index);
    itsAD.J /= detach(itsAD.J);
}

void Shape::getPoint(int tri_index, const Vector2 &_barycentric, Intersection &its, Float &J) const
{
    Vector2 barycentric = detach(_barycentric);
    Vector geom_normal;
    std::tie(its.p, geom_normal, J) = getPoint(tri_index, barycentric);
    Vector shading_normal = getShadingNormal(tri_index, barycentric);
    if (geom_normal.dot(shading_normal) < 0.)
        geom_normal = -geom_normal;

    its.geoFrame = Frame(geom_normal);
    its.shFrame = Frame(shading_normal);

    its.uv = getUV(tri_index, barycentric);
}

bool Shape::isSihoulette(const Edge &edge, const Vector &p) const
{
    if (edge.f0 == -1 || edge.f1 == -1)
        return true;

    const Vector &v0 = vertices[edge.v0],
                 &v1 = vertices[edge.v1];

    bool frontfacing0 = false;
    const Vector3i &ind0 = indices[edge.f0];
    for (int i = 0; i < 3; i++)
    {
        if (ind0[i] != edge.v0 && ind0[i] != edge.v1)
        {
            const Vector &v = vertices[ind0[i]];
            Vector n0 = (v0 - v).cross(v1 - v).normalized();
            frontfacing0 = n0.dot(p - v) > 0.0f;
            break;
        }
    }

    bool frontfacing1 = false;
    const Vector3i &ind1 = indices[edge.f1];
    for (int i = 0; i < 3; i++)
    {
        if (ind1[i] != edge.v0 && ind1[i] != edge.v1)
        {
            const Vector &v = vertices[ind1[i]];
            Vector n1 = (v1 - v).cross(v0 - v).normalized();
            frontfacing1 = n1.dot(p - v) > 0.0f;
            break;
        }
    }
    return (frontfacing0 && !frontfacing1) ||
           (!frontfacing0 && frontfacing1);
}

#ifdef FORWARD
const Vector3 Shape::getVertex(int index) const
{
    // global transformation
    if (requires_grad && vertex_idx == -1)
        return transform(index, param);
    // vertex transformation
    else if (requires_grad && vertex_idx >= 0 && index == vertex_idx)
        return transform(index, param);
    else
        return vertices[index];
}

Vector Shape::transform(int index, Float param) const
{
    Vector v = vertices[index];
    if (is_tranlate)
        return param * translation + detach(v); //! detach is neccessary here
    // else if (is_rotate) // segfault
    //     return Eigen::AngleAxis<Float>(param, rotation.normalized())
    //                .toRotationMatrix() *
    //            detach(v);
    else if (is_scale)
        return param * scale.cwiseProduct(detach(v));
    else if (is_velocity)
    {
        Vector velocity = velocities[index];
        return param * detach(velocity) + detach(v); // 1st order Taylor expansion
    }
    else
        return v;
}
#endif

Vector Shape::getShadingNormal(int tri_index, const Vector2 &barycentric) const
{
    Vector geom_normal = getFaceNormal(tri_index);
    const Vector3i &ind = indices[tri_index];
    Vector shading_normal;
    if (!m_use_face_normals)
    {
        /* interpolation */
        const Float u = barycentric[0],
                    v = barycentric[1];
        const Vector &n0 = getVertexNormal(ind[0]),
                     &n1 = getVertexNormal(ind[1]),
                     &n2 = getVertexNormal(ind[2]);
        shading_normal = ((1.0f - u - v) * n0 + u * n1 + v * n2)
                             .normalized();
    }
    else
        shading_normal = geom_normal;
    return shading_normal;
}

/* ==================== Normal precomputation ==================== */
#ifdef NORMAL_PREPROCESS
namespace
{
    void precompute_normal(Shape &shape)
    {
        /* compute face normal */
        shape.computeFaceNormals();
        /* compute vertex normal */
        if (!shape.m_use_face_normals)
            shape.computeVertexNormals();
    }
}

void d_precompute_normal(const Shape &shape, Shape &d_shape)
{
    __enzyme_autodiff((void *)precompute_normal,
                      enzyme_dup, &shape, &d_shape);
}
#endif

bool Shape::save(const std::string &name) const
{
    std::ofstream ofs(name);
    if (!ofs.is_open())
        throw std::runtime_error("failed to open file");
    for (int i = 0; i < num_vertices; i++)
    {
        ofs << "v " << vertices[i].x() << " " << vertices[i].y() << " " << vertices[i].z() << std::endl;
    }
    if (uvs.size() > 0)
    {
        for (int i = 0; i < num_vertices; i++)
        {
            ofs << "vt " << uvs[i].x() << " " << uvs[i].y() << std::endl;
        }
    }
    if (uv_indices.size() > 0)
    {
        assert(uv_indices.size() == num_triangles);
        for (int i = 0; i < num_triangles; i++)
        {
            ofs << "f "
                << uv_indices[i][0] << "/" << uv_indices[i][0] << " "
                << uv_indices[i][1] << "/" << uv_indices[i][1] << " "
                << uv_indices[i][2] << "/" << uv_indices[i][2] << std::endl;
        }
    }
    else
    {
        for (int i = 0; i < num_triangles; i++)
        {
            ofs << "f " << indices[i][0] << " " << indices[i][1] << " " << indices[i][2] << std::endl;
        }
    }
    ofs.close();
    return true;
}
