use interpolate::Interpolate;

/// Find the intersection of the zero level set of a linear function with a tetrahedron.
///
/// The level set is determined by the function values `u` at vertices `v`. The resulting
/// intersection emits a vertex (at most 4), and each triangle emits a face with indices into the
/// emitted vertices.
fn tetrahedron<T, FV, FF>(u: [f64; 4], v: [T; 4], mut emit_vertex: FV, mut emit_face: FF)
where
    T: Interpolate<f64> + Copy,
    FV: FnMut(T),
    FF: FnMut([u32; 3]),
{
    let u0 = u[0];
    let u1 = u[1];
    let u2 = u[2];
    let u3 = u[3];

    let v0 = v[0];
    let v1 = v[1];
    let v2 = v[2];
    let v3 = v[3];

    // avoid problems with level set going through an edge
    let tiny = 1e-15;

    // START GENERATED: generated by `scripts/tetrahedron_cuts.py`
    if u0 >= 0. {
        let u0 = u0 + tiny;
        if u1 >= 0. {
            let u1 = u1 + tiny;
            if u2 >= 0. {
                let u2 = u2 + tiny;
                if u3 >= 0. {
                } else {
                    emit_vertex(v3.interpolate(&v0, u3, u0));
                    emit_vertex(v3.interpolate(&v1, u3, u1));
                    emit_vertex(v3.interpolate(&v2, u3, u2));
                    emit_face([0, 1, 2]);
                }
            } else {
                if u3 >= 0. {
                    emit_vertex(v2.interpolate(&v0, u2, u0));
                    emit_vertex(v2.interpolate(&v1, u2, u1));
                    emit_vertex(v2.interpolate(&v3, u2, u3));
                    emit_face([0, 1, 2]);
                } else {
                    emit_vertex(v2.interpolate(&v0, u2, u0));
                    emit_vertex(v2.interpolate(&v1, u2, u1));
                    emit_vertex(v3.interpolate(&v0, u3, u0));
                    emit_vertex(v3.interpolate(&v1, u3, u1));
                    emit_face([0, 1, 2]);
                    emit_face([2, 1, 3]);
                }
            }
        } else {
            if u2 >= 0. {
                let u2 = u2 + tiny;
                if u3 >= 0. {
                    emit_vertex(v1.interpolate(&v0, u1, u0));
                    emit_vertex(v1.interpolate(&v2, u1, u2));
                    emit_vertex(v1.interpolate(&v3, u1, u3));
                    emit_face([0, 1, 2]);
                } else {
                    emit_vertex(v1.interpolate(&v0, u1, u0));
                    emit_vertex(v1.interpolate(&v2, u1, u2));
                    emit_vertex(v3.interpolate(&v0, u3, u0));
                    emit_vertex(v3.interpolate(&v2, u3, u2));
                    emit_face([0, 1, 2]);
                    emit_face([2, 1, 3]);
                }
            } else {
                if u3 >= 0. {
                    emit_vertex(v1.interpolate(&v0, u1, u0));
                    emit_vertex(v1.interpolate(&v3, u1, u3));
                    emit_vertex(v2.interpolate(&v0, u2, u0));
                    emit_vertex(v2.interpolate(&v3, u2, u3));
                    emit_face([0, 1, 2]);
                    emit_face([2, 1, 3]);
                } else {
                    emit_vertex(v1.interpolate(&v0, u1, u0));
                    emit_vertex(v2.interpolate(&v0, u2, u0));
                    emit_vertex(v3.interpolate(&v0, u3, u0));
                    emit_face([0, 1, 2]);
                }
            }
        }
    } else {
        if u1 >= 0. {
            let u1 = u1 + tiny;
            if u2 >= 0. {
                let u2 = u2 + tiny;
                if u3 >= 0. {
                    emit_vertex(v0.interpolate(&v1, u0, u1));
                    emit_vertex(v0.interpolate(&v2, u0, u2));
                    emit_vertex(v0.interpolate(&v3, u0, u3));
                    emit_face([0, 1, 2]);
                } else {
                    emit_vertex(v0.interpolate(&v1, u0, u1));
                    emit_vertex(v0.interpolate(&v2, u0, u2));
                    emit_vertex(v3.interpolate(&v1, u3, u1));
                    emit_vertex(v3.interpolate(&v2, u3, u2));
                    emit_face([0, 1, 2]);
                    emit_face([2, 1, 3]);
                }
            } else {
                if u3 >= 0. {
                    emit_vertex(v0.interpolate(&v1, u0, u1));
                    emit_vertex(v0.interpolate(&v3, u0, u3));
                    emit_vertex(v2.interpolate(&v1, u2, u1));
                    emit_vertex(v2.interpolate(&v3, u2, u3));
                    emit_face([0, 1, 2]);
                    emit_face([2, 1, 3]);
                } else {
                    emit_vertex(v0.interpolate(&v1, u0, u1));
                    emit_vertex(v2.interpolate(&v1, u2, u1));
                    emit_vertex(v3.interpolate(&v1, u3, u1));
                    emit_face([0, 1, 2]);
                }
            }
        } else {
            if u2 >= 0. {
                let u2 = u2 + tiny;
                if u3 >= 0. {
                    emit_vertex(v0.interpolate(&v2, u0, u2));
                    emit_vertex(v0.interpolate(&v3, u0, u3));
                    emit_vertex(v1.interpolate(&v2, u1, u2));
                    emit_vertex(v1.interpolate(&v3, u1, u3));
                    emit_face([0, 1, 2]);
                    emit_face([2, 1, 3]);
                } else {
                    emit_vertex(v0.interpolate(&v2, u0, u2));
                    emit_vertex(v1.interpolate(&v2, u1, u2));
                    emit_vertex(v3.interpolate(&v2, u3, u2));
                    emit_face([0, 1, 2]);
                }
            } else {
                if u3 >= 0. {
                    emit_vertex(v0.interpolate(&v3, u0, u3));
                    emit_vertex(v1.interpolate(&v3, u1, u3));
                    emit_vertex(v2.interpolate(&v3, u2, u3));
                    emit_face([0, 1, 2]);
                } else {
                }
            }
        }
    }
    // END GENERATED
}

