import numpy as np
import pytest
from dynasor.qpoints import get_spherical_qpoints, get_supercell_qpoints_along_path
from ase.build import bulk


def test_get_spherical_qpoints():

    cell = np.diag([2.2, 3.76, 4.01])
    q_max = 20

    # without pruning
    q_points = get_spherical_qpoints(cell, q_max)
    assert q_points.shape[0] == 4465
    assert q_points.shape[1] == 3
    assert np.max(np.linalg.norm(q_points, axis=1)) <= q_max

    # with pruning
    max_points = 250
    q_points = get_spherical_qpoints(cell, q_max, max_points=max_points)

    assert q_points.shape[0] < max_points + 100
    assert q_points.shape[1] == 3
    assert np.max(np.linalg.norm(q_points, axis=1)) <= q_max


def test_get_spherical_qpoints_with_seed():
    cell = np.diag([2.2, 3.76, 4.01])
    q_max = 30
    max_points = 1500
    q_points1 = get_spherical_qpoints(cell, q_max, max_points=max_points, seed=42)
    q_points2 = get_spherical_qpoints(cell, q_max, max_points=max_points, seed=43)
    q_points3 = get_spherical_qpoints(cell, q_max, max_points=max_points, seed=42)
    assert np.allclose(q_points1, q_points3)
    assert q_points1.shape != q_points2.shape


def test_get_spherical_qpoints_left_handed_cell():
    right_handed_cell = np.diag([5.0, 5.0, 5.0])
    left_handed_cell = np.diag([-5.0, 5.0, 5.0])

    kwargs = dict(q_max=5.0, max_points=10, seed=42)
    right_handed_qpoints = get_spherical_qpoints(right_handed_cell, **kwargs)
    left_handed_qpoints = get_spherical_qpoints(left_handed_cell, **kwargs)

    assert len(left_handed_qpoints) == len(right_handed_qpoints)


def test_get_spherical_qpoints_values():
    cell = np.diag([2.2, 3.76, 4.01])
    q_max = 10
    max_points = 20
    target = [[-5.71198664,  0,          -3.13375826],
              [-2.85599332, -6.68423969,  4.70063739],
              [-2.85599332, -3.34211984, -1.56687913],
              [-2.85599332,  0,          -4.70063739],
              [-2.85599332,  0,           0],
              [-2.85599332,  5.01317977, -3.13375826],
              [-2.85599332,  5.01317977,  1.56687913],
              [0,           -6.68423969, -1.56687913],
              [0,           -5.01317977, -6.26751652],
              [0,            0,          -1.56687913],
              [0,            0,           0],
              [0,            0,           6.26751652],
              [0,            1.67105992,  0],
              [0,            1.67105992,  1.56687913],
              [0,            3.34211984, -1.56687913],
              [0,            3.34211984,  0],
              [0,            6.68423969, -3.13375826],
              [2.85599332,  -3.34211984,  1.56687913],
              [2.85599332,   0,          -1.56687913],
              [2.85599332,   0,           9.40127477],
              [2.85599332,   8.35529961,  3.13375826],
              [5.71198664,  -5.01317977,  3.13375826],
              [8.56797996,   1.67105992,  0]]
    actual = get_spherical_qpoints(cell, q_max, max_points=max_points, seed=42)
    assert np.allclose(actual, target)


def test_get_spherical_qpoints_invalid_q_min_greater_than_q_max():
    cell = np.diag([5.0, 5.0, 5.0])
    with pytest.raises(ValueError, match='q_min must be smaller than q_max'):
        get_spherical_qpoints(cell, q_max=5.0, q_min=10.0)


def test_get_spherical_qpoints_invalid_q_min_equal_to_q_max():
    cell = np.diag([5.0, 5.0, 5.0])
    with pytest.raises(ValueError, match='q_min must be smaller than q_max'):
        get_spherical_qpoints(cell, q_max=5.0, q_min=5.0)


def test_get_spherical_qpoints_invalid_negative_q_min():
    cell = np.diag([5.0, 5.0, 5.0])
    with pytest.raises(ValueError, match='q_min cannot be negative'):
        get_spherical_qpoints(cell, q_max=5.0, q_min=-1.0)


def test_get_spherical_qpoints_empty_result():
    """q_min > 0 with q_max below the smallest non-zero reciprocal vector returns empty array."""
    # Large cell: smallest non-zero |q| = 2π/100 ≈ 0.063 rad/Å
    cell = np.diag([100.0, 100.0, 100.0])
    q_points = get_spherical_qpoints(cell, q_max=0.05, q_min=0.01)
    assert q_points.shape == (0, 3)


def test_get_spherical_qpoints_max_points_no_valid_prune_distance():
    """When the cubic has no real solution, pruning is skipped and all points are returned."""
    # This cell/q_max combination causes D_t < 0 in _get_prune_distance,
    # meaning no valid q_prune exists; the function should return all points without error.
    cell = np.diag([4, 3, 18])
    q_max = 3
    q_points_unpruned = get_spherical_qpoints(cell, q_max)
    q_points_pruned = get_spherical_qpoints(cell, q_max, max_points=100)
    # Pruning is silently skipped, so all points are returned
    assert len(q_points_pruned) == len(q_points_unpruned)
    assert np.all(np.linalg.norm(q_points_pruned, axis=1) <= q_max)


