use crate::mat3::Mat3;
use crate::math::{atan, distance, ln, mag3};
use crate::sources::element::{edge_csys_fast, transform_fast};
use crate::sources::{edge_csys, transform};
use crate::vec3::Vec3;
use std::f64::consts::PI;

// Refer to the gmsh or ansys documentation for node numbering on a tet element:
// https://gmsh.info/doc/texinfo/gmsh.html#Node-ordering
// https://ansyshelp.ansys.com/public/account/secured?returnurl=//////////Views/Secured/corp/v242/en/ans_elem/Hlp_E_SOLID236.html
const NODE_MAP: [[[usize; 3]; 3]; 4] = [
    [[0, 2, 1], [2, 1, 0], [1, 0, 2]], // back face
    [[0, 3, 2], [3, 2, 0], [2, 0, 3]], // left face
    [[1, 2, 3], [2, 3, 1], [3, 1, 2]], // right face
    [[0, 1, 3], [1, 3, 0], [3, 0, 1]], // bottom face
];

const INV_4PI: f64 = 1.0 / (4.0 * PI);

// Compute the indefinite edge integral of a finite element given the coordinates
// of the target point in the edge reference frame.
// Uses inlined/vectorizable approximate functions for `ln()` and `atan2()`.
// `mag3()` is also inlined and uses fma instructions.
#[inline(always)]
fn edge_integral(x: f64, y: f64, z: f64) -> f64 {
    let r = mag3(x, y, z);
    let x_plus_r = x + r;
    let ln_x_plus_r = if x_plus_r.abs() > 1e-8 {
        ln(x_plus_r)
    } else {
        0.0
    };
    let y_inv = 1.0 / y;
    let za = z.abs();
    let za_over_y = if y.abs() > 1e-12 * r { za * y_inv } else { 0.0 };
    let atan_1 = if y.abs() > 1e-12 * r && r > 1e-8 {
        atan(x * za_over_y / r)
    } else {
        0.0
    };
    let atan_2 = if y.abs() > 1e-12 * r {
        atan(x * y_inv)
    } else {
        0.0
    };

    ln_x_plus_r + za_over_y * (atan_1 - atan_2)
}

/// Compute magnetic field intensity generated by a 4-node tetrahedral finite element
///
/// This function operates on n-number of target points and
/// is vectorized for efficient computation.
///
/// # Arguments
/// - `nodes`: (m), coordinates of each corner node in 3D space
/// - `jdensity`: (A/m2), current density vector, assumed constant throughout the element
/// - `targets`: (m),  target locations in 3d space
/// - `f`: workspace allocated by the caller
/// - `h`: (A/m), preallocated arrays for h-field results
///
/// # Reference  
///
/// CRYO-06-034: "MAGNUM - A Fortran Library for the Calculation of Magnetic Configurations"  
/// [https://supermagnet.sourceforge.io/notes/CRYO-06-034.pdf](https://supermagnet.sourceforge.io/notes/CRYO-06-034.pdf)
pub fn h_field_tet4(
    nodes: &[Vec3; 4],
    jdensity: &Vec3,
    targets: (&[f64], &[f64], &[f64]),
    f: &mut [Vec3],
    h: (&mut [f64], &mut [f64], &mut [f64]),
) {
    assert_eq!(f.len(), targets.0.len());

    // Allocate memory for the result of the sum before crossing with J
    // TODO: replace?
    // let mut f: Vec<Vec3> = vec![Vec3([0.0; 3]); targets.0.len()];

    for face in 0..4 {
        for edge in 0..3 {
            // normal on face

            // Node numbers from the map
            let n0: usize = NODE_MAP[face][edge][0];
            let n1: usize = NODE_MAP[face][edge][1];
            let n2: usize = NODE_MAP[face][edge][2];

            // Local coordinate system tensor; transform the target point into the local system
            let (xhat, yhat, zhat) = edge_csys_fast(&nodes[n0], &nodes[n1], &nodes[n2]);
            let ni: Vec3 = zhat;
            let xpq2 = (nodes[n1] - nodes[n0]).mag(); // length of the edge

            // Compute the effect of this edge by transforming each target into the local reference frame
            for i in 0..targets.0.len() {
                let target_local: Vec3 = Vec3([
                    targets.0[i] - nodes[n1][0],
                    targets.1[i] - nodes[n1][1],
                    targets.2[i] - nodes[n1][2],
                ]);
                let p = transform_fast(&target_local, &xhat, &yhat, &zhat);
                let s_face = p[1]
                    * (edge_integral(xpq2 - p[0], p[1], p[2]) - edge_integral(-p[0], p[1], p[2]));
                f[i] += ni * s_face;
            }
        }
    }

    // Cross with J/4pi and accumulate
    for i in 0..targets.0.len() {
        let _h = jdensity.cross(&f[i]) * INV_4PI;
        h.0[i] += _h[0];
        h.1[i] += _h[1];
        h.2[i] += _h[2];
    }
}

