"""IO mixin for reading and writing VASP files.

Provides methods for:
- Writing input files (INCAR, POSCAR, KPOINTS, POTCAR)
- Reading input files
- Reading output files (OUTCAR, vasprun.xml)
- Parsing results
"""

from __future__ import annotations

import os
import re
from typing import TYPE_CHECKING, Any

import numpy as np

if TYPE_CHECKING:
    from ase import Atoms


class IOMixin:
    """Mixin providing VASP file I/O methods.

    This mixin requires the following attributes on the class:
    - directory: str - Path to calculation directory
    - parameters: dict - VASP parameters
    - atoms: Atoms - ASE Atoms object
    - results: dict - Calculation results
    - sort: list - Atom sorting indices
    - resort: list - Reverse sorting indices
    """

    # Type declarations for attributes provided by the main class
    directory: str
    parameters: dict[str, Any]
    atoms: Atoms | None
    results: dict[str, Any]
    sort: list[int]
    resort: list[int]

    # =========================================================================
    # Input File Writing
    # =========================================================================

    def write_input(
        self,
        atoms: Atoms | None = None,
        properties: list[str] | None = None,
        system_changes: list[str] | None = None,
    ) -> None:
        """Write all VASP input files.

        Args:
            atoms: ASE Atoms object (uses self.atoms if None).
            properties: Requested properties (unused, for ASE compatibility).
            system_changes: System changes (unused, for ASE compatibility).
        """
        if atoms is not None:
            self.atoms = atoms

        if self.atoms is None:
            raise ValueError(
                "No atoms specified for calculation.\n\n"
                "Atoms must be provided either:\n"
                "  1. At initialization: calc = Vasp('dir', atoms=atoms, ...)\n"
                "  2. As a property: calc.atoms = atoms\n"
                "  3. When calling write_input: calc.write_input(atoms=atoms)\n\n"
                "Example:\n"
                "  from ase.build import bulk\n"
                "  atoms = bulk('Si')\n"
                "  calc = Vasp('si-calc', atoms=atoms, xc='PBE')\n"
            )

        # Ensure directory exists
        os.makedirs(self.directory, exist_ok=True)

        # Check for NEB calculation
        if hasattr(self, "neb_images") and self.neb_images is not None:
            self._write_neb_input()
        else:
            # Write standard input files
            self.write_poscar()
            self.write_incar()
            self.write_kpoints()
            self.write_potcar()

    def _write_neb_input(self) -> None:
        """Write input files for NEB calculation."""
        from ase.io import write

        images = self.neb_images
        n_images = len(images)

        # Set IMAGES parameter (intermediate images, not including endpoints)
        self.parameters["images"] = n_images - 2

        # Write POSCAR for each image in subdirectories
        for i, img in enumerate(images):
            subdir = os.path.join(self.directory, f"{i:02d}")
            os.makedirs(subdir, exist_ok=True)

            # Sort atoms for VASP
            sorted_atoms = img[self.sort]
            poscar_path = os.path.join(subdir, "POSCAR")
            write(poscar_path, sorted_atoms, format="vasp", vasp5=True)

        # Write INCAR, KPOINTS, POTCAR in main directory
        self.write_incar()
        self.write_kpoints()
        self.write_potcar()

    def write_poscar(self, atoms: Atoms | None = None) -> None:
        """Write POSCAR file.

        Args:
            atoms: ASE Atoms object (uses self.atoms if None).
        """
        if atoms is None:
            atoms = self.atoms

        from ase.io import write

        poscar_path = os.path.join(self.directory, "POSCAR")

        # Sort atoms by species
        sorted_atoms = atoms[self.sort]
        write(poscar_path, sorted_atoms, format="vasp", vasp5=True)

    def write_incar(self) -> None:
        """Write INCAR file from parameters."""
        incar_path = os.path.join(self.directory, "INCAR")

        with open(incar_path, "w") as f:
            f.write("# INCAR generated by vasp-ase\n")

            # Get species order for dict-type parameters
            if self.atoms is not None:
                symbols = self.atoms.get_chemical_symbols()
                species_order = []
                for s in symbols:
                    if s not in species_order:
                        species_order.append(s)
            else:
                species_order = []

            for key, val in sorted(self.parameters.items()):
                # Skip special keywords
                if key in ("xc", "pp", "setups", "kpts", "gamma", "kpts_nintersections"):
                    continue

                # Format value
                if isinstance(val, bool):
                    val_str = ".TRUE." if val else ".FALSE."
                elif isinstance(val, dict):
                    # Dict mapping element symbols to values (e.g., rwigs, magmom)
                    # Convert to space-separated values in species order
                    if species_order:
                        values = [str(val.get(s, 0)) for s in species_order]
                        val_str = " ".join(values)
                    else:
                        # Fallback: just use values in dict order
                        val_str = " ".join(str(v) for v in val.values())
                elif isinstance(val, (list, tuple, np.ndarray)):
                    val_str = " ".join(str(v) for v in val)
                elif isinstance(val, str):
                    val_str = val
                else:
                    val_str = str(val)

                f.write(f"{key.upper()} = {val_str}\n")

    def write_kpoints(self) -> None:
        """Write KPOINTS file."""
        # Check if KSPACING is set (no KPOINTS needed)
        if "kspacing" in self.parameters:
            return

        kpoints_path = os.path.join(self.directory, "KPOINTS")
        kpts = self.parameters.get("kpts", (1, 1, 1))
        gamma = self.parameters.get("gamma", False)

        with open(kpoints_path, "w") as f:
            f.write("KPOINTS generated by vasp-ase\n")

            if isinstance(kpts, (list, tuple)) and len(kpts) == 3:
                # Monkhorst-Pack grid
                f.write("0\n")
                if gamma:
                    f.write("Gamma\n")
                else:
                    f.write("Monkhorst-Pack\n")
                f.write(f"{kpts[0]} {kpts[1]} {kpts[2]}\n")
                f.write("0 0 0\n")
            else:
                # Explicit k-points
                f.write(f"{len(kpts)}\n")
                f.write("Reciprocal\n")
                for kpt in kpts:
                    if len(kpt) == 3:
                        f.write(f"{kpt[0]} {kpt[1]} {kpt[2]} 1.0\n")
                    else:
                        f.write(f"{kpt[0]} {kpt[1]} {kpt[2]} {kpt[3]}\n")

    def write_potcar(self) -> None:
        """Write POTCAR file by concatenating pseudopotentials.

        Checks for pseudopotential path in these environment variables (in order):
        - VASP_PP_PATH
        - VASP_PP_BASE
        - VASP_PSP_DIR

        If none are set, writes a placeholder POTCAR for testing.
        """
        potcar_path = os.path.join(self.directory, "POTCAR")

        # Check multiple environment variables for PP path
        pp_path = None
        for env_var in ["VASP_PP_PATH", "VASP_PP_BASE", "VASP_PSP_DIR"]:
            pp_path = os.environ.get(env_var)
            if pp_path:
                break

        if not pp_path:
            # For testing without real POTCARs, create placeholder
            with open(potcar_path, "w") as f:
                symbols = self.atoms.get_chemical_symbols()
                unique = []
                for s in symbols:
                    if s not in unique:
                        unique.append(s)
                for symbol in unique:
                    f.write(f"  PAW_PBE {symbol} (placeholder for testing)\n")
                    f.write(f"   TITEL  = PAW_PBE {symbol}\n")
                    f.write("   End of Dataset\n\n")
            return

        pp = self.parameters.get("pp", "PBE")
        setups = self.parameters.get("setups", {})

        # Get unique symbols in order
        symbols = self.atoms.get_chemical_symbols()
        sorted_symbols = [symbols[i] for i in self.sort]
        unique_symbols = []
        for s in sorted_symbols:
            if s not in unique_symbols:
                unique_symbols.append(s)

        with open(potcar_path, "w") as potcar:
            for symbol in unique_symbols:
                # Check for special setup
                setup = setups.get(symbol, "")
                if setup:
                    pot_name = f"{symbol}_{setup}"
                else:
                    pot_name = symbol

                # Find POTCAR file
                pot_file = os.path.join(pp_path, f"potpaw_{pp}", pot_name, "POTCAR")
                if not os.path.exists(pot_file):
                    pot_file = os.path.join(pp_path, f"potpaw_{pp}", symbol, "POTCAR")

                if not os.path.exists(pot_file):
                    raise FileNotFoundError(
                        f"POTCAR not found for element '{symbol}'.\n"
                        f"Searched: {pot_file}\n"
                        f"PP path: {pp_path}\n"
                        f"PP type: {pp}\n\n"
                        f"Solutions:\n"
                        f"  1. Set one of these environment variables:\n"
                        f"     export VASP_PP_PATH=/path/to/vasp/potentials\n"
                        f"     export VASP_PP_BASE=/path/to/vasp/potentials\n"
                        f"     export VASP_PSP_DIR=/path/to/vasp/potentials\n"
                        f"  2. Ensure directory structure is:\n"
                        f"     $VASP_PP_PATH/potpaw_{pp}/{symbol}/POTCAR\n\n"
                        f"See: https://www.vasp.at/wiki/index.php/POTCAR"
                    )

                with open(pot_file) as pf:
                    potcar.write(pf.read())

    # =========================================================================
    # Input File Reading
    # =========================================================================

    def read_incar(self) -> dict[str, Any]:
        """Read INCAR file into dictionary.

        Returns:
            Dict of INCAR parameters.
        """
        incar_path = os.path.join(self.directory, "INCAR")
        params = {}

        if not os.path.exists(incar_path):
            return params

        with open(incar_path) as f:
            for line in f:
                line = line.strip()

                # Skip comments and empty lines
                if not line or line.startswith("#") or line.startswith("!"):
                    continue

                # Remove inline comments
                if "#" in line:
                    line = line.split("#")[0]
                if "!" in line:
                    line = line.split("!")[0]

                # Parse key = value
                if "=" in line:
                    key, val = line.split("=", 1)
                    key = key.strip().lower()
                    val = val.strip()
                    params[key] = self._parse_incar_value(val)

        return params

    def _parse_incar_value(self, val: str) -> Any:
        """Parse an INCAR value string to Python type."""
        val = val.strip()

        # Boolean
        if val.upper() in (".TRUE.", "T", ".T."):
            return True
        if val.upper() in (".FALSE.", "F", ".F."):
            return False

        # Try integer
        try:
            return int(val)
        except ValueError:
            pass

        # Try float
        try:
            return float(val)
        except ValueError:
            pass

        # Array of values
        if " " in val:
            parts = val.split()
            try:
                return [int(p) for p in parts]
            except ValueError:
                try:
                    return [float(p) for p in parts]
                except ValueError:
                    pass

        return val

    def read_kpoints(self) -> dict[str, Any]:
        """Read KPOINTS file.

        Returns:
            Dict with 'kpts', 'gamma', and other k-point info.
        """
        kpoints_path = os.path.join(self.directory, "KPOINTS")

        if not os.path.exists(kpoints_path):
            return {}

        with open(kpoints_path) as f:
            lines = f.readlines()

        result = {}
        if len(lines) < 3:
            return result

        nkpts = int(lines[1].strip())
        mode = lines[2].strip()[0].upper()

        if nkpts == 0:
            # Automatic grid
            result["gamma"] = mode == "G"
            kpts = [int(x) for x in lines[3].split()[:3]]
            result["kpts"] = tuple(kpts)
        else:
            # Explicit k-points
            kpts = []
            for line in lines[3 : 3 + nkpts]:
                parts = line.split()
                kpt = [float(x) for x in parts[:3]]
                if len(parts) > 3:
                    kpt.append(float(parts[3]))
                kpts.append(kpt)
            result["kpts"] = kpts

        return result

    # =========================================================================
    # Output File Reading
    # =========================================================================

    def read_results(self) -> None:
        """Read calculation results from output files.

        Populates self.results with energy, forces, stress, etc.
        """
        self.results = {}

        # Try vasprun.xml first (more complete)
        vasprun_path = os.path.join(self.directory, "vasprun.xml")
        if os.path.exists(vasprun_path):
            self._read_vasprun_xml()
        else:
            # Fall back to OUTCAR
            self._read_outcar()

        # Read atoms from CONTCAR if available
        contcar = os.path.join(self.directory, "CONTCAR")
        if os.path.exists(contcar):
            self._read_contcar()

    def _read_vasprun_xml(self) -> None:
        """Read results from vasprun.xml."""
        from xml.etree import ElementTree

        vasprun_path = os.path.join(self.directory, "vasprun.xml")

        try:
            tree = ElementTree.parse(vasprun_path)
            root = tree.getroot()

            # Find last calculation (for relaxations)
            calcs = root.findall(".//calculation")
            if not calcs:
                return
            calc = calcs[-1]

            # Energy (use last SCF iteration value)
            energy_elems = calc.findall('.//energy/i[@name="e_0_energy"]')
            if energy_elems:
                self.results["energy"] = float(energy_elems[-1].text)

            # Free energy (use last SCF iteration value)
            free_energy_elems = calc.findall('.//energy/i[@name="e_fr_energy"]')
            if free_energy_elems:
                self.results["free_energy"] = float(free_energy_elems[-1].text)

            # Forces
            forces_elem = calc.find('.//varray[@name="forces"]')
            if forces_elem is not None:
                forces = []
                for v in forces_elem.findall("v"):
                    forces.append([float(x) for x in v.text.split()])
                # Unsort forces to match original atom order
                forces_array = np.array(forces)[self.resort]
                self.results["forces"] = forces_array

            # Stress
            stress_elem = calc.find('.//varray[@name="stress"]')
            if stress_elem is not None:
                stress = []
                for v in stress_elem.findall("v"):
                    stress.append([float(x) for x in v.text.split()])
                stress_array = np.array(stress)
                # Convert to Voigt notation and kBar to eV/A^3
                voigt = np.array(
                    [
                        stress_array[0, 0],
                        stress_array[1, 1],
                        stress_array[2, 2],
                        stress_array[1, 2],
                        stress_array[0, 2],
                        stress_array[0, 1],
                    ]
                )
                self.results["stress"] = voigt * -0.1  # kBar to eV/A^3

            # Fermi level
            fermi_elem = root.find('.//dos/i[@name="efermi"]')
            if fermi_elem is not None:
                self.results["fermi_level"] = float(fermi_elem.text)

            # Magnetic moment from XML - try different paths
            mag_elem = calc.find('.//separator[@name="magnetization"]')
            if mag_elem is not None:
                for i in mag_elem.findall("i"):
                    if "total" in i.attrib.get("name", ""):
                        self.results["magmom"] = float(i.text)

            # Read magmom from OUTCAR as fallback (more reliable)
            self._read_magnetic_moments()

            # Dipole moment - read from OUTCAR (more reliable than vasprun.xml)
            self._read_dipole_moment()

        except ElementTree.ParseError:
            # Fall back to OUTCAR
            self._read_outcar()

    def _read_outcar(self) -> None:
        """Read results from OUTCAR."""
        outcar_path = os.path.join(self.directory, "OUTCAR")

        if not os.path.exists(outcar_path):
            return

        with open(outcar_path) as f:
            content = f.read()

        # Energy (last occurrence)
        energy_matches = re.findall(r"free  energy   TOTEN\s*=\s*([-\d.]+)\s*eV", content)
        if energy_matches:
            self.results["energy"] = float(energy_matches[-1])

        # Fermi level
        fermi_match = re.search(r"E-fermi\s*:\s*([-\d.]+)", content)
        if fermi_match:
            self.results["fermi_level"] = float(fermi_match.group(1))

        # Forces (from last ionic step)
        forces_pattern = r"TOTAL-FORCE \(eV/Angst\)\s*-+\s*((?:\s*[-\d.]+\s+[-\d.]+\s+[-\d.]+\s+[-\d.]+\s+[-\d.]+\s+[-\d.]+\s*\n)+)"
        forces_matches = re.findall(forces_pattern, content)
        if forces_matches:
            forces = []
            for line in forces_matches[-1].strip().split("\n"):
                parts = line.split()
                if len(parts) >= 6:
                    forces.append([float(parts[3]), float(parts[4]), float(parts[5])])
            if forces:
                forces_array = np.array(forces)[self.resort]
                self.results["forces"] = forces_array

        # Stress
        stress_pattern = (
            r"in kB\s+([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)"
        )
        stress_matches = re.findall(stress_pattern, content)
        if stress_matches:
            stress = [float(x) for x in stress_matches[-1]]
            # Convert kBar to eV/A^3
            self.results["stress"] = np.array(stress) * -0.1

        # Magnetic moments
        self._read_magnetic_moments()

    def _read_magnetic_moments(self) -> None:
        """Read magnetic moments from OUTCAR into results."""
        outcar_path = os.path.join(self.directory, "OUTCAR")
        if not os.path.exists(outcar_path):
            return

        with open(outcar_path) as f:
            content = f.read()

        # Total magnetic moment
        match = re.search(r"number of electron\s+[\d.]+\s+magnetization\s+([-\d.]+)", content)
        if match:
            self.results["magmom"] = float(match.group(1))

        # Per-atom magnetic moments
        pattern = r"magnetization \(x\)\s*\n\s*\n((?:\s*\d+\s+[-\d.]+\s*\n)+)"
        matches = re.findall(pattern, content, re.DOTALL)
        if matches:
            magmoms = []
            for line in matches[-1].strip().split("\n"):
                parts = line.split()
                if len(parts) >= 2:
                    magmoms.append(float(parts[1]))
            if magmoms and hasattr(self, "resort"):
                self.results["magmoms"] = np.array(magmoms)[self.resort]

    def _read_dipole_moment(self) -> None:
        """Read dipole moment from OUTCAR."""
        outcar_path = os.path.join(self.directory, "OUTCAR")

        if not os.path.exists(outcar_path):
            return

        with open(outcar_path) as f:
            content = f.read()

        # Find all dipolmoment lines and take the last one
        pattern = r"dipolmoment\s+([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)"
        matches = re.findall(pattern, content)
        if matches:
            # Take the last occurrence (final converged value)
            dipole = [float(x) for x in matches[-1]]
            self.results["dipole"] = np.array(dipole)

    def _read_contcar(self) -> None:
        """Read atomic positions from CONTCAR."""
        from ase.io import read

        contcar = os.path.join(self.directory, "CONTCAR")
        if os.path.exists(contcar) and os.path.getsize(contcar) > 0:
            try:
                atoms = read(contcar, format="vasp")
                # Unsort to original order
                self.atoms.positions = atoms.positions[self.resort]
                self.atoms.cell = atoms.cell
            except Exception:
                pass

    def load_atoms(self) -> Atoms:
        """Load atoms from CONTCAR or POSCAR.

        Returns:
            ASE Atoms object with this calculator attached.
        """
        from ase.io import read

        contcar = os.path.join(self.directory, "CONTCAR")
        poscar = os.path.join(self.directory, "POSCAR")

        if os.path.exists(contcar) and os.path.getsize(contcar) > 0:
            atoms = read(contcar, format="vasp")
        elif os.path.exists(poscar):
            atoms = read(poscar, format="vasp")
        else:
            raise FileNotFoundError(f"No POSCAR/CONTCAR in {self.directory}")

        # Attach calculator and set up sorting
        atoms.calc = self
        self.atoms = atoms
        self._setup_sorting(atoms)

        # Apply bader charges if they were computed
        if "bader_charges" in self.results:
            atoms.set_initial_charges(self.results["bader_charges"])

        return atoms
