import numpy as np
import taco_format


def normalize_read(result):
    """Ensure taco_format.read() always returns a list."""
    return result if isinstance(result, list) else [result]


def test_stress_and_extra(tmp_path):
    from ase import Atoms
    from ase.calculators.singlepoint import SinglePointCalculator

    # Create a simple atoms object
    atoms = Atoms("H2", positions=[[0, 0, 0], [1, 0, 0]])
    atoms.set_velocities([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
    # Set stress and extra arrays
    stress = np.random.rand(6)
    extra1 = np.array([42.0, 43.0])
    extra2 = np.array([[1.0, 2.0], [3.0, 4.0]])
    extra3 = np.arange(12).reshape(2, 2, 3)
    # Attach as calculator results (simulate as if these are available)
    atoms.arrays["extra1"] = extra1
    atoms.arrays["extra2"] = extra2
    atoms.arrays["extra3"] = extra3
    calc = SinglePointCalculator(
        atoms, energy=1.23, forces=np.ones((2, 3)), stress=stress
    )
    atoms.calc = calc
    # Write and read
    filename = tmp_path / "test_stress_extra.taco"
    taco_format.write(str(filename), atoms)
    atoms_list = normalize_read(taco_format.read(str(filename)))
    atoms2 = atoms_list[0]
    # Check stress
    np.testing.assert_allclose(atoms2.get_stress(), stress)
    # Check extra arrays
    np.testing.assert_allclose(atoms2.arrays["extra1"], extra1)
    np.testing.assert_allclose(atoms2.arrays["extra2"], extra2)
    np.testing.assert_allclose(atoms2.arrays["extra3"], extra3)
    # Check energy and forces
    assert abs(atoms2.get_potential_energy() - 1.23) < 1e-6
    np.testing.assert_allclose(atoms2.get_forces(), np.ones((2, 3)))
