use faer::{
    ColRef, Side,
    linalg::solvers::Solve,
    sparse::{SparseColMat, Triplet},
};

use crate::{graph::Graph, triangulation::causal_triangulation2d::CausalTriangulation2D};

fn compute_boundary_shift_sum(
    triangulation: &CausalTriangulation2D,
    path: &[usize],
) -> (Vec<f64>, Vec<f64>) {
    let origin_triangle = *path.last().expect("Path should not be empty");
    let mut xbound = vec![0.0; triangulation.num_vertices()];
    let mut prev = path[0];
    for next in path.iter().rev().copied() {
        let triangle = &triangulation.triangles[prev];
        let index = triangle
            .neighbours
            .iter()
            .position(|&nbr| nbr == next)
            .expect("Next path element must be neighbour for valid paths");
        let left = triangle.vertices[(index + 2) % 3];
        let right = triangle.vertices[(index + 1) % 3];
        xbound[left] += 1.0;
        xbound[right] -= 1.0;
        prev = next;
    }

    let mut ybound = vec![0.0; triangulation.num_vertices()];
    let mut prev = origin_triangle;
    loop {
        let next = triangulation.get_right_neighbour(prev);

        let triangle = &triangulation.triangles[prev];
        let index = triangle
            .neighbours
            .iter()
            .position(|&nbr| nbr == next)
            .expect("Next path element must be neighbour for valid paths");
        let up = triangle.vertices[(index + 2) % 3];
        let down = triangle.vertices[(index + 1) % 3];
        ybound[down] += 1.0;
        ybound[up] -= 1.0;

        if next == origin_triangle {
            break;
        }
        prev = next;
    }

    (xbound, ybound)
}

/// Construct the graph Laplacian
///
/// Note we exclude the origin as that will have fixed coordinates
fn construct_laplacian(
    triangulation: &CausalTriangulation2D,
    graph: &Graph,
    origin_vertex: usize,
) -> SparseColMat<usize, f64> {
    // Laplacian matrix will be `n x n`
    let n = graph.len();
    // Construct Laplacian with triplets
    let mut triplets = Vec::with_capacity(graph.len() + triangulation.triangles.len() * 3 / 2);
    // Add diagonal based on degrees
    triplets.extend(graph.nodes().iter().enumerate().map(|(i, node)| {
        Triplet::new(
            i,
            i,
            if i == origin_vertex {
                1.0
            } else {
                node.degree() as f64
            },
        )
    }));
    // Add -1 for all neighbours
    triplets.extend(
        graph
            .nodes()
            .iter()
            .enumerate()
            .filter(|&(label, _)| label != origin_vertex)
            .flat_map(|(label, node)| {
                node.get_neighbours()
                    .iter()
                    .copied()
                    .filter(|&nbr| nbr != origin_vertex)
                    .map(move |nbr| Triplet::new(label, nbr, -1.0))
            }),
    );

    // Use `faer` to construct sparse matrix
    SparseColMat::<usize, f64>::try_new_from_triplets(n, n, &triplets)
        .expect("Laplacian should be proper")
}

/// Construct a toroidal Tutte embedding for the vertex graph of the `triangulation`
///
/// The coordinates are normalized to a unit torus: [0, 1) x [0, 1)
pub fn construct_tutte_embedding(
    triangulation: &CausalTriangulation2D,
    origin_vertex: usize,
) -> Vec<(f64, f64)> {
    // Construct time-like loop
    let origin_triangle = 0;
    let path = triangulation.identify_timelike_loop(origin_triangle);

    let vertex_graph = triangulation.vertex_graph();
    construct_tutte_embedding_with_graph_with_path(
        triangulation,
        &vertex_graph,
        origin_vertex,
        &path,
    )
}

/// Construct a toroidal Tutte embedding for the `vertex_graph` of the `triangulation`
///
/// The coordinates are normalized to a unit torus: [0, 1) x [0, 1).
/// `vertex_graph` must be constructed with `triangulation.vertex_graph()`
/// `path` must be constructed with `triangulation.identify_timelike_loop(any_origin)`
pub fn construct_tutte_embedding_with_graph_with_path(
    triangulation: &CausalTriangulation2D,
    vertex_graph: &Graph,
    origin_vertex: usize,
    path: &[usize],
) -> Vec<(f64, f64)> {
    // Pick any origin triangle, should be independent of choice
    let mut bound_shift_sum = compute_boundary_shift_sum(triangulation, path);
    // Fix origin in position
    bound_shift_sum.0[origin_vertex] = 0.0;
    bound_shift_sum.1[origin_vertex] = 0.0;
    let bx = ColRef::from_slice(&bound_shift_sum.0[..]).as_mat();
    let by = ColRef::from_slice(&bound_shift_sum.1[..]).as_mat();

    // Use Cholesky decomposition to solve linear triangulation system
    let llt = construct_laplacian(triangulation, vertex_graph, origin_vertex)
        .sp_cholesky(Side::Lower)
        .expect("Laplacian must be positive definite hermitian");

    // Solve sparse linear system
    let coordx = llt.solve(&bx);
    let coordy = llt.solve(&by);

    // // Shift coordinates in [0, 1) range
    // std::iter::zip(
    //     coordx.col_as_slice(0).iter().map(|v| {
    //         let frac = v.fract();
    //         if frac < 0.0 { 1.0 + frac } else { frac }
    //     }),
    //     coordy.col_as_slice(0).iter().map(|v| {
    //         let frac = v.fract();
    //         if frac < 0.0 { 1.0 + frac } else { frac }
    //     }),
    // )
    // .collect()
    std::iter::zip(
        coordx.col_as_slice(0).iter().copied(),
        coordy.col_as_slice(0).iter().copied(),
    )
    .collect()
}
