import pytest
import numpy as np

from dynasor.core.reciprocal import calc_rho_q, calc_rho_j_q
from dynasor.core.rho_j_q_numba import rho_q_single


# Setup test arrays and inputs

@pytest.fixture
def xvq():
    Nx = 4
    Nq = 5
    np.random.seed(42)
    x = np.random.normal(size=(Nx, 3))
    v = np.random.normal(size=(Nx, 3))
    q = np.random.normal(size=(Nq, 3))
    return x, v, q


def test_rho_j_q(xvq):
    """Regression test: rho from calc_rho_j_q matches calc_rho_q and known values."""
    x, v, q = xvq
    rho_q, rho_j_q = calc_rho_j_q(x, v, q)
    rho_q2 = calc_rho_q(x, q)
    assert np.allclose(rho_q, rho_q2)
    Nx = x.shape[0]

    # Expected values computed from a reference run with numpy.random.seed(42)
    rho_q_target = Nx**0.5 * np.array([1.6509947-0.7272682j,
                                       1.74291595+0.81511274j,
                                       1.16879217-0.98689153j,
                                       1.00661299-1.40743052j,
                                       1.06588947+0.44672131j])
    rho_j_q_target = Nx**0.5 * np.array(
        [[-0.72463296+0.13571632j, -1.56724203+1.27085398j, -0.29899287+0.34897803j],
         [-0.60871115-0.37439687j, -1.96631866-0.60202257j, -0.62226018-0.17672867j],
         [-0.40649785+0.10510835j, -1.53524612+0.6491588j,  -0.31563418+1.30276366j],
         [-0.50810342+0.40748936j, -0.76088817+1.75737831j, -0.11795971+0.63627665j],
         [-0.4026586-0.18057006j, -1.68535262+0.40104027j, -0.01598231-0.55236761j]])

    assert np.allclose(rho_q, rho_q_target)
    assert np.allclose(rho_j_q, rho_j_q_target)


# Analytic tests for calc_rho_q
# rho(q) = sum_i exp(i * q . x_i)

def test_calc_rho_q_at_q_zero_equals_number_of_particles():
    """rho(q=0) = N for any positions, since exp(0) = 1 for each particle."""
    N = 7
    x = np.random.default_rng(0).normal(size=(N, 3))
    q = np.zeros((1, 3))
    rho = calc_rho_q(x, q)
    assert np.allclose(rho, [N])


def test_calc_rho_q_single_particle_at_origin():
    """Single particle at origin: exp(i * q . 0) = 1 for all q."""
    x = np.zeros((1, 3))
    q = np.array([[1.0, 2.0, 3.0],
                  [0.5, -1.0, 0.0],
                  [0.0, 0.0, 1.0]])
    rho = calc_rho_q(x, q)
    assert np.allclose(rho, np.ones(3, dtype=complex))


def test_calc_rho_q_single_particle_at_known_position():
    """Single particle at x: rho(q) = exp(i * q . x)."""
    x = np.array([[1.0, 2.0, 3.0]])
    q = np.array([[0.5, 0.0, 0.0],
                  [0.0, 1.0, 0.0],
                  [1.0, 1.0, 1.0]])
    rho = calc_rho_q(x, q)
    expected = np.array([np.exp(1j * 0.5),
                         np.exp(1j * 2.0),
                         np.exp(1j * 6.0)])
    assert np.allclose(rho, expected)


def test_calc_rho_q_two_particles_symmetric_positions():
    """Two particles at +x and -x: rho(q) = 2*cos(q.x), purely real."""
    x0 = np.array([1.5, 0.0, 0.0])
    x = np.array([x0, -x0])
    q = np.array([[1.0, 0.0, 0.0],
                  [2.0, 0.0, 0.0],
                  [0.0, 1.0, 0.0]])  # last q is perpendicular → cos(0)*2 = 2
    rho = calc_rho_q(x, q)
    expected = np.array([2 * np.cos(1.0 * 1.5),
                         2 * np.cos(2.0 * 1.5),
                         2 * np.cos(0.0)])
    assert np.allclose(rho.imag, 0.0)
    assert np.allclose(rho.real, expected)


def test_calc_rho_q_conjugate_symmetry():
    """rho(-q) = conj(rho(q)) for real particle positions."""
    rng = np.random.default_rng(1)
    x = rng.normal(size=(6, 3))
    q = rng.normal(size=(5, 3))
    rho_plus = calc_rho_q(x, q)
    rho_minus = calc_rho_q(x, -q)
    assert np.allclose(rho_minus, np.conj(rho_plus))


def test_calc_rho_q_superposition():
    """rho of combined particle set equals sum of individual rho values."""
    rng = np.random.default_rng(2)
    x1 = rng.normal(size=(3, 3))
    x2 = rng.normal(size=(5, 3))
    x_combined = np.vstack([x1, x2])
    q = rng.normal(size=(4, 3))

    rho1 = calc_rho_q(x1, q)
    rho2 = calc_rho_q(x2, q)
    rho_combined = calc_rho_q(x_combined, q)
    assert np.allclose(rho_combined, rho1 + rho2)


def test_calc_rho_q_magnitude_single_particle():
    """|rho(q)| = 1 for a single particle at any position and any q."""
    rng = np.random.default_rng(3)
    x = rng.normal(size=(1, 3))
    q = rng.normal(size=(10, 3))
    rho = calc_rho_q(x, q)
    assert np.allclose(np.abs(rho), 1.0)


