import numpy as np
import pytest
from ase import Atoms
from ase.build import bulk

from dynasor.tools.structures import (
    get_displacements,
    get_displacements_from_u,
    find_permutation,
    get_offset_index,
)


@pytest.fixture
def fcc_al():
    return bulk('Al', 'fcc', a=4.05, cubic=True)


# --- get_displacements ---

def test_get_displacements_basic(fcc_al):
    """Displaced atoms return the displacement vectors."""
    ideal = fcc_al.copy()
    displaced = fcc_al.copy()
    displaced.positions[0] += [0.1, 0.0, 0.0]
    u = get_displacements(displaced, ideal)
    assert np.isclose(u[0, 0], 0.1)
    assert np.allclose(u[1:], 0.0)


def test_get_displacements_no_mic():
    ideal = Atoms('Al', positions=[[0, 0, 0]], cell=10 * np.eye(3), pbc=True)
    displaced = ideal.copy()
    displaced.positions[0, 0] = 8.0

    u = get_displacements(displaced, ideal, check_mic=False)
    assert np.isclose(u[0, 0], 8.0)


def test_get_displacements_mismatched_numbers(fcc_al):
    """Different atomic numbers raise ValueError."""
    ideal = fcc_al.copy()
    other = fcc_al.copy()
    other.numbers[0] = 29  # Cu instead of Al
    with pytest.raises(ValueError, match='Atomic numbers do not match'):
        get_displacements(other, ideal)


def test_get_displacements_mismatched_cell(fcc_al):
    """Different cell metrics raise ValueError."""
    ideal = fcc_al.copy()
    other = fcc_al.copy()
    other.cell *= 1.1
    with pytest.raises(ValueError, match='Cells do not match'):
        get_displacements(other, ideal)


# --- get_displacements_from_u ---

def test_get_displacements_from_u_no_mic():
    """check_mic=False returns u unchanged even when |u| > half the cell length."""
    cell = 10.0 * np.eye(3)
    # displacement of 8 Å in a 10 Å cell: mic would fold it to -2 Å
    u = np.array([[8.0, 0.0, 0.0]])
    result = get_displacements_from_u(u, cell, check_mic=False)
    assert np.allclose(result, u)
    # confirm that mic=True would give a different (shorter) result
    result_mic = get_displacements_from_u(u.copy(), cell, check_mic=True)
    assert not np.allclose(result_mic, u)


# --- find_permutation ---

def test_find_permutation_cell_mismatch():
    """find_permutation raises ValueError when cells differ."""
    a = bulk('Al', 'fcc', a=4.05, cubic=True)
    b = bulk('Al', 'fcc', a=4.10, cubic=True)
    with pytest.raises(ValueError, match='Cells do not match'):
        find_permutation(a, b)


def test_find_permutation_duplicate_raises():
    """Duplicate assignments in the permutation raise Exception."""
    # Two atoms at the same position → both map to the same reference atom
    cell = 10.0 * np.eye(3)
    atoms = Atoms('AlAl', positions=[[0, 0, 0], [0, 0, 0]], cell=cell, pbc=True)
    atoms_ref = Atoms('AlAl', positions=[[0, 0, 0], [5, 0, 0]], cell=cell, pbc=True)
    with pytest.raises(Exception, match='Duplicates in permutation'):
        find_permutation(atoms, atoms_ref)


def test_find_permutation_occupation_mismatch_raises():
    """Different species at matched lattice sites raise Exception."""
    cell = 10.0 * np.eye(3)
    atoms = Atoms('CuAl', positions=[[0, 0, 0], [5, 0, 0]], cell=cell, pbc=True)
    atoms_ref = Atoms('AlCu', positions=[[0, 0, 0], [5, 0, 0]], cell=cell, pbc=True)
    with pytest.raises(Exception, match='Matching lattice sites have different occupation'):
        find_permutation(atoms, atoms_ref)


# --- get_offset_index ---

def test_get_offset_index_non_atoms_primitive_raises(fcc_al):
    """Non-Atoms primitive raises ValueError."""
    with pytest.raises(ValueError):
        get_offset_index('not_atoms', fcc_al)


def test_get_offset_index_non_atoms_supercell_raises(fcc_al):
    """Non-Atoms supercell raises ValueError."""
    with pytest.raises(ValueError):
        get_offset_index(fcc_al, 'not_atoms')