/// Finds the isosurface at `level` of a function given by its values `u` on a regular grid  using
/// the marching tetrahedra algorithm.
///
/// `dim` is the dimension of the array `u` assumed to be in _row-major order_ (C order).
///
/// Returns vertices, faces and normals of the generated triangular mesh.
pub fn marching_tetrahedra(
    u: &[f64],
    dim: (usize, usize, usize),
    level: f64,
) -> (Vec<[f64; 3]>, Vec<[u32; 3]>, Vec<[f64; 3]>) {
    let (verts, faces, normals, _) =
        marching_tetrahedra_with_data(u, dim, level, &vec![(); u.len()]);

    (verts, faces, normals)
}

/// As `marching_tetrahedra`, but also linearly interpolates the provided data for each vertex.
pub fn marching_tetrahedra_with_data<T>(
    u: &[f64],
    dim: (usize, usize, usize),
    level: f64,
    data: &[T],
) -> (Vec<[f64; 3]>, Vec<[u32; 3]>, Vec<[f64; 3]>, Vec<T>)
where
    T: Interpolate<f64> + Default + Copy,
{
    let (ni, nj, nk) = dim;
    assert_eq!(ni * nj * nk, u.len());
    assert_eq!(ni * nj * nk, data.len());

    let mut verts: Vec<[f64; 3]> = Vec::new();
    let mut normals: Vec<[f64; 3]> = Vec::new();
    let mut faces: Vec<[u32; 3]> = Vec::new();
    let mut interp_data: Vec<T> = Vec::new();

    // permutations of [0, 1, 2]
    let perms = [
        [0, 1, 2],
        [0, 2, 1],
        [1, 0, 2],
        [1, 2, 0],
        [2, 0, 1],
        [2, 1, 0],
    ];

    // inverse permutation (index where the element goes)
    let inv_perms = {
        let mut inv_perms = [[0; 3]; 6];
        for i in 0..6 {
            for j in 0..3 {
                inv_perms[i][perms[i][j]] = j;
            }
        }
        inv_perms
    };

    let strides = [nj * nk, nk, 1];

    let vert_offsets = {
        let mut vert_offsets = [0; 8];

        let mut c = 0;
        for i in 0..2 {
            for j in 0..2 {
                for k in 0..2 {
                    vert_offsets[c] = i * nj * nk + j * nk + k;
                    c += 1;
                }
            }
        }
        vert_offsets
    };

    for i in 1..ni {
        for j in 1..nj {
            for k in 1..nk {
                // split each cube into 6 tetrahedra, emit triangles
                // vertex position
                let s = (i - 1) * nj * nk + (j - 1) * nk + (k - 1);
                let ps = [(i - 1) as f64, (j - 1) as f64, (k - 1) as f64];

                let n_above = vert_offsets
                    .iter()
                    .filter(|&offset| u[s + offset] >= level)
                    .count();

                if n_above == 0 || n_above == 8 {
                    continue;
                }

                for (perm, inv_perm) in perms.iter().zip(inv_perms.iter()) {
                    // find tetrahedron data by walking along the edges of a cube
                    // in the order given by the permutation
                    let (us, vs) = {
                        let mut us = [0.; 4];
                        let mut vs = [([0.; 3], T::default()); 4];
                        let mut vi = s;
                        let mut vp = ps;
                        us[0] = u[vi] - level;
                        vs[0] = (vp, data[vi]);
                        for m in 0..3 {
                            let t = perm[m];
                            vi += strides[t];
                            vp[t] += 1.;
                            us[m + 1] = u[vi] - level;
                            vs[m + 1] = (vp, data[vi]);
                        }
                        (us, vs)
                    };

                    let cur = verts.len() as u32;
                    tetrahedron(
                        us,
                        vs,
                        |(v, d)| {
                            verts.push(v);
                            interp_data.push(d);
                        },
                        |f| {
                            faces.push([f[0] + cur, f[1] + cur, f[2] + cur]);
                        },
                    );

                    // normals
                    let n = [us[1] - us[0], us[2] - us[1], us[3] - us[2]];
                    let n = [n[inv_perm[0]], n[inv_perm[1]], n[inv_perm[2]]];

                    for _ in 0..verts.len() - normals.len() {
                        normals.push(n);
                    }
                }
            }
        }
    }

    (verts, faces, normals, interp_data)
}