/// DEPRECATED
/// Compute the magnetic flux density generated by an arbitrary tetrahedron in 3D space
///
/// # Arguments
/// - `x`: x coordinates of each corner node in 3D space
/// - `y`: y coordinates of each corner node in 3D space
/// - `z`: z coordinates of each corner node in 3D space
/// - `jdensity`: current density vector, assumed constant throughout the element
///
/// # Returns
///     magnetic flux density vector generated by the tetrahedron
///
/// # Reference  
///
/// CRYO-06-034: "MAGNUM - A Fortran Library for the Calculation of Magnetic Configurations"
/// https://supermagnet.sourceforge.io/notes/CRYO-06-034.pdf
///
#[allow(dead_code)]
fn hfield_tetrahedron(
    nx: &[f64],      // 4-length
    ny: &[f64],      // 4-length
    nz: &[f64],      // 4-length
    jdensity: &Vec3, // 3-length
    target: &Vec3,   // 3-length
) -> Vec3 {
    // Compute a total of 12 line integrals, along each edge (3 on each face)
    // Refer to the gmsh or ansys documentation for node numbering on a tet element:
    // https://gmsh.info/doc/texinfo/gmsh.html#Node-ordering
    // https://ansyshelp.ansys.com/public/account/secured?returnurl=//////////Views/Secured/corp/v242/en/ans_elem/Hlp_E_SOLID236.html
    // let NODE_MAP: [[[usize;3];3];4] = [
    //     [ [0,2,1], [2,1,0], [1,0,2] ],      // back face
    //     [ [0,3,2], [3,2,0], [2,0,3] ],      // left face
    //     [ [1,2,3], [2,3,1], [3,1,2] ],      // right face
    //     [ [0,1,3], [1,3,0], [3,0,1] ],      // bottom face
    // ];

    let mut s = [0.0; 3];
    // outer loop over faces
    for i in 0..4 {
        let mut ni = [0.0; 3]; // normal on face 
        let mut s_face = 0.0;

        for j in 0..3 {
            // Node numbers from the map
            let n0 = NODE_MAP[i][j][0];
            let n1 = NODE_MAP[i][j][1];
            let n2 = NODE_MAP[i][j][2];

            // Local coordinate system tensor; transform the target point into the local system
            let (xhat, yhat, zhat) = edge_csys(
                &[nx[n0], nx[n1], nx[n2]],
                &[ny[n0], ny[n1], ny[n2]],
                &[nz[n0], nz[n1], nz[n2]],
            );
            ni = zhat;
            let target_local = [target[0] - nx[n1], target[1] - ny[n1], target[2] - nz[n1]];
            let p = transform(&target_local, &xhat, &yhat, &zhat);
            let xpq2 = distance(&[nx[n1], ny[n1], nz[n1]], &[nx[n0], ny[n0], nz[n0]]); // length of the edge, basically

            s_face +=
                p[1] * (edge_integral(xpq2 - p[0], p[1], p[2]) - edge_integral(-p[0], p[1], p[2]));
        }
        s[0] += ni[0] * s_face;
        s[1] += ni[1] * s_face;
        s[2] += ni[2] * s_face;
    }

    jdensity.cross(&Vec3(s)) * INV_4PI
}

#[allow(unused)]
fn tet_surface_integral(nodes: &[Vec3; 4], target: &Vec3) -> Vec3 {
    let mut s = Vec3([0.0; 3]);

    // Node mapping as previous
    // let NODE_MAP: [[[usize;3];3];4] = [
    //     [ [0,2,1], [2,1,0], [1,0,2] ],      // back face
    //     [ [0,3,2], [3,2,0], [2,0,3] ],      // left face
    //     [ [1,2,3], [2,3,1], [3,1,2] ],      // right face
    //     [ [0,1,3], [1,3,0], [3,0,1] ],      // bottom face
    // ];

    // Outer loop over each face
    for i in 0..4 {
        let mut ni = [0.0; 3]; // normal on face 
        let mut s_face = 0.0;

        for j in 0..3 {
            // Node numbers from the map
            let n0 = NODE_MAP[i][j][0];
            let n1 = NODE_MAP[i][j][1];
            let n2 = NODE_MAP[i][j][2];

            // Local coordinate system tensor; transform the target point into the local system
            let (xhat, yhat, zhat) = edge_csys(
                &[nodes[n0][0], nodes[n1][0], nodes[n2][0]],
                &[nodes[n0][1], nodes[n1][1], nodes[n2][1]],
                &[nodes[n0][2], nodes[n1][2], nodes[n2][2]],
            );
            ni = zhat;
            let target_local = [
                target[0] - nodes[n1][0],
                target[1] - nodes[n1][1],
                target[2] - nodes[n1][2],
            ];
            let p = transform(&target_local, &xhat, &yhat, &zhat);
            let xppq2 = distance(nodes[n1].to_slice(), nodes[n0].to_slice()); // length of the edge, basically

            s_face +=
                p[1] * (edge_integral(xppq2 - p[0], p[1], p[2]) - edge_integral(-p[0], p[1], p[2]));
        }
        s[0] += ni[0] * s_face;
        s[1] += ni[1] * s_face;
        s[2] += ni[2] * s_face;
    }
    s
}