def test_get_spherical_qpoints_chunk_with_no_matches():
    """chunk_size=1 with q_min > 0 forces chunks containing only the excluded origin."""
    cell = np.diag([5.0, 5.0, 5.0])
    q_points = get_spherical_qpoints(cell, q_max=2.0, q_min=0.5, chunk_size=1)
    norms = np.linalg.norm(q_points, axis=1)
    assert np.all(norms >= 0.5)
    assert np.all(norms <= 2.0)


def test_get_supercell_qpts_along_path_cubic():
    prim = bulk('Al', 'fcc', a=4.05)
    supercell = bulk('Al', 'fcc', a=4.05, cubic=True).repeat(4)
    coordinates = dict(
        X=[0.5, 0.5, 0],
        G=[0, 0, 0],
        L=[0.5, 0.5, 0.5],
        W=[0.5, 0.25, 0.75],
    )

    # connected path
    path = [('X', 'G'), ('G', 'L')]
    qpoints = get_supercell_qpoints_along_path(path, coordinates, prim.cell, supercell.cell)
    assert len(qpoints) == len(path)

    qpoints_target = [
        np.array([[0,          0,          1.55140378],
                  [0,          0,          1.16355283],
                  [0,          0,          0.77570189],
                  [0,          0,          0.38785094],
                  [0,          0,          0]]),
        np.array([[0,          0,          0],
                  [0.38785094, 0.38785094, 0.38785094],
                  [0.77570189, 0.77570189, 0.77570189]])
        ]

    assert len(qpoints) == len(qpoints_target)
    for q_segment1, q_segment2 in zip(qpoints, qpoints_target):
        assert np.allclose(q_segment1, q_segment2)

    # disconnected path
    path = [('X', 'G'), ('G', 'L'), ('W', 'X')]
    qpoints = get_supercell_qpoints_along_path(path, coordinates, prim.cell, supercell.cell)
    assert len(qpoints) == len(path)

    qpoints_target = [
        np.array([[0,          0,          1.55140378],
                  [0,          0,          1.16355283],
                  [0,          0,          0.77570189],
                  [0,          0,          0.38785094],
                  [0,          0,          0]]),
        np.array([[0,          0,          0],
                  [0.38785094, 0.38785094, 0.38785094],
                  [0.77570189, 0.77570189, 0.77570189]]),
        np.array([[0.77570189, 1.55140378, 0],
                  [0.38785094, 0.77570189, 0.77570189],
                  [0, 0, 1.55140378]])
        ]

    assert len(qpoints) == len(qpoints_target)
    for q_segment1, q_segment2 in zip(qpoints, qpoints_target):
        assert np.allclose(q_segment1, q_segment2)


def test_get_supercell_qpts_along_path_hex():
    prim = bulk('Al', 'hcp', a=3.05, c=5.0)
    supercell = prim.repeat(4)
    coordinates = dict(
        X=[0.5, 0.5, 0],
        G=[0, 0, 0],
        L=[0.5, 0.5, 0.5],
        W=[0.5, 0.25, 0.75],
    )

    # connected path
    path = [('X', 'G'), ('G', 'L')]
    qpoints = get_supercell_qpoints_along_path(path, coordinates, prim.cell, supercell.cell)
    assert len(qpoints) == len(path)

    qpoints_target = [
        np.array([[1.03003038, 1.78406495, 0],
                  [0.51501519, 0.89203247, 0],
                  [0, 0, 0]]),
        np.array([[0, 0, 0],
                  [0.51501519, 0.89203247, 0.31415927],
                  [1.03003038, 1.78406495, 0.62831853]]),
        ]
    assert len(qpoints) == len(qpoints_target)
    for q_segment1, q_segment2 in zip(qpoints, qpoints_target):
        assert np.allclose(q_segment1, q_segment2)


def test_get_supercell_qpoints_along_path():
    prim = bulk('Al', 'fcc', a=4.0)
    supercell = bulk('Al', 'fcc', a=4.0, cubic=True).repeat(1)
    path = [('X', 'G'), ('A', 'B'), ('G', 'L')]
    coordinates = dict(X=[0.5, 0.5, 0], G=[0, 0, 0], L=[0.5, 0.5, 0.5], A=[1/6, 1/6, 1/6],
                       B=[2/6, 2/6, 2/6])
    with pytest.warns(UserWarning, match='No q-points along path!'):
        qpoints = get_supercell_qpoints_along_path(path, coordinates, prim.cell, supercell.cell)
    assert len(qpoints) == 3
    assert qpoints[1].shape == (0, 3)


def test_get_supercell_qpts_along_path_errors():
    prim = bulk('Al', 'fcc', a=4.0)
    supercell = bulk('Al', 'fcc', a=4.05, cubic=True).repeat(4)
    path = [('X', 'G'), ('G', 'L')]
    coordinates = dict(X=[0.5, 0.5, 0], G=[0, 0, 0], L=[0.5, 0.5, 0.5])
    with pytest.raises(ValueError, match='Please check that the supercell metric'):
        qpoints, labels = get_supercell_qpoints_along_path(
            path, coordinates, prim.cell, supercell.cell)

    coordinates = dict(X=[0.5, 0.5, 0], G=[0, 0, 0])
    with pytest.raises(ValueError, match='Please check that the supercell metric'):
        qpoints, labels = get_supercell_qpoints_along_path(
            path, coordinates, prim.cell, supercell.cell)