/// Emits the triangles of the level set of u intersecting tetrahedra of the mesh.
///
/// For each triangle, emits the coordinates of its vertices and the linearly interpolated data
/// at these vertices.
///
/// The coordinate system is chosen so that the node (i, j, k) with index i * dim.1 * dim.2 + j *
/// dim.2 + k has coordinate
/// (i, j, k).
pub fn marching_tetrahedra_with_data_emit<F, D>(
    u: &[f64],
    data: &[D],
    dim: (usize, usize, usize),
    level: f64,
    mut emit: F,
) where
    F: FnMut([[f64; 3]; 3], [D; 3]),
    D: Interpolate<f64> + Default + Copy,
{
    let (ni, nj, nk) = dim;
    assert_eq!(ni * nj * nk, u.len());
    assert_eq!(ni * nj * nk, data.len());

    let mut verts: Vec<[f64; 3]> = Vec::with_capacity(4);
    let mut faces: Vec<[u32; 3]> = Vec::with_capacity(2);
    let mut interp_data: Vec<D> = Vec::with_capacity(4);

    // permutations of [0, 1, 2]
    let perms = [
        [0, 1, 2],
        [0, 2, 1],
        [1, 0, 2],
        [1, 2, 0],
        [2, 0, 1],
        [2, 1, 0],
    ];

    let strides = [nj * nk, nk, 1];

    let vert_offsets = {
        let mut vert_offsets = [0; 8];

        let mut c = 0;
        for i in 0..2 {
            for j in 0..2 {
                for k in 0..2 {
                    vert_offsets[c] = i * nj * nk + j * nk + k;
                    c += 1;
                }
            }
        }
        vert_offsets
    };

    for i in 1..ni {
        for j in 1..nj {
            for k in 1..nk {
                // split each cube into 6 tetrahedra, emit triangles
                // vertex position
                let s = (i - 1) * nj * nk + (j - 1) * nk + (k - 1);
                let ps = [(i - 1) as f64, (j - 1) as f64, (k - 1) as f64];

                let n_above = vert_offsets
                    .iter()
                    .filter(|&offset| u[s + offset] >= level)
                    .count();

                if n_above == 0 || n_above == 8 {
                    continue;
                }

                for perm in &perms {
                    // find tetrahedron data by walking along the edges of a cube
                    // in the order given by the permutation
                    let (us, vs) = {
                        let mut us = [0.; 4];
                        let mut vs = [([0.; 3], D::default()); 4];
                        let mut vi = s;
                        let mut vp = ps;
                        us[0] = u[vi] - level;
                        vs[0] = (vp, data[vi]);
                        for m in 0..3 {
                            let t = perm[m];
                            vi += strides[t];
                            vp[t] += 1.;
                            us[m + 1] = u[vi] - level;
                            vs[m + 1] = (vp, data[vi]);
                        }
                        (us, vs)
                    };

                    verts.clear();
                    interp_data.clear();
                    faces.clear();
                    tetrahedron(
                        us,
                        vs,
                        |(v, d)| {
                            verts.push(v);
                            interp_data.push(d);
                        },
                        |f| {
                            faces.push(f);
                        },
                    );
                    for f in &faces {
                        emit(
                            [
                                verts[f[0] as usize],
                                verts[f[1] as usize],
                                verts[f[2] as usize],
                            ],
                            [
                                interp_data[f[0] as usize],
                                interp_data[f[1] as usize],
                                interp_data[f[2] as usize],
                            ],
                        );
                    }
                }
            }
        }
    }
}