/// Compute the h field generated by a magnetized tetrahedral element
///
/// This computes equations 58 and 60 in CRYO-06-034, assuming the target point is outside the element:  
/// $$ H(\vec{r}) = \frac{-M}{4\pi} \cdot \nabla \int_V \frac{\vec{r}}{r^3}dV $$
pub fn hmag_tet4(
    nodes: &[Vec3; 4], // source nodes
    mvector: &Vec3,    // source element
    j_invt: &[Mat3],   // at each *target* element
    target_nodes: (&[f64], &[f64], &[f64]),
    target_element_connectivity: &[[u32; 4]],
    f: &mut [Vec3],
    h: (&mut [f64], &mut [f64], &mut [f64]),
) {
    // Temporary storage is for nodes, hmag field is calculated at elements
    assert_eq!(f.len(), target_nodes.0.len());
    assert_eq!(h.0.len(), target_element_connectivity.len());

    // Ensure f is zeroed
    for v in f.iter_mut() {
        *v = Vec3::default();
    }

    // Compute the value of the all edge integrals at each target node
    for face in 0..4 {
        for edge in 0..3 {
            // Node numbers from the map
            let n0: usize = NODE_MAP[face][edge][0];
            let n1: usize = NODE_MAP[face][edge][1];
            let n2: usize = NODE_MAP[face][edge][2];

            // Local coordinate system tensor; transform the target point into the local system
            let (xhat, yhat, zhat) = edge_csys_fast(&nodes[n0], &nodes[n1], &nodes[n2]);
            let ni: Vec3 = zhat; // normal on face 
            let xpq2 = (nodes[n1] - nodes[n0]).mag(); // length of the edge

            // Compute the effect of this edge by transforming each target into the local reference frame
            for i in 0..target_nodes.0.len() {
                let target_local: Vec3 = Vec3([
                    target_nodes.0[i] - nodes[n1][0],
                    target_nodes.1[i] - nodes[n1][1],
                    target_nodes.2[i] - nodes[n1][2],
                ]);
                let p = transform_fast(&target_local, &xhat, &yhat, &zhat);
                let s_face = p[1]
                    * (edge_integral(xpq2 - p[0], p[1], p[2]) - edge_integral(-p[0], p[1], p[2]));
                f[i] += ni * s_face;
            }
        }
    }

    // Take gradient at each element, dot with -M/4pi, and accumulate
    let scalar: f64 = -1.0 / (4.0 * PI);
    let m_4pi: Vec3 = *mvector * scalar;
    let n_elements = target_element_connectivity.len();
    for i in 0..n_elements {
        let n_idx = target_element_connectivity[i];

        // Edge/surface integral at each of the four nodes
        let i0 = f[n_idx[0] as usize];
        let i1 = f[n_idx[1] as usize];
        let i2 = f[n_idx[2] as usize];
        let i3 = f[n_idx[3] as usize];

        // Compute individual `g` vectors for each component of the gradient evaluation
        let gx = Vec3([i1[0] - i0[0], i2[0] - i0[0], i3[0] - i0[0]]);
        let gy = Vec3([i1[1] - i0[1], i2[1] - i0[1], i3[1] - i0[1]]);
        let gz = Vec3([i1[2] - i0[2], i2[2] - i0[2], i3[2] - i0[2]]);

        // Compute gradient
        let grad_ix = j_invt[i].mul_vec(&gx);
        let grad_iy = j_invt[i].mul_vec(&gy);
        let grad_iz = j_invt[i].mul_vec(&gz);

        let jacobian = Mat3::from_cols(&grad_ix, &grad_iy, &grad_iz);
        let _h_mag = jacobian.mul_vec(&m_4pi);
        h.0[i] += _h_mag[0];
        h.1[i] += _h_mag[1];
        h.2[i] += _h_mag[2];
    }
}
