from __future__ import annotations

import re
import subprocess
import warnings
from shutil import which
from types import SimpleNamespace

import numpy as np
import pytest
from pytest import approx

from pymatgen.command_line import enumlib_caller
from pymatgen.command_line.enumlib_caller import EnumError, EnumlibAdaptor
from pymatgen.core import Element, Structure
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.transformations.site_transformations import RemoveSitesTransformation
from pymatgen.transformations.standard_transformations import SubstitutionTransformation
from pymatgen.util.testing import TEST_FILES_DIR, MatSciTest

ENUM_CMD = which("enum.x") or which("multienum.x")
MAKESTR_CMD = which("makestr.x") or which("makeStr.x") or which("makeStr.py")

ENUMLIB_TEST_FILES_DIR: str = f"{TEST_FILES_DIR}/command_line/enumlib"


class TestEnumlibJlWarning:
    """Tests for the advisory Enumlib.jl recommendation.

    These are independent of any installed enum.x: the ``--version`` probe is
    monkeypatched to simulate the Fortran and Julia engines.
    """

    @staticmethod
    def _patch_version_output(monkeypatch, stdout="", stderr=""):
        """Fake ``subprocess.run`` so the version probe returns fixed output."""

        def fake_run(*_args, **_kwargs):
            return SimpleNamespace(stdout=stdout, stderr=stderr, returncode=0)

        monkeypatch.setattr(enumlib_caller.subprocess, "run", fake_run)
        # Clear the per-process probe cache so each simulation is honored.
        enumlib_caller._is_enumlib_jl.cache_clear()

    def test_warns_for_fortran_enum(self, monkeypatch):
        # Legacy Fortran enum.x: no "Enumlib.jl" token -> warning fires.
        self._patch_version_output(monkeypatch, stderr="Unknown option --version")
        with pytest.warns(UserWarning, match="Enumlib.jl"):
            enumlib_caller._warn_if_not_enumlib_jl("/fake/fortran/enum.x")

    def test_no_warn_for_julia_enum(self, monkeypatch):
        # Julia Enumlib.jl engine: token present -> no warning.
        self._patch_version_output(monkeypatch, stdout="enum.x (Enumlib.jl) 0.3.2")
        with warnings.catch_warnings():
            warnings.simplefilter("error")  # any UserWarning would fail the test
            enumlib_caller._warn_if_not_enumlib_jl("/fake/julia/enum.x")
        assert enumlib_caller._is_enumlib_jl("/fake/julia/enum.x") is True

    def test_probe_never_raises(self, monkeypatch):
        # A crashing probe must be swallowed and reported as "not Enumlib.jl".
        def boom(*_args, **_kwargs):
            raise subprocess.TimeoutExpired(cmd="enum.x", timeout=10)

        monkeypatch.setattr(enumlib_caller.subprocess, "run", boom)
        enumlib_caller._is_enumlib_jl.cache_clear()
        assert enumlib_caller._is_enumlib_jl("/fake/hanging/enum.x") is False


@pytest.mark.skipif(not (ENUM_CMD and MAKESTR_CMD), reason="enumlib not present.")
class TestEnumlibAdaptor(MatSciTest):
    def test_init(self):
        struct = self.get_structure("LiFePO4")
        sub_trans = SubstitutionTransformation({"Li": {"Li": 0.5}})
        adaptor = EnumlibAdaptor(sub_trans.apply_transformation(struct), 1, 2)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 86
        for struct_trafo in structures:
            assert struct_trafo.composition.get_atomic_fraction(Element("Li")) == approx(0.5 / 6.5)
        adaptor = EnumlibAdaptor(sub_trans.apply_transformation(struct), 1, 2, refine_structure=True)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 52

        sub_trans = SubstitutionTransformation({"Li": {"Li": 0.25}})
        adaptor = EnumlibAdaptor(sub_trans.apply_transformation(struct), refine_structure=True)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 1
        for struct_trafo in structures:
            assert struct_trafo.composition.get_atomic_fraction(Element("Li")) == approx(0.25 / 6.25)

        # Make sure it works for completely disordered structures.
        struct = Structure(np.eye(3) * 10, [{"Fe": 0.5}], [[0, 0, 0]])
        adaptor = EnumlibAdaptor(struct, 1, 2)
        adaptor.run()
        assert len(adaptor.structures) == 3

        # Make sure it works properly when symmetry is broken by ordered sites.
        struct = self.get_structure("LiFePO4")
        sub_trans = SubstitutionTransformation({"Li": {"Li": 0.25}})
        struct_trafo = sub_trans.apply_transformation(struct)
        # REmove some ordered sites to break symmetry.
        remove_trans = RemoveSitesTransformation([4, 7])
        struct_trafo = remove_trans.apply_transformation(struct_trafo)
        adaptor = EnumlibAdaptor(struct_trafo, enum_precision_parameter=0.01)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 4

        struct = Structure(
            np.eye(3) * 3,
            [{"Si": 0.5}] * 2,
            [[0, 0, 0], [0.5, 0.5, 0.5]],
        )
        adaptor = EnumlibAdaptor(struct, 1, 3, enum_precision_parameter=0.01)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 10

        struct = Structure.from_file(f"{ENUMLIB_TEST_FILES_DIR}/EnumerateTest.json.gz")
        adaptor = EnumlibAdaptor(struct)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 2

    def test_rounding_errors(self):
        # It used to be that a rounding issue would result in this structure
        # showing that Cu3Te2 satisfies an ordering of this structure.
        # This has been fixed by multiplying the base by 100.
        struct = Structure.from_file(f"{TEST_FILES_DIR}/cif/Cu7Te5.cif")
        adaptor = EnumlibAdaptor(struct, 1, 2)
        with pytest.raises(EnumError, match="Unable to enumerate structure"):
            adaptor.run()
        adaptor = EnumlibAdaptor(struct, 1, 5)
        adaptor.run()
        assert len(adaptor.structures) == 197

    def test_partial_disorder(self):
        struct = Structure.from_file(filename=f"{TEST_FILES_DIR}/cif/garnet.cif")
        spga = SpacegroupAnalyzer(struct, 0.1)
        prim = spga.find_primitive()
        struct = prim.copy()
        struct["Al3+"] = {"Al3+": 0.5, "Ga3+": 0.5}
        adaptor = EnumlibAdaptor(struct, enum_precision_parameter=0.01)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 7
        for struct in structures:
            assert struct.formula == "Ca12 Al4 Ga4 Si12 O48"
        struct = prim.copy()
        struct["Ca2+"] = {"Ca2+": 1 / 3, "Mg2+": 2 / 3}
        adaptor = EnumlibAdaptor(struct, enum_precision_parameter=0.01)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 20
        for struct in structures:
            assert struct.formula == "Ca4 Mg8 Al8 Si12 O48"

        struct = prim.copy()
        struct["Si4+"] = {"Si4+": 1 / 3, "Ge4+": 2 / 3}
        adaptor = EnumlibAdaptor(struct, enum_precision_parameter=0.01)
        adaptor.run()
        structures = adaptor.structures
        assert len(structures) == 18
        for struct in structures:
            assert struct.formula == "Ca12 Al8 Si4 Ge8 O48"

    def test_timeout(self):
        struct = Structure.from_file(f"{ENUMLIB_TEST_FILES_DIR}/test_timeout.json.gz")

        adaptor = EnumlibAdaptor(struct, max_cell_size=10, timeout=0.05)  # timeout in minute

        with pytest.raises(TimeoutError, match=re.escape("Enumeration took more than timeout 0.05 minutes")):
            adaptor.run()
