from __future__ import annotations

import math
import re
import warnings
from pathlib import Path

import numpy as np
import pytest
from pytest import approx

from pymatgen.core.structure import Molecule, Structure
from pymatgen.io.cp2k.sets import SETTINGS, Cp2kValidationError, DftSet, GaussianTypeOrbitalBasisSet, GthPotential
from pymatgen.util.testing import TEST_FILES_DIR, MatSciTest

CP2K_TEST_DIR = f"{TEST_FILES_DIR}/io/cp2k"
# Local test-files dir (independent of the installed package's TEST_FILES_DIR).
LOCAL_CP2K_TEST_DIR = Path(__file__).parent.parent.parent.parent / "test-files" / "io" / "cp2k"

Si_structure = Structure(
    lattice=[[0, 2.734364, 2.734364], [2.734364, 0, 2.734364], [2.734364, 2.734364, 0]],
    species=["Si", "Si"],
    coords=[[0, 0, 0], [0.25, 0.25, 0.25]],
)
molecule = Molecule(species=["Si"], coords=[[0, 0, 0]])

# Basis set / potential with objects
Si_gto_basis = """
Si SZV-MOLOPT-GTH SZV-MOLOPT-GTH-q4
1
2 0 1 6 1 1
    2.693604434572  0.015333179500 -0.005800105400
    1.359613855428 -0.283798205000 -0.059172026000
    0.513245176029 -0.228939692700  0.121487149900
    0.326563011394  0.728834000900  0.423382421100
    0.139986977410  0.446205299300  0.474592116300
    0.068212286977  0.122025292800  0.250129397700
"""
Si_gth_pot = """Si GTH-BLYP-q4 GTH-BLYP
2    2
0.44000000    1    -6.25958674
2
0.44465247    2     8.31460936    -2.33277947
                                    3.01160535
0.50279207    1     2.33241791"""

BASIS_AND_POTENTIAL: dict[str, dict[str, GaussianTypeOrbitalBasisSet | GthPotential]] = {
    "Si": {"basis": GaussianTypeOrbitalBasisSet.from_str(Si_gto_basis), "potential": GthPotential.from_str(Si_gth_pot)},
    "Fe": {
        "basis": GaussianTypeOrbitalBasisSet.from_str(Si_gto_basis.replace("Si", "Fe")),
        "potential": GthPotential.from_str(Si_gth_pot.replace("Si", "Fe")),
    },
}


