import numpy as np

from yt.testing import assert_array_almost_equal
from yt.utilities.math_utils import (
    get_cyl_r,
    get_cyl_r_component,
    get_cyl_theta,
    get_cyl_theta_component,
    get_cyl_z,
    get_cyl_z_component,
    get_sph_phi,
    get_sph_phi_component,
    get_sph_r,
    get_sph_r_component,
    get_sph_theta,
    get_sph_theta_component,
)

# Randomly generated coordinates in the domain [[-1,1],[-1,1],-1,1]]
coords = np.array(
    [
        [-0.41503037, -0.22102472, -0.55774212],
        [0.73828247, -0.17913899, 0.64076921],
        [0.08922066, -0.94254844, -0.61774511],
        [0.10173242, -0.95789145, 0.16294352],
        [0.73186508, -0.3109153, 0.75728738],
        [0.8757989, -0.41475119, -0.57039201],
        [0.58040762, 0.81969082, 0.46759728],
        [-0.89983356, -0.9853683, -0.38355343],
    ]
).T


def test_spherical_coordinate_conversion():
    normal = [0, 0, 1]
    real_r = [
        0.72950559,
        0.99384957,
        1.13047198,
        0.97696269,
        1.09807968,
        1.12445067,
        1.10788685,
        1.38843954,
    ]
    real_theta = [
        2.44113629,
        0.87012028,
        2.14891444,
        1.4032274,
        0.80979483,
        2.10280198,
        1.13507735,
        1.85068416,
    ]
    real_phi = [
        -2.65224483,
        -0.23804243,
        -1.47641858,
        -1.46498842,
        -0.40172325,
        -0.4422801,
        0.95466734,
        -2.31085392,
    ]

    calc_r = get_sph_r(coords)
    calc_theta = get_sph_theta(coords, normal)
    calc_phi = get_sph_phi(coords, normal)

    assert_array_almost_equal(calc_r, real_r)
    assert_array_almost_equal(calc_theta, real_theta)
    assert_array_almost_equal(calc_phi, real_phi)

    normal = [1, 0, 0]
    real_theta = [
        2.17598842,
        0.73347681,
        1.49179079,
        1.46647589,
        0.8412984,
        0.67793705,
        1.0193883,
        2.27586987,
    ]
    real_phi = [
        -1.94809584,
        1.843405,
        -2.56143151,
        2.97309903,
        1.96037671,
        -2.1995016,
        0.51841239,
        -2.77038877,
    ]

    calc_theta = get_sph_theta(coords, normal)
    calc_phi = get_sph_phi(coords, normal)

    assert_array_almost_equal(calc_theta, real_theta)
    assert_array_almost_equal(calc_phi, real_phi)


def test_cylindrical_coordinate_conversion():
    normal = [0, 0, 1]
    real_r = [
        0.47021498,
        0.75970506,
        0.94676179,
        0.96327853,
        0.79516968,
        0.96904193,
        1.00437346,
        1.3344104,
    ]
    real_theta = [
        -2.65224483,
        -0.23804243,
        -1.47641858,
        -1.46498842,
        -0.40172325,
        -0.4422801,
        0.95466734,
        -2.31085392,
    ]
    real_z = [
        -0.55774212,
        0.64076921,
        -0.61774511,
        0.16294352,
        0.75728738,
        -0.57039201,
        0.46759728,
        -0.38355343,
    ]

    calc_r = get_cyl_r(coords, normal)
    calc_theta = get_cyl_theta(coords, normal)
    calc_z = get_cyl_z(coords, normal)

    assert_array_almost_equal(calc_r, real_r)
    assert_array_almost_equal(calc_theta, real_theta)
    assert_array_almost_equal(calc_z, real_z)

    normal = [1, 0, 0]
    real_r = [
        0.59994016,
        0.66533898,
        1.12694569,
        0.97165149,
        0.81862843,
        0.70524152,
        0.94368441,
        1.05738542,
    ]
    real_theta = [
        -1.94809584,
        1.843405,
        -2.56143151,
        2.97309903,
        1.96037671,
        -2.1995016,
        0.51841239,
        -2.77038877,
    ]
    real_z = [
        -0.41503037,
        0.73828247,
        0.08922066,
        0.10173242,
        0.73186508,
        0.8757989,
        0.58040762,
        -0.89983356,
    ]

    calc_r = get_cyl_r(coords, normal)
    calc_theta = get_cyl_theta(coords, normal)
    calc_z = get_cyl_z(coords, normal)

    assert_array_almost_equal(calc_r, real_r)
    assert_array_almost_equal(calc_theta, real_theta)
    assert_array_almost_equal(calc_z, real_z)


def test_spherical_coordinate_projections():
    normal = [0, 0, 1]
    theta = get_sph_theta(coords, normal)
    phi = get_sph_phi(coords, normal)
    zero = np.tile(0, coords.shape[1])

    # Purely radial field
    vecs = np.array(
        [np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]
    )
    assert_array_almost_equal(zero, get_sph_theta_component(vecs, theta, phi, normal))
    assert_array_almost_equal(zero, get_sph_phi_component(vecs, phi, normal))

    # Purely toroidal field
    vecs = np.array([-np.sin(phi), np.cos(phi), zero])
    assert_array_almost_equal(zero, get_sph_theta_component(vecs, theta, phi, normal))
    assert_array_almost_equal(zero, get_sph_r_component(vecs, theta, phi, normal))

    # Purely poloidal field
    vecs = np.array(
        [np.cos(theta) * np.cos(phi), np.cos(theta) * np.sin(phi), -np.sin(theta)]
    )
    assert_array_almost_equal(zero, get_sph_phi_component(vecs, phi, normal))
    assert_array_almost_equal(zero, get_sph_r_component(vecs, theta, phi, normal))


def test_cylindrical_coordinate_projections():
    normal = [0, 0, 1]
    theta = get_cyl_theta(coords, normal)
    z = get_cyl_z(coords, normal)
    zero = np.tile(0, coords.shape[1])

    # Purely radial field
    vecs = np.array([np.cos(theta), np.sin(theta), zero])
    assert_array_almost_equal(zero, get_cyl_theta_component(vecs, theta, normal))
    assert_array_almost_equal(zero, get_cyl_z_component(vecs, normal))

    # Purely toroidal field
    vecs = np.array([-np.sin(theta), np.cos(theta), zero])
    assert_array_almost_equal(zero, get_cyl_z_component(vecs, normal))
    assert_array_almost_equal(zero, get_cyl_r_component(vecs, theta, normal))

    # Purely z field
    vecs = np.array([zero, zero, z])
    assert_array_almost_equal(zero, get_cyl_theta_component(vecs, theta, normal))
    assert_array_almost_equal(zero, get_cyl_r_component(vecs, theta, normal))
