from bfieldtools import contour

import pytest
import numpy as np

from numpy.testing import (
    assert_array_almost_equal,
    assert_array_equal,
    assert_allclose,
    assert_equal,
)


def setup_contour_input():
    """Load example mesh and create scalars data"""
    from bfieldtools.utils import load_example_mesh

    mesh = load_example_mesh("unit_disc")

    r = np.linalg.norm(mesh.vertices, axis=1)
    scalars = r**2

    return mesh, scalars


def compare_contour_direction_to_rotated_gradient(mesh, scalars, polyline):
    """Check inner product between the polyline edges
        the rotated gradient vectors closes to the initial points of those
        edges. These should point to the same direction.


    Parameters
    ----------
    mesh : trimesh
        mesh
    scalars : ndarray
        stream function
    polyline : ndarray (N, 3)
        coordinates of points representing a polyline

    """
    from bfieldtools.mesh_calculus import gradient

    edges = polyline[1:] - polyline[:-1]
    g = gradient(scalars, mesh, rotated=True).T
    fc = mesh.vertices[mesh.faces].mean(axis=1)

    norm = np.linalg.norm
    p = polyline
    # Find closest face centers to polyline nodes
    f_inds = np.argmin(norm(p[:, None, :] - fc[None, :, :], axis=-1), axis=1)

    g_poly = g[f_inds]
    assert np.all(np.sum(g_poly[:-1] * edges, axis=1) > 0)


def compare_magnetic_field_directions(mesh, scalars, polys, test_point):
    """Check direction of magnetic calculated from mesh and stream function
        versus the direction of the field generated by contours


    Parameters
    ----------
    mesh : trimesh
        mesh
    scalars : ndarray
        stream function
    polys : list of ndarray (N, 3)
        coordinates of points representing a many polylines
    test_point: ndarray(1, 3)
        coordinates of a test point

    """
    from bfieldtools.line_magnetics import magnetic_field
    from bfieldtools.mesh_conductor import magnetic_field_coupling

    B_mesh = magnetic_field_coupling(mesh, test_point)
    b_stream_func = B_mesh @ scalars
    b_polys = np.array([magnetic_field(p, test_point) for p in polys])
    b_poly_sum = b_polys.sum(axis=0)

    assert np.sum(b_poly_sum * b_stream_func) > 0


def test_scalar_contour_direction():
    """Test the direction of scalar_contour and the field direction
    generated by the scalar contour
    """
    mesh, scalars = setup_contour_input()
    N = 10

    polys, vals = contour.scalar_contour(
        mesh, scalars, N_contours=N, return_values=True
    )

    compare_contour_direction_to_rotated_gradient(mesh, scalars, polys[-1])

    test_point = np.array([[0, 0, 1]])
    compare_magnetic_field_directions(mesh, scalars, polys, test_point)


def test_simplify_contour():

    assert True


if __name__ == "__main__":
    test_scalar_contour_direction()