class TestDftSet(MatSciTest):
    def test_dft_set(self) -> None:
        """Test various DFT set configurations."""
        SETTINGS["PMG_CP2K_DATA_DIR"] = CP2K_TEST_DIR

        # Basis sets / potentials searching
        basis_and_potential: dict[str, str | dict[str, str | None] | None] = {
            "basis_type": "SZV",
            "potential_type": "Pseudopotential",
            "functional": None,
        }
        dft_set = DftSet(Si_structure, basis_and_potential=basis_and_potential, xc_functionals="PBE")

        # Basis sets / potentials by name
        basis_and_potential = {"Si": {"basis": "SZV-GTH-q4", "potential": "GTH-PBE-q4"}}
        dft_set = DftSet(Si_structure, basis_and_potential=basis_and_potential, xc_functionals="PBE")

        # Basis sets / potentials by name with ADMM
        basis_and_potential = {
            "Si": {
                "basis": "SZV-GTH-q4",
                "potential": "GTH-PBE-q4",
                "aux_basis": "cFIT3",
            }
        }
        dft_set = DftSet(Si_structure, basis_and_potential=basis_and_potential, xc_functionals="PBE")
        basis_sets = dft_set["force_eval"]["subsys"]["Si_1"].get("basis_set")
        assert any("AUX_FIT" in b.values for b in basis_sets)
        assert any("cFIT3" in b.values for b in basis_sets)

        # Basis sets / potentials by hash value
        basis_and_potential = {
            "Si": {
                "basis": "30767c18f6e7e46c1b56c1d34ff6007d",
                "potential": "21e2f468a18404ff6119fe801da81e43",
            }
        }
        dft_set = DftSet(Si_structure, basis_and_potential=basis_and_potential, xc_functionals="PBE")

        set_kwargs = dict.fromkeys(("print_pdos", "print_dos", "print_v_hartree", "print_e_density"), False)
        dft_set = DftSet(
            Si_structure,
            basis_and_potential=BASIS_AND_POTENTIAL,
            xc_functionals="PBE",
            **set_kwargs,
        )
        assert dft_set.cutoff == approx(150)

        # Test that printing will activate sections
        assert not dft_set.check("motion")
        dft_set.activate_motion()
        assert dft_set.check("motion")
        assert not dft_set.check("force_eval/dft/print/pdos")
        dft_set.print_pdos()
        assert dft_set.check("force_eval/dft/print/pdos")
        assert not dft_set.check("force_eval/dft/print/v_hartree_cube")
        dft_set.print_v_hartree()
        assert dft_set.check("force_eval/dft/print/v_hartree_cube")

        # Test property activators
        assert not dft_set.check("force_eval/properties")
        dft_set.activate_nmr()
        dft_set.activate_epr()
        dft_set.activate_hyperfine()
        dft_set.activate_polar()
        dft_set.activate_tddfpt()
        assert dft_set.check("force_eval/properties/linres/localize")
        assert dft_set.check("force_eval/properties/linres/nmr/print/chi_tensor")
        assert dft_set.check("force_eval/properties/linres/epr/print/g_tensor")
        assert dft_set.check("force_eval/properties/tddfpt")
        assert dft_set.check("force_eval/dft/print/hyperfine_coupling_tensor")

        # For at least up to v2022.1, DOS doesn't work without kpoints
        assert not dft_set.check("force_eval/dft/print/dos")
        dft_set.print_dos()
        assert not dft_set.check("force_eval/dft/print/dos")

        assert not dft_set.check("force_eval/dft/xc/hf")
        dft_set.activate_hybrid()
        assert dft_set.check("force_eval/dft/xc/hf")
        assert dft_set.check("force_eval/dft/auxiliary_density_matrix_method")

        # Validator will trip for kpoints + hfx
        dft_set |= {"force_eval": {"dft": {"kpoints": {}}}}
        with pytest.raises(
            Cp2kValidationError,
            match=re.escape("CP2K v2022.1: Does not support hartree fock with kpoints"),
        ):
            dft_set.validate()

        # Test non-periodic settings
        dft_set = DftSet(molecule, basis_and_potential=basis_and_potential, xc_functionals="PBE")
        assert dft_set.check("force_eval/dft/poisson")
        assert dft_set["force_eval"]["dft"]["poisson"].get("periodic").values[0].upper() == "NONE"

    def test_kind_magnetization(self) -> None:
        """Test different ways of setting kind (i.e. element-specific) initial magnetization."""
        fe_structure = Structure(
            lattice=np.eye(3) * 3,
            species=["Fe", "Fe"],
            coords=[[0, 0, 0], [0.5, 0.5, 0.5]],
        )

        # Test 1: Setting via site properties
        dft_set = DftSet(
            fe_structure.copy().add_site_property("magmom", [2.0, -2.0]),
            basis_and_potential=BASIS_AND_POTENTIAL,
            xc_functionals="PBE",
        )
        fe1_kind = dft_set["force_eval"]["subsys"]["Fe_1"]
        fe2_kind = dft_set["force_eval"]["subsys"]["Fe_2"]
        assert {fe1_kind["magnetization"].values[0], fe2_kind["magnetization"].values[0]} == {2.0, -2.0}

        # Test 2: Setting via element defaults
        element_defaults = {"Fe": {"magnetization": 4.0}}
        dft_set = DftSet(
            fe_structure,
            basis_and_potential=BASIS_AND_POTENTIAL,
            xc_functionals="PBE",
            element_defaults=element_defaults,
        )
        fe_kind = dft_set["force_eval"]["subsys"]["Fe_1"]
        assert math.isclose(fe_kind["magnetization"].values[0], 4.0)

        # Test 3: Site properties take precedence over element defaults
        fe_structure.add_site_property("magmom", [2.0, -2.0])
        dft_set = DftSet(
            fe_structure,
            basis_and_potential=BASIS_AND_POTENTIAL,
            xc_functionals="PBE",
            element_defaults=element_defaults,
        )
        fe1_kind = dft_set["force_eval"]["subsys"]["Fe_1"]
        fe2_kind = dft_set["force_eval"]["subsys"]["Fe_2"]
        assert {fe1_kind["magnetization"].values[0], fe2_kind["magnetization"].values[0]} == {2.0, -2.0}

    def test_cell_parameters(self) -> None:
        """Test that cell parameters are properly set when provided."""
        # Test with cell parameters
        cell_params = {
            "SYMMETRY": "CUBIC",
            "MULTIPLE_UNIT_CELL": [2, 2, 2],
        }

        dft_set = DftSet(Si_structure, basis_and_potential=BASIS_AND_POTENTIAL, xc_functionals="PBE", cell=cell_params)

        # Check that cell parameters were properly set
        subsys = dft_set["force_eval"]["subsys"]
        assert subsys.check("cell")
        cell_section = subsys["cell"]
        assert cell_section["symmetry"].values[0] == "CUBIC"
        assert cell_section["multiple_unit_cell"].values == ([2, 2, 2],)

        # Test without cell parameters (default behavior)
        dft_set = DftSet(Si_structure, basis_and_potential=BASIS_AND_POTENTIAL, xc_functionals="PBE")

        # Check that only default cell parameters exist
        subsys = dft_set["force_eval"]["subsys"]
        assert subsys.check("cell")
        cell_section = subsys["cell"]
        assert not cell_section.check("symmetry")
        assert not cell_section.check("multiple_unit_cell")

    def test_no_yaml_unsafe_deprecation_warning(self) -> None:
        """Regression test for https://github.com/materialsproject/pymatgen-core/issues/4444.

        Loading CP2K element data files must not emit a PendingDeprecationWarning
        about ruamel.yaml's deprecated 'unsafe' loader type.
        """
        SETTINGS["PMG_CP2K_DATA_DIR"] = str(LOCAL_CP2K_TEST_DIR)
        basis_and_potential: dict[str, str | dict] = {
            "basis_type": "SZV",
            "potential_type": "Pseudopotential",
            "functional": None,
        }

        with warnings.catch_warnings():
            warnings.simplefilter("error", PendingDeprecationWarning)
            # Constructing DftSet with a data dir triggers YAML loading of the
            # per-element files (e.g. test-files/io/cp2k/Si). The old code used
            # YAML(typ="unsafe") which raises a PendingDeprecationWarning from
            # ruamel.yaml. After the fix this must not raise.
            DftSet(Si_structure, basis_and_potential=basis_and_potential, xc_functionals="PBE")

    def test_yaml_data_loading_roundtrip(self) -> None:
        """CP2K element data files use MSONable-format YAML (no unsafe Python tags).

        Verify that the data can be loaded with the default (safe) ruamel.yaml
        loader and that the resulting dicts reconstruct valid objects.
        """
        from ruamel.yaml import YAML

        si_file = LOCAL_CP2K_TEST_DIR / "Si"
        yaml = YAML()
        with open(si_file, encoding="utf-8") as fh:
            data = yaml.load(fh)

        assert "basis_sets" in data
        assert "potentials" in data

        # Reconstruct a GaussianTypeOrbitalBasisSet from the YAML data
        basis_hash = next(iter(data["basis_sets"]))
        basis = GaussianTypeOrbitalBasisSet.from_dict(data["basis_sets"][basis_hash])
        assert isinstance(basis, GaussianTypeOrbitalBasisSet)
        assert basis.element is not None
        assert basis.name is not None

        # Reconstruct a GthPotential from the YAML data
        pot_hash = next(iter(data["potentials"]))
        potential = GthPotential.from_dict(data["potentials"][pot_hash])
        assert isinstance(potential, GthPotential)
        assert potential.element is not None
        assert potential.name is not None
