"""Tests for ceraviz.structure module."""

import numpy as np

from ceraviz.structure import parse_oxidation_states, make_supercell
from ceraviz.types import CrystalTD, LatticeTD


def _make_simple_structure() -> CrystalTD:
    """Create a simple cubic structure for testing."""
    lattice: LatticeTD = {
        "abc": (4.0, 4.0, 4.0),
        "angles": (90.0, 90.0, 90.0),
        "vectors": np.eye(3) * 4.0,
        "volume": 64.0,
        "pbc": (True, True, True),
    }
    atoms = [
        {"element": "Na", "frac": np.array([0.0, 0.0, 0.0]), "cart": np.array([0.0, 0.0, 0.0])},
        {"element": "Cl", "frac": np.array([0.5, 0.5, 0.5]), "cart": np.array([2.0, 2.0, 2.0])},
    ]
    return {
        "lattice": lattice,
        "atoms": atoms,
        "nsites": 2,
        "elements": ["Na", "Cl"],
    }


class TestParseOxidationStates:
    def test_positive_oxidation(self):
        result = parse_oxidation_states(["Fe3+"])
        assert result == {"Fe": 3}

    def test_negative_oxidation(self):
        result = parse_oxidation_states(["O2-"])
        assert result == {"O": -2}

    def test_multiple(self):
        result = parse_oxidation_states(["Fe3+", "O2-"])
        assert result == {"Fe": 3, "O": -2}

    def test_single_charge(self):
        result = parse_oxidation_states(["Na+"])
        assert result == {"Na": 1}

    def test_empty(self):
        assert parse_oxidation_states([]) == {}
        assert parse_oxidation_states(None) == {}

    def test_with_spaces(self):
        result = parse_oxidation_states(["Fe 3+"])
        assert result == {"Fe": 3}


class TestMakeSupercell:
    def test_identity(self):
        struct = _make_simple_structure()
        result = make_supercell(struct, 1, 1, 1)
        assert result["nsites"] == 2
        assert len(result["atoms"]) == 2

    def test_2x2x2(self):
        struct = _make_simple_structure()
        result = make_supercell(struct, 2, 2, 2)
        assert result["nsites"] == 2 * 8  # 2 atoms × 2³ cells
        assert len(result["atoms"]) == 16

    def test_lattice_scaling(self):
        struct = _make_simple_structure()
        result = make_supercell(struct, 3, 2, 1)
        assert result["lattice"]["abc"] == (12.0, 8.0, 4.0)
        assert result["lattice"]["volume"] == 64.0 * 6

    def test_preserves_elements(self):
        struct = _make_simple_structure()
        result = make_supercell(struct, 2, 2, 2)
        assert result["elements"] == ["Na", "Cl"]

    def test_preserves_angles(self):
        struct = _make_simple_structure()
        result = make_supercell(struct, 2, 2, 2)
        assert result["lattice"]["angles"] == (90.0, 90.0, 90.0)