def test_calc_rho_q_all_particles_at_same_position():
    """N particles at the same position: rho(q) = N * exp(i * q . x0)."""
    N = 5
    x0 = np.array([1.0, -2.0, 0.5])
    x = np.tile(x0, (N, 1))
    q = np.array([[1.0, 0.0, 0.0],
                  [0.0, 1.0, 0.0],
                  [1.0, 1.0, 1.0]])
    rho = calc_rho_q(x, q)
    expected = N * np.array([np.exp(1j * x0[0]),
                             np.exp(1j * (-2.0)),
                             np.exp(1j * (1.0 - 2.0 + 0.5))])
    assert np.allclose(rho, expected)


def test_rho_q_single_function():
    """rho_q_single returns the same result as calc_rho_q for one q-point."""
    rng = np.random.default_rng(7)
    x = rng.normal(size=(5, 3))
    q_vec = rng.normal(size=(3,))

    result_single = rho_q_single(x, q_vec)
    result_batch = calc_rho_q(x, q_vec[np.newaxis, :])[0]

    assert np.allclose(result_single, result_batch)


def test_rho_q_noncontiguous_input():
    """calc_rho_q and calc_rho_j_q must handle non-contiguous arrays correctly."""
    rng = np.random.default_rng(8)
    x_big = rng.normal(size=(8, 3))
    v_big = rng.normal(size=(8, 3))
    q_big = rng.normal(size=(6, 3))

    x_nc = x_big[::2]   # non-contiguous
    v_nc = v_big[::2]
    q_nc = q_big[::2]

    assert not x_nc.flags['C_CONTIGUOUS']

    rho_nc = calc_rho_q(x_nc, q_nc)
    rho_c = calc_rho_q(np.ascontiguousarray(x_nc), np.ascontiguousarray(q_nc))
    assert np.allclose(rho_nc, rho_c)

    rho_nc2, j_nc = calc_rho_j_q(x_nc, v_nc, q_nc)
    rho_c2, j_c = calc_rho_j_q(
        np.ascontiguousarray(x_nc),
        np.ascontiguousarray(v_nc),
        np.ascontiguousarray(q_nc),
    )
    assert np.allclose(rho_nc2, rho_c2)
    assert np.allclose(j_nc, j_c)


# Analytic tests for calc_rho_j_q
# rho(q) = sum_i exp(i * q . x_i)
# j(q)   = sum_i v_i * exp(i * q . x_i)

def test_calc_rho_j_q_rho_matches_calc_rho_q():
    """The rho returned by calc_rho_j_q must equal calc_rho_q."""
    rng = np.random.default_rng(4)
    x = rng.normal(size=(8, 3))
    v = rng.normal(size=(8, 3))
    q = rng.normal(size=(6, 3))
    rho_from_rho_j_q, _ = calc_rho_j_q(x, v, q)
    rho_from_rho_q = calc_rho_q(x, q)
    assert np.allclose(rho_from_rho_j_q, rho_from_rho_q)


def test_calc_rho_j_q_current_at_q_zero_equals_total_momentum():
    """j(q=0) = sum_i v_i, since exp(0) = 1 for each particle."""
    rng = np.random.default_rng(5)
    x = rng.normal(size=(10, 3))
    v = rng.normal(size=(10, 3))
    q = np.zeros((1, 3))
    _, j_q = calc_rho_j_q(x, v, q)
    assert np.allclose(j_q[0], v.sum(axis=0))


def test_calc_rho_j_q_zero_velocities_gives_zero_current():
    """All v=0 → j(q) = 0 for all q."""
    rng = np.random.default_rng(6)
    x = rng.normal(size=(5, 3))
    v = np.zeros((5, 3))
    q = rng.normal(size=(4, 3))
    _, j_q = calc_rho_j_q(x, v, q)
    assert np.allclose(j_q, 0.0)


def test_calc_rho_j_q_single_particle_at_origin():
    """Single particle at origin: j(q) = v for all q, since exp(0)=1."""
    v0 = np.array([[1.0, -2.0, 3.0]])
    x = np.zeros((1, 3))
    q = np.array([[1.0, 0.0, 0.0],
                  [0.0, 1.0, 0.0],
                  [1.5, -0.5, 2.0]])
    _, j_q = calc_rho_j_q(x, v0, q)
    expected = np.tile(v0, (3, 1)).astype(complex)
    assert np.allclose(j_q, expected)


def test_calc_rho_j_q_single_particle_at_known_position():
    """Single particle at x with velocity v: j(q) = v * exp(i * q . x)."""
    x0 = np.array([[2.0, 0.0, 0.0]])
    v0 = np.array([[0.0, 1.0, 0.0]])
    q = np.array([[1.0, 0.0, 0.0],
                  [0.5, 0.0, 0.0]])
    _, j_q = calc_rho_j_q(x0, v0, q)
    expected = np.array([v0[0] * np.exp(1j * 2.0),
                         v0[0] * np.exp(1j * 1.0)])
    assert np.allclose(j_q, expected)


def test_calc_rho_j_q_current_superposition():
    """j of combined set equals sum of individual j values."""
    rng = np.random.default_rng(7)
    x1 = rng.normal(size=(3, 3))
    v1 = rng.normal(size=(3, 3))
    x2 = rng.normal(size=(4, 3))
    v2 = rng.normal(size=(4, 3))
    x_combined = np.vstack([x1, x2])
    v_combined = np.vstack([v1, v2])
    q = rng.normal(size=(5, 3))

    _, j1 = calc_rho_j_q(x1, v1, q)
    _, j2 = calc_rho_j_q(x2, v2, q)
    _, j_combined = calc_rho_j_q(x_combined, v_combined, q)
    assert np.allclose(j_combined, j1 + j2)
