System.py

"""
System.py — QE-backed system and image utilities

This module provides a thin interface around Quantum ESPRESSO (QE) for setting
up calculation directories, writing inputs, running the QE plane wave executable `pw.x`, and parsing
results. It defines the data structures used by the NEB driver to manage
system-level configuration and per-image state (cell + atomic positions), and
to persist each image's history to disk as JSON.

Key classes
-----------
- PWXconfig: Runtime configuration for QE (cutoffs, smearing, k-point control,
  dispersion, pseudopotential class, relativity, pressure, etc.). The object
  auto‑persists to `<sysdir>/PWXconfig.json` whenever a field is changed.
- System: Owns a working directory (`sysdir`), ensures a clean layout, and
  carries a shared `PWXconfig`. Provides `get_image(...)` to construct `Image`
  objects bound to this system.
- Image: Represents a single structure state with both cell (3×3) and atomic
  positions (3×Nat). Knows how to write QE inputs, launch SCF / relax runs,
  and parse outputs back into energies, forces, stress, etc. Maintains a
  JSON history per image.
- OutputReturn: Lightweight container for parsed results (etot, evdw, h,
  walltime, cell, cell_vol, apos, stress, forces).
- IConfig: Placeholder for image‑level configuration (reserved for future use).

Helper functions
----------------
- read_json(path)  → Python object
- write_json(path, data)  → pretty‑printed JSON

Minimal example
---------------
```python
from neb.System import System
import numpy as np

s = System('work/sys1')        # creates/cleans directory, writes PWXconfig.json
s.pwx_config.atoms = ['Cu', 'Cu']
cell = np.eye(3)
apos = np.array([[0.0, 0.5],
                 [0.0, 0.5],
                 [0.0, 0.5]])  # fractional → will be converted to Å by Image
I = s.get_image(cell, apos)
out = I.calc_state(band_iter=0, incr_iter=False)
print(out.etot, out.forces.shape)
```

File layout (within `<sysdir>`)
-------------------------------
- `PWXconfig.json`  : Auto‑updated mirror of `PWXconfig`.
- `<id>[_id_mod]/` : Per‑image working directory (QE scratch, inputs/outputs).
- `<id>[_id_mod].json` : Per‑image history (cell, apos, energy, forces, iter).

Units & conventions
-------------------
- Internal arrays use Å, eV, eV/Å; pressure exposed as kbar via
  `PWXconfig.press` (converted internally to eV/ų). XML parsing converts from
  QE (Bohr/Hartree) using fixed constants noted inline.

Dependencies
------------
- Python ≥ 3.10, NumPy, psutil, BeautifulSoup4, lxml, SciPy (Rotation, optimizers),
  ASE (masses, numbers), and a working QE 7.x (`pw.x`) in PATH or via MPI runner.
"""
import glob
import subprocess
import os
from os import mkdir, remove, rmdir
from shutil import copytree, rmtree, copyfile
from copy import deepcopy
from dataclasses import dataclass, field, asdict
from os.path import isdir
import json
from shutil import rmtree
import psutil
import warnings
import re
import signal

def custom_formatwarning(msg, category, filename, lineno, line=None):
    return f"{filename}:{lineno}: {category.__name__}: {msg}\n"
warnings.formatwarning = custom_formatwarning

from bs4 import BeautifulSoup
from scipy.spatial.transform import Rotation as R
from scipy.optimize import differential_evolution
import numpy as np
from ase.data import atomic_masses, atomic_numbers


@dataclass
class PWXconfig:
    """
    PWXconfig — mapping of Python fields to Quantum ESPRESSO (QE) pw.x input keywords.

    This dataclass mirrors the subset of QE runtime options used by this package
    and writes an always-current mirror to `<sysdir>/PWXconfig.json`. Fields map
    to QE namelists/sections as follows (see `pw.x` input reference):

    • &CONTROL
      - `pseudo_dir`      → `pseudo_dir`
      - thresholds        → `etot_conv_thr`, `forc_conv_thr` (set in write step)
      - 'verbosity'       → 'verbosity'

    • &SYSTEM
      - `ecutwfc`         → `ecutwfc` (wavefunction cutoff, Ry)
      - `ecutrho`         → `ecutrho` (charge density cutoff, Ry)
      - `occupations`     → `occupations`
      - `smearing`        → `smearing`
      - `degauss`         → `degauss` (Ry)
      - `noncolin`        → `noncolin`
      - `lspinorb`        → `lspinorb` (requires `noncolin = .true.`)
      - `vdw_corr`        ← from `disp` (e.g., 'DFT-D3')
      - `dftd3_version`   ← from `dft3_version`
      - `nosym`           → `nosym`
      - `ecfixed`, `qcutz`, `q2sigma` are advanced/rarely used options passed
        as-is into &SYSTEM when set (retain here for compatibility).

    • &ELECTRONS
      - `mixing_beta`     → `mixing_beta`
      - `conv_thr`        → `conv_thr`

    • &CELL (vc-relax / variable-cell runs)
      - `press` (kbar)    → exposed as property; internally stored in eV/Å^3 as
        `_press` and converted back when writing inputs.
      - `press_conv_thr`  → `press_conv_thr`

    • K-POINTS
      - If `kspacing` is set (Å⁻¹), a Monkhorst–Pack mesh is computed from the
        reciprocal lattice lengths and written as `K_POINTS automatic`.
      - Otherwise `kgrid = [n1, n2, n3, 0, 0, 0]` is written verbatim.

    Pseudopotentials
    ----------------
    `atoms` holds the species sequence per atom in the structure (e.g., ['Cu','O',...]).
    Pseudopotentials are looked up from `<pseudo_dir>/pps.json` using the keys
    `func` (e.g., 'PBE'), `ppclass` ('NCPP'|'USPP'|'PAW'), and `rel` ('scalar'|'fully').
    """
    save_dir: str
    atoms: list[str] = None
    nproc: int = int(os.getenv("SLURM_NTASKS", psutil.cpu_count(logical=False)))  # MPI ranks for pw.x
    print('nproc=', nproc)
    # pseudopotential parameters - look at https://pseudopotentials.quantum-espresso.org/home/naming-convention
    # --------------------------------------------------------------------------------------------------
    func: str = 'pbe'  # exchange–correlation-functional: pbe | pz | pw91 ........
    pptype: str = 'rrkjus'  # PP class: rrkjus | rrkj | kjpaw ........
    ppfullrel: str = False  # Fully relativistic PP: true | false
    #--------------------------------------------------------------------------------------------------
    pseudo_dir: str = os.path.dirname(os.path.abspath(__file__))+'/pseudos/PSlibrary'  # &CONTROL:pseudo_dir
    max_seconds: int = 1e6       # &CONTROL: calculation time limit
    etot_conv_thr: float = 1e-5  # &CONTROL: etot_conv_thr
    forc_conv_thr: float = 1e-4  # &CONTROL: forc_conv_thr (eV/Å converted in writer if needed)
    verbosity: str = 'low'       # &CONTROL: verbosity of output
    disk_io: str = 'minimal'
    # --------------------------------------------------------------------------------------------------
    ecutwfc: float = 80    # &SYSTEM: ecutwfc (Ry)
    ecutrho: float = 400   # &SYSTEM: ecutrho (Ry)
    disp: str = 'DFT-D3'   # &SYSTEM: vdw_corr
    dft3_version: int = 4  # &SYSTEM: dftd3_version
    occupations: str = 'smearing'  # &SYSTEM: occupations
    noncolin: bool = False       # &SYSTEM: noncolin (needed for fully relativistic)
    lspinorb: bool = False       # &SYSTEM: lspinorb (requires noncolin)
    smearing: str = 'gauss'      # &SYSTEM: smearing
    degauss: float = 0.02        # &SYSTEM: degauss (Ry)
    ecfixed: float = 0           # &SYSTEM: advanced option (passed through)
    qcutz: float = 0             # &SYSTEM: advanced option (passed through)
    q2sigma: float = 0.1         # &SYSTEM: advanced option (passed through)
    nosym: bool = False          # &SYSTEM: nosym
    # --------------------------------------------------------------------------------------------------
    mixing_beta: float = 0.7     # &ELECTRONS: mixing_beta
    conv_thr: float = 1e-6       # &ELECTRONS: conv_thr
    diagonalization: str = 'david'
    # --------------------------------------------------------------------------------------------------
    press_conv_thr: float = 5e-2 # &CELL: press_conv_thr (kbar)
    # --------------------------------------------------------------------------------------------------
    kspacing: float = None  # K_POINTS automatic derived from spacing (Å⁻¹)
    kgrid: list[float] = field(default_factory=lambda: [6, 6, 6, 0, 0, 0])  # K_POINTS automatic (explicit grid)
    _press: float = 0            # internal storage in eV/Å^3; use property `press` in kbar
    # --------------------------------------------------------------------------------------------------
    _initialized: bool = field(default=False, init=False, repr=False)

    @property
    def press(self):
        return self._press * 1.602176634e3      # eV/A^3 to kbar
    @press.setter
    def press(self, value):
        self._press = value * 6.24150907446e-4  # kbar to eV/A^3

    def __post_init__(self):
        self._initialized = True
        write_json(self.save_dir + '/PWXconfig.json', asdict(self))

    def __setattr__(self, name, value):
        super().__setattr__(name, value)
        if getattr(self, '_initialized', False):
            write_json(self.save_dir + '/PWXconfig.json', asdict(self))

    def update_nproc(self):
        n = psutil.cpu_count(logical=False)
        self.nproc = n

@dataclass
class DOSXconfig:
    """
    DOSXconfig — configuration kept for Quantum ESPRESSO `dos.x`.

    What dos.x expects
    ------------------
    `dos.x` reads a single `&DOS` namelist with (most commonly used) keys:
     - `Emin`, `Emax` (eV)         : absolute energy window
     - `DeltaE` (eV)               : energy step
    Notes on fields below
    ---------------------
    Several attributes mirror PWXconfig for convenience and a consistent JSON shape.
    They are **not** used directly by `dos.x`. When writing `_write_dosx_input(...)`, only the
    relevant `&DOS` options should be emitted (`Emin/Emax/DeltaE`, `fildos`, `ngauss/degauss`, `bz_sum`).
    """
    #--------------------------------------------------------------------------------------------------
    Emin: float = None             # &DOS: lower energy bound (eV, relative to E_F)
    Emax: float = None              # &DOS: upper energy bound (eV, relative to E_F)
    DeltaE: float = 0.01            # &DOS: energy step (eV)
    # --------------------------------------------------------------------------------------------------


class System:
    """
    System — working directory manager and shared QE configuration

    Purpose
    -------
    Owns a calculation workspace (`sysdir`) and a shared `PWXconfig`. Ensures the
    directory exists and is clean, then provides convenience constructors for
    per-structure `Image` objects that read/write within this workspace.

    Lifecycle & Side Effects
    ------------------------
    - On construction, `System(sysdir)` **deletes** an existing directory of the
      same name (if present) and recreates it empty.
    - Writes `<sysdir>/PWXconfig.json` immediately via the embedded `PWXconfig`.
    - All images created via `get_image(...)` will store their history JSON and
      QE I/O files inside `<sysdir>`.

    Attributes
    ----------
    - sysdir (str): Absolute or relative path to the workspace, normalized to
      end with a trailing '/'.
    - pwx_config (PWXconfig): Shared runtime configuration for all images.

    Methods
    -------
    - get_image(cell, apos=None, apos_cart=None) -> Image:
      Construct an `Image` bound to this `System`.

    Notes
    -----
    This class is intentionally small: orchestration (NEB, interpolation, etc.)
    lives in `neb.NEB`, while `Image` handles QE I/O per structure.
    """
    def __init__(self, sysdir: str):
        """Create/clean a workspace and initialize shared QE configuration.

        Parameters
        ----------
        sysdir : str
            Path to the working directory. If it exists, it will be removed and
            recreated empty.

        Side Effects
        ------------
        - Creates `sysdir/` (fresh).
        - Instantiates `PWXconfig(save_dir=sysdir)` and writes `PWXconfig.json`.
        """
        self.sysdir = sysdir
        if isdir(sysdir):
            rmtree(sysdir)
        os.mkdir(sysdir)

        self.pwx_config = PWXconfig(sysdir)
        self.dosx_config = DOSXconfig(sysdir)

    def get_image(self, *args, **kwargs) -> 'Image':
        """Construct an `Image` bound to this `System`.

        Parameters
        ----------
        cell : np.ndarray, shape (3,3)
            Cell matrix in Å (columns are a1, a2, a3).
        apos : np.ndarray, optional, shape (3, Nat)
            Fractional positions; converted internally to Å as `cell @ apos`.
        apos_cart : np.ndarray, optional, shape (3, Nat)
            Cartesian positions in Å. Mutually exclusive with `apos`.

        Returns
        -------
        Image
            A new image whose files are placed under `<sysdir>`.
        """
        return Image(self, *args, **kwargs)


@dataclass
class IConfig:
    # =====================================
    # Currently not in use
    # =====================================
    pass

@dataclass
class OutputReturn:
    """
    OutputReturn — parsed results from a QE `pw.x` run

    This lightweight container is returned by `Image.calc_state(...)` and
    `Image.opt_image(...)`. It carries numeric results already converted to
    standard units used in this codebase.

    Fields & units
    --------------
    - etot (float): Electronic total energy (eV).
    - evdw (float): Dispersion correction term (eV); 0.0 if not present.
    - h (float): Enthalpy = etot + P·V with P taken from `PWXconfig._press` (eV).
    - walltime (float): QE wall time (seconds).
    - cell (np.ndarray): Cell matrix, shape (3,3), columns are a1,a2,a3 (Å).
    - cell_vol (float): Cell volume (Å^3).
    - apos (np.ndarray): Atomic positions, shape (3, Nat) in Å; origin-shifted so atom 0 is at (0,0,0).
    - stress (np.ndarray): Stress tensor, shape (3,3) in eV/Å^3.
    - forces (np.ndarray): Combined cell+atomic forces, shape (3, 3+Nat) in eV/Å.
        First 3 columns correspond to cell “forces” derived from stress; the
        remaining Nat columns are atomic forces in Cartesian coordinates.

    Notes
    -----
    - Unit conversions are performed from QE internal units (Bohr/Hartree) in
      `Image._read_output(...)`. See that method for constants.
    - Some fields may be `None` if QE did not print the corresponding quantity
      (e.g., `evdw` when no dispersion is enabled).
    """
    etot: float = None      # eV
    evdw: float = None      # eV
    h: float = None         # eV
    walltime: float = None # seconds
    cell: np.ndarray = None   # (3,3) Å
    cell_vol: float = None  # Å^3
    apos: np.ndarray = None   # (3,Nat) Å
    stress: np.ndarray = None # (3,3) eV/Å^3
    forces: np.ndarray = None # (3,3+Nat) eV/Å (cell|atomic)

class Image:
    """
    Image — a single structural state (cell + atomic positions) with QE I/O

    Purpose
    -------
    Encapsulates one configuration’s geometry and calculated quantities. Each
    instance has its own working directory `<sysdir>/<id>[_id_mod]/` and JSON
    history `<sysdir>/<id>[_id_mod].json`. Provides helpers to run QE (`pw.x`),
    parse results, and maintain state needed by NEB.

    Key attributes
    --------------
    - cell : (3,3) Å   — column vectors a1,a2,a3
    - apos : (3,Nat) Å — Cartesian atomic positions
    - r    : (3,3+Nat) — concatenated [cell|apos] convenience view
    - v,f  : (3,3+Nat) — update vector and forces (cell|atomic split)
    - state_history     — dict with lists for cell/apos/energy/forces/iteration
    - hdir, save_file   — per-image directory and JSON file path

    File layout
    -----------
    - `<hdir>/<prefix>.in/.out/.save`  — QE input/output/scratch files
    - `<save_file>`                    — persistent JSON history
    """
    def __init__(self, system: 'System', cell: np.ndarray, apos: np.ndarray = None, apos_cart: np.ndarray = None):
        """Create a new Image bound to a `System`.

        Parameters
        ----------
        system : System
            The owning `System`; supplies `sysdir` and shared `PWXconfig`.
        cell : np.ndarray, shape (3,3)
            Cell matrix in Å (columns are a1,a2,a3).
        apos : np.ndarray, optional, shape (3,Nat)
            Fractional positions; if given, converted to Cartesian as `cell @ apos`.
        apos_cart : np.ndarray, optional, shape (3,Nat)
            Cartesian positions in Å. Mutually exclusive with `apos`.

        Side effects
        ------------
        - Creates a per-image directory `<sysdir>/<id>/` and JSON history file.
        - Initializes force/velocity buffers and writes initial history.
        """
        if apos is None and apos_cart is None:
            print("either 'apos' or 'apos_cart' must be specified")
            exit(1)
        elif apos is not None and apos_cart is not None:
            print("'apos' and 'apos_cart' cannot be specified together")
            exit(1)
        elif apos is not None and apos_cart is None:
            apos = cell @ apos
        else:
            apos = apos_cart

        self.pwx_config = system.pwx_config
        self.dosx_config = system.dosx_config
        self.sysdir = system.sysdir

        self._id_mod = ''
        self.hdir = f'{self.sysdir}/{id(self)}'
        self.save_file = f'{self.sysdir}/{id(self)}.json'

        if not isdir(self.hdir):
            os.mkdir(self.hdir)

        self._energy = None
        self._enthalpy = None
        self._cell = cell
        self._apos = apos
        self._cell_prev = cell
        self._apos_prev = apos
        self._v_cell = np.zeros((3, 3))
        self._v_apos = np.zeros(apos.shape)
        self._f_cell = np.zeros((3, 3))
        self._f_apos = np.zeros(apos.shape)
        self._f_cell_prev = np.zeros((3, 3))
        self._f_apos_prev = np.zeros(apos.shape)
        self.calc_state_subcall_count = 0
        self.iter_count = -1
        self.state_history = {'iteration': self.iter_count, 'id_mod': None, 'pressure': self.pwx_config._press, 'cell': [],
                              'apos': [], 'energy': [], 'forces': []}
        self.save_history()

    def __copy__(self):
        """Shallow-copy the Image, duplicating metadata/history and directory layout."""
        cls = self.__class__
        result = cls.__new__(cls)
        result.__dict__.update(self.__dict__)
        result._id_mod = '_copy'
        result.state_history = deepcopy(self.state_history)
        result.__update_hdir(mv_files=False)
        return result

    def cleanup(self):
        """Remove this image's working directory and JSON history if they exist."""
        try:
            rmtree(self.hdir)
            remove(self.save_file)
        except OSError:
            pass

    def __update_hdir(self, mv_files: bool = True):
        """Refresh `hdir/save_file` after `id_mod` changes; optionally move files over."""
        old_dir = self.hdir
        old_save_file = self.save_file
        self.hdir = f'{self.sysdir}/{id(self)}{self.id_mod}'
        self.save_file = f'{self.sysdir}/{id(self)}{self.id_mod}.json'
        if isdir(old_dir) and mv_files:
            copytree(old_dir, self.hdir)
            copyfile(old_save_file, self.save_file)
            rmtree(old_dir)
            remove(old_save_file)
        else:
            mkdir(self.hdir)
            self.save_history()

    def save_history(self):
        """Write current `state_history` to `<save_file>` as pretty JSON."""
        write_json(self.save_file, self.state_history)

    def erase_history(self, iteration: int = -1):
        """Clear history arrays and set iteration counter, then persist to JSON."""
        self.state_history['cell'] = []
        self.state_history['apos'] = []
        self.state_history['energy'] = []
        self.state_history['forces'] = []
        self.state_history['iteration'] = iteration
        self.save_history()

    def opt_image(self,nosym: bool = False) -> OutputReturn:
        """Run a QE `vc-relax` for this image and update geometry from the result."""
        print(f'Optimizing Image {id(self)}{self.id_mod}')
        prefix = f"{self.iter_count}_relax"
        self.pwx_config.nosym = nosym
        self._construct_pwx_input('vc-relax', prefix)
        self._runpwx(prefix)
        self.pwx_config.nosym = False
        output = self._read_pwx_output(prefix)
        self.r = np.concatenate((output.cell,output.apos), 1)
        return output

    def calc_state(self, new_state = False, rm_old_files = True) -> OutputReturn:
        """Run a QE SCF (`pw.x`) and parse results.

        Parameters
        ----------
        new_state : bool
            If True, save results in image attributes.
        rm_old_files : bool
            If True, remove old files matching `[0-9]*_[0-9]*.*` in `hdir` first.


        Returns
        -------
        OutputReturn
            Parsed results (energies, forces, stress, cell, positions, walltime).
        """
        prefix = f'{self.iter_count}_{self.calc_state_subcall_count}'
        if rm_old_files:
            for f in glob.glob(f"{self.hdir}/[-0-9]*_[0-9]*.*"):
                try:
                    remove(f)
                except IsADirectoryError:
                    rmdir(f)
        self._construct_pwx_input('scf', prefix)
        ret = self._runpwx(prefix)

        # only here to restart the calculation after a bug where pw.x is frozen
        if ret['timed_out']: # if pw.x gets terminated by a sigint
            print(f'Restarting calculation {id(self)}{self.id_mod}/{prefix} because it timed out')
            self.pwx_config.nosym = True
            self._construct_pwx_input('scf', prefix)
            ret = self._runpwx(prefix)
            self.pwx_config.nosym = False

        # case for failed diagonalization
        if ret['qe_ec'] in [41,42,43,44,45,46,47,48,49,50,51,52]:  # ec's for non pos. def. S matrix
            print(f'Restarting calculation {id(self)}{self.id_mod}/{prefix} because of failed diagonalization.')
            old_diag = self.pwx_config.diagonalization
            self.pwx_config.diagonalization = 'cg'
            self._construct_pwx_input('scf', prefix)
            ret = self._runpwx(prefix)
            self.pwx_config.diagonalization = old_diag

        # case for symmetry-related failures (e.g. sym_rho_init_shell, lone vector)
        if (ret['qe_routine'] is not None
                and ret['qe_routine'].startswith('sym_')
                and not self.pwx_config.nosym):
            print(
                f"Restarting calculation {id(self)}{self.id_mod}/{prefix} without symmetry "
                f"because QE failed in routine {ret['qe_routine']} ({ret['qe_ec']}): {ret['qe_msg']}"
            )
            old_nosym = self.pwx_config.nosym
            self.pwx_config.nosym = True
            self._construct_pwx_input('scf', prefix)
            ret = self._runpwx(prefix)
            self.pwx_config.nosym = old_nosym

        output = self._read_pwx_output(prefix)
        self.calc_state_subcall_count += 1

        if new_state:
            self.e = output.etot
            self.h = output.h
            self.f = output.forces


        return output

    def calc_nscf(self, kgrid_fct = 2):
        """Run a QE `nscf` for this image"""
        print(f'Running NSCF cycle for image {id(self)}{self.id_mod}')
        prefix = f'{self.iter_count}_nscf'

        # recalculate charge density
        self._construct_pwx_input('scf', prefix)
        self._runpwx(prefix, keep_bnds=True) # keep_bnds important to not mess with the formatting that QE expects

        old_kgrid = self.pwx_config.kgrid
        old_occ = self.pwx_config.occupations
        for i in range(3):
            self.pwx_config.kgrid[i] = int(np.ceil(self.pwx_config.kgrid[i] * kgrid_fct))
        self.pwx_config.occupations = 'tetrahedra_opt'

        self._construct_pwx_input('nscf', prefix)
        self._runpwx(prefix,  keep_bnds=True)

        self.pwx_config.kgrid = old_kgrid
        self.pwx_config.occupations = old_occ

    def calc_bandgap(self):
        """Compute the band gap (eV) by running `dos.x` on this image's latest NSCF result.

        Returns
        -------
        dict
            {
              'Eg': float | None,           # eV; 0.0 for metals/overlap
              'VBM': float | None,   # eV (absolute); None if not determined
              'CBM': float | None,   # eV (absolute); None if not determined
              'Ef': float | None,# eV (absolute)
              'is_metal': bool
            }
        """
        print(f'Evaluating band gap for image {id(self)}{self.id_mod}')

        nscf_calcs = glob.glob(f"{self.hdir}/[-0-9]*_nscf.save/")
        # Choose the NSCF result with the highest numeric prefix, e.g., "12_nscf.save"
        if not nscf_calcs:
            warnings.warn(f"No NSCF results found in '{self.hdir}' matching '*_nscf.save'", UserWarning)
            return {'Eg': None, 'VBM': None, 'CBM': None, 'Efermi': None, 'is_metal': None, 'path': None}
        def _iter_from_path(p: str) -> int:
            return int(p.split('/')[-2].split('_')[0])
        latest_calc_ind = _iter_from_path(max(nscf_calcs, key=_iter_from_path))
        if latest_calc_ind != self.iter_count:
            warnings.warn('NSCF calculation is not up to date! Make a new one to use the latest band geometry.', UserWarning)
        nscf_prefix = f"{latest_calc_ind}_nscf"

        # 1) Read Fermi energy from XML: <espresso><output><band_structure><fermi_energy>
        # QE XML energies are in Hartree; convert to eV.
        energy_conf_fact = 2 * 13.6057039763  # Ha -> eV
        nscf_xml = f"{self.hdir}/{nscf_prefix}.save/data-file-schema.xml"
        try:
            with open(nscf_xml, 'r') as f:
                xml_text = f.read()
        except FileNotFoundError:
            warnings.warn(f"NSCF XML not found: {nscf_xml}", UserWarning)
            return {'Eg': None, 'VBM': None, 'CBM': None, 'Efermi': None, 'is_metal': None, 'path': None}
        xml_file = BeautifulSoup(xml_text, 'xml')
        Ef = energy_conf_fact * float(xml_file.espresso.output.band_structure.fermi_energy.string)

        # 2) Write dos.x input and run it
        dos_prefix = f"{latest_calc_ind}_dos"
        self.dosx_config.Emin = Ef-20
        self.dosx_config.Emax = Ef+20
        self._construct_dosx_input(prefix=dos_prefix, nscf_prefix=nscf_prefix)
        rc = self._rundosx(dos_prefix)
        if rc != 0:
            warnings.warn(f"dos.x returned non-zero exit status {rc}", UserWarning)

        # 3) Load the DOS file produced by dos.x
        dos_path = f"{self.hdir}/{dos_prefix}.dat"
        try:
            data = np.loadtxt(dos_path, comments='#', usecols=(0, 1))
        except Exception as e:
            warnings.warn(f"Failed to read DOS file '{dos_path}': {e}", UserWarning)
            return {'Eg': None, 'VBM': None, 'CBM': None, 'Efermi': Ef, 'is_metal': None}

        # Columns from dos.x: first is energy (scaled, such that Ef=0), second is DOS
        E = data[:, 0].astype(float)
        DOS  = data[:, 1].astype(float)

        # --- Adaptive DOS threshold ---
        # Absolute floor plus a fraction of max(DOS) to be scale-invariant.
        tol = 1e-6

        # Helper: linear interpolation of the x where y crosses 'thr' between (x0,y0) and (x1,y1)
        def _interp_cross(x0, y0, x1, y1, thr):
            if y1 == y0:
                return x0
            t = (thr - y0) / (y1 - y0)
            return x0 + t * (x1 - x0)

        # Identify indices with DOS >= tol on each side of E=0
        left_idx  = np.where((E <= Ef) & (DOS >= tol))[0]
        right_idx = np.where((E > Ef) & (DOS >= tol))[0]

        VBM = None
        CBM = None

        # VBM: highest energy below 0 with significant DOS; optionally interpolate to the threshold
        if left_idx.size > 0:
            k = left_idx[-1]
            VBM = E[k]

        # CBM: lowest energy above 0 with significant DOS; optionally interpolate to the threshold
        if right_idx.size > 0:
            k = right_idx[0]
            CBM = E[k-1]

        # Decide metallicity and compute gap. Allow tiny smeared DOS at E=0 if a finite gap is found.
        Eg = float(CBM - VBM)
        if Eg < 1e-6:
            Eg = 0.0
            is_metal = True
        else:
            is_metal = False

        return {
            'Eg': Eg,
            'VBM': VBM if not is_metal else None,
            'CBM': CBM if not is_metal else None,
            'Ef': Ef,
            'is_metal': is_metal
        }

    def minimize_distance(self, I_ref: 'Image'):
        """Rotate/reflect/permute this image to best match a reference cell (least-squares)."""
        cell1 = I_ref.cell
        cell2 = self.cell

        reflections = []
        for sx in [1, -1]:
            for sy in [1, -1]:
                for sz in [1, -1]:
                    reflections.append(np.diag([sx, sy, sz]))

        P = np.array([[0, 1, 0],
                      [0, 0, 1],
                      [1, 0, 0]])
        RP = np.array([[0, 0, 1],
                       [0, 1, 0],
                       [1, 0, 0]])
        permutations = [np.identity(3), P, P@P, RP, P@RP, P@P@RP]

        best_results = []
        for refl in reflections:
            for perm in permutations:
                def dist_rot_only(x):  # Euler angles only
                    rot = R.from_euler(seq='xyz', angles=x, degrees=True).as_matrix()
                    cell2_transformed = rot @ refl @ (perm @ cell2.T).T
                    return np.linalg.norm(cell1 - cell2_transformed)

                result = differential_evolution(dist_rot_only, [(-180, 180)] * 3, tol=1e-12, polish=True)
                best_results.append((result.fun, result.x, refl, perm))

        # Find the one with minimal distance
        best_result = min(best_results, key=lambda t: t[0])
        # Keep only results that are within 1e-8 of the best minimum
        best_results = [res for res in best_results if abs(res[0] - best_result[0]) <= 1e-8]

        # Keep only solutions with the minimal number of reflection + permutation changes
        def count_non_identity(mat):
            return np.sum(~np.isclose(mat, np.eye(3)))

        min_complexity = min(count_non_identity(res[2] @ res[3]) for res in best_results)
        best_results = [res for res in best_results if count_non_identity(res[2] @ res[3]) == min_complexity]
        best_result = best_results[0]

        x0 = best_result[1]
        rot = R.from_euler(seq='xyz', angles=x0, degrees=True).as_matrix()
        refl = best_result[2]
        perm = best_result[3]

        self.cell = rot @ refl @ (perm @ self.cell.T).T
        self.apos = rot @ refl @ self.apos

    def kpts(self) -> str:
        """Compute a Monkhorst–Pack grid string from `PWXconfig.kspacing` and the cell."""
        n1, n2, n3 = self._kpts(self.cell, self.pwx_config.kspacing)
        return f'{n1} {n2} {n3} 0 0 0'

    def _append_history(self):
        """Append current e/cell/apos/forces to history, increment iteration, and persist."""
        self.state_history['energy'].append(self.e)
        self.state_history['cell'].append(list(self.cell.flatten()))
        self.state_history['apos'].append(list(self.apos.flatten()))
        self.state_history['forces'].append(list(self.f.flatten()))
        self.state_history['iteration'] = self.iter_count
        self.save_history()

    def increase_iteration(self):
        self.iter_count += 1
        self.calc_state_subcall_count = 0
        self._append_history()

    def _construct_pwx_input(self, calc, prefix):
        """Assemble QE pw.x input dictionaries and write `<hdir>/<prefix>.in`."""
        pwx_config = self.pwx_config
        species = list(set(self.pwx_config.atoms))
        pwx_options = {
            '&control': {
                'calculation': calc,
                'pseudo_dir': pwx_config.pseudo_dir,
                'prefix': prefix,
                'outdir': self.hdir,
                'disk_io': pwx_config.disk_io,
                'tstress': True,
                'tprnfor': True,
                'max_seconds': int(pwx_config.max_seconds),
                'etot_conv_thr': pwx_config.etot_conv_thr,
                'forc_conv_thr' : pwx_config.forc_conv_thr,
                'verbosity': pwx_config.verbosity
            },
            '&system': {
                'ecutwfc': pwx_config.ecutwfc,
                'ecutrho': pwx_config.ecutrho,
                'vdw_corr': pwx_config.disp,
                'dftd3_version': pwx_config.dft3_version,
                'ibrav': 0,
                'nat': len(self.pwx_config.atoms),
                'ntyp': len(list(set(self.pwx_config.atoms))),
                'occupations': pwx_config.occupations,
                'noncolin': pwx_config.noncolin,
                'lspinorb': pwx_config.lspinorb,
                'smearing': pwx_config.smearing,
                'degauss': pwx_config.degauss,
                'ecfixed': pwx_config.ecfixed,
                'qcutz': pwx_config.qcutz,
                'q2sigma': pwx_config.q2sigma,
                'nosym' : pwx_config.nosym
            },
            '&electrons': {
                'mixing_beta': pwx_config.mixing_beta,
                'conv_thr': pwx_config.conv_thr,
                'diagonalization': pwx_config.diagonalization
            },
            '&ions': {
            },
            '&cell': {
                'press': pwx_config.press,
                'press_conv_thr': pwx_config.press_conv_thr
            },
            '&fcp': {
            },
            '&rism': {
            },
            'atomic_species': [
                f'{species[i]} {str(atomic_masses[atomic_numbers[species[i]]])} '
                f'{self._find_pp(species[i])}' for i in range(len(species))
            ],
            'atomic_positions angstrom': [
                f'{self.pwx_config.atoms[i]} {self.apos[0,i]} {self.apos[1,i]} {self.apos[2,i]}' for i in range(len(self.pwx_config.atoms))
            ],
            'cell_parameters angstrom': [
                ' '.join([str(y) for y in x]) for x in list(self.cell.transpose())
            ],
            'k_points automatic': [self.kpts()] if pwx_config.kspacing is not None
                                                else [' '.join([str(x) for x in pwx_config.kgrid])]
        }

        self._write_qe_input(prefix=prefix,option_dict=pwx_options)

    def _construct_dosx_input(self, prefix, nscf_prefix):
        """Assemble QE dos.x input and write `<hdir>/<prefix>.in`."""
        dosx_config = self.dosx_config
        dosx_options = {
            '&dos': {
                'outdir': self.hdir,
                'prefix': nscf_prefix,
                'fildos': f'{self.hdir}/{prefix}.dat',
                'Emin': dosx_config.Emin,
                'Emax': dosx_config.Emax,
                'DeltaE': dosx_config.DeltaE,
            }
        }
        self._write_qe_input(prefix=prefix, option_dict=dosx_options)

    def _write_qe_input(self, prefix: str, option_dict: dict):
        fh = open(f'{self.hdir}/{prefix}.in', 'w')
        for key, value in option_dict.items():
            if key.startswith('&'):
                fh.write(key.upper() + '\n')
                for k, v in value.items():
                    if isinstance(v, str):
                        v = "'" + v + "'"
                    elif isinstance(v, bool):
                        v = '.true.' if v else '.false.'
                    fh.write('   ' + k + ' = ' + str(v) + '\n')
                fh.write('/\n\n')
            else:
                l = key.split()
                l[0] = l[0].upper()
                key = ' '.join(l)
                fh.write(key + '\n')
                for v in value:
                    fh.write(str(v) + '\n')
                fh.write('\n')
        fh.close()

    def _rundosx(self, prefix) -> int:
        file = f'{self.hdir}/{prefix}'
        cp = subprocess.run(['dos.x', '-in', f'{file}.in'], capture_output=True, text=True)
        with open(f'{file}.out', 'w') as f:
            f.write(cp.stdout + cp.stderr)
        return cp.returncode

    def _runpwx(self, prefix, keep_bnds: bool = False) -> dict:
        """Launch `pw.x` via `mpirun -np nproc` with a watchdog.

        If the process exceeds `timeout_s` seconds, send SIGINT to the whole
        process group (so `mpirun` and all ranks receive it). If it still
        doesn't exit within `grace_s` seconds, escalate to SIGKILL.

        Parameters
        ----------
        prefix : str
            File prefix written to `<hdir>/<prefix>.in/.out/.save`.
        keep_bnds : bool
            Keep band_structure section in XML if True.
        """
        file = f'{self.hdir}/{prefix}'
        timeout_s = int(self.pwx_config.max_seconds)

        # Start under a new process group so we can signal mpirun and all ranks.
        proc = subprocess.Popen(
            ['mpirun', '-np', str(self.pwx_config.nproc), 'pw.x', '-in', f'{file}.in'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            start_new_session=True  # posix: setsid(); allows os.killpg
        )

        timed_out = False
        try:
            stdout, stderr = proc.communicate(timeout=timeout_s)
        except subprocess.TimeoutExpired:
            timed_out = True
            os.killpg(proc.pid, signal.SIGINT)
            stdout, stderr = proc.communicate(timeout=5)

        # Write combined output
        with open(f'{file}.out', 'w') as f:
            f.write((stdout or '') + (stderr or ''))

        # Inspect QE-style error:
        #   Error in routine cdiaghg (50):
        #   S matrix not positive definite
        combined_out = (stdout or '') + '\n' + (stderr or '')
        qe_match = re.search(
            r"Error in routine\s+(\S+)\s*\((\d+)\):\s*\n\s*(.+)",
            combined_out
        )
        if qe_match:
            qe_routine = qe_match.group(1)
            qe_ec = int(qe_match.group(2))
            qe_msg = qe_match.group(3).strip()
        else:
            qe_routine = None
            qe_ec = None
            qe_msg = None

        # Inspect stderr for `forrtl: error (N): ...`
        match = re.search(r"forrtl: error \((\d+)\): (.*)", stderr or '')
        if match:
            fortran_ec = int(match.group(1))
            fortran_msg = match.group(2).strip()
        else:
            fortran_ec = None
            fortran_msg = None

        # Post-processing: optionally shrink save folder
        if not keep_bnds:  # delete bands structure for storage efficiency
            output_file = f'{self.hdir}/{prefix}.save/data-file-schema.xml'
            try:
                with open(output_file, 'r') as f:
                    data = f.read()
                xml_data = BeautifulSoup(data, 'xml')
                xml_data.espresso.output.band_structure.clear()
                with open(output_file, 'w') as f:
                    f.write(xml_data.prettify())
            except FileNotFoundError:
                warnings.warn(f"rm_band_structure: File '{output_file}' not found", UserWarning)

        return {
            'rc': proc.returncode,
            'f_ec': fortran_ec,
            'f_msg': fortran_msg,
            'qe_routine': qe_routine,
            'qe_ec': qe_ec,
            'qe_msg': qe_msg,
            'timed_out': timed_out,
        }

    def _read_pwx_output(self, prefix) -> OutputReturn:
        """Parse QE XML (`data-file-schema.xml`) into `OutputReturn` with eV/Å units."""
        output_file = f'{self.hdir}/{prefix}.save/data-file-schema.xml'
        try:
            f = open(output_file, 'r')
            data = f.read()
        except FileNotFoundError:
            warnings.warn(f'read_output: File \'{output_file}\' not found', UserWarning)
            return OutputReturn()
        f.close()
        xml_data = BeautifulSoup(data, 'xml')
        espresso_xml = xml_data.espresso

        energy_conf_fact = 2 * 13.6057039763                        # hartree to eV
        length_conv_fact = 0.52917721067121                         # bohr to angstrom
        force_conf_fact = energy_conf_fact / length_conv_fact       # Ha/bohr to eV/A
        press_conv_fact = energy_conf_fact / length_conv_fact ** 3  # Ha/bohr^3 to eV/A^3


        int_energy = energy_conf_fact * float(espresso_xml.output.total_energy.etot.string)

        try:
            disp_energy = energy_conf_fact * float(espresso_xml.output.total_energy.vdW_term.string)
        except AttributeError:
            disp_energy = 0

        walltime = float(espresso_xml.timing_info.total.wall.string)

        a1 = [float(x) * length_conv_fact for x in espresso_xml.output.atomic_structure.cell.a1.string.split()]
        a2 = [float(x) * length_conv_fact for x in espresso_xml.output.atomic_structure.cell.a2.string.split()]
        a3 = [float(x) * length_conv_fact for x in espresso_xml.output.atomic_structure.cell.a3.string.split()]
        cell = np.transpose(np.asarray([a1,a2,a3]))
        cell_vol = np.linalg.det(cell)

        enthalpy = int_energy + self.pwx_config._press * cell_vol

        apos = [[float(pos) * length_conv_fact for pos in child.string.split()] for child in
                espresso_xml.output.atomic_structure.atomic_positions.find_all(recurlive=False)]
        apos = np.asarray(apos).transpose()
        for i in range(apos.shape[1]):
            apos[:,i] -= apos[:,0]

        aforces = force_conf_fact *  np.array([[float(x) for x in x.split()] for x in espresso_xml.output.forces.string.split('\n')[1:-1]]).transpose()

        for i in range(aforces.shape[1]):
            aforces[:,i] -= aforces[:,0]

        stress = press_conv_fact * np.array([float(x) for x in espresso_xml.output.stress.string.split()]).reshape((3, 3))

        cellforces = cell_vol * (stress - self.pwx_config._press * np.identity(3)) @ (np.linalg.inv(cell).transpose())
        forces = np.concatenate((cellforces, aforces), 1)

        output = OutputReturn()
        output.etot = int_energy
        output.evdw = disp_energy
        output.h = enthalpy
        output.cell_vol = cell_vol
        output.apos = apos
        output.cell = cell
        output.stress = stress
        output.forces = forces
        output.walltime = walltime

        return output

    def _find_pp(self, element):
        """Return the *best* matching UPF path for `element`.

        Matching filters (required): func (field3), pptype (field5), and
        full-relativistic flag (field1 'rel' if ppfullrel is True).

        Preference among matches (more exact core configuration):
        1) 'l' present in field4 (long projectors) — higher is better
        2) More semicore letters included among 's','p','d' — higher is better
        3) 'n' present (NLCC) — higher is better
        4) Higher semantic version in field6 (e.g., _1.0.0) — higher is better
        5) Filename (lexicographic) as final tiebreaker
        """
        base_dir = self.pwx_config.pseudo_dir
        want_func = str(self.pwx_config.func).lower()          # e.g. 'pbe'
        want_type = str(self.pwx_config.pptype).lower()        # e.g. 'rrkjus'
        want_rel = bool(self.pwx_config.ppfullrel)             # True => require 'rel'

        def parse_version(vstr: str):
            if not vstr:
                return tuple()
            parts = vstr.split('.')
            out = []
            for p in parts:
                try:
                    out.append(int(p))
                except ValueError:
                    out.append(p)
            return tuple(out)

        def core_token_from_tokens(tokens: list[str]) -> str:
            """Extract field4 token (concatenated flags like 'spnl') if present."""
            # field4 is typically the token made of only letters from this set
            CORE_CHARS = set('spdnfl')
            for t in tokens:
                tl = t.lower()
                if tl and set(tl) <= CORE_CHARS and tl not in {'rel', want_func, want_type}:
                    return tl
            return ''

        def core_score(core_token: str) -> tuple:
            """Compute preference score: (has_l, semicore_count, has_n)."""
            if not core_token:
                return (0, 0, 0)
            has_l = 1 if 'l' in core_token else 0
            semi = sum(1 for c in 'spd' if c in core_token)
            has_n = 1 if 'n' in core_token else 0
            return (has_l, semi, has_n)

        candidates = []  # (score_tuple, version_tuple, filename, filename)

        for root, _, files in os.walk(base_dir):
            for fn in files:
                if not (fn.endswith('.UPF') or fn.endswith('.upf')):
                    continue
                stem = fn[:-4]
                if '.' not in stem:
                    continue
                sym, desc = stem.split('.', 1)
                if sym != element:
                    continue

                # Split optional trailing version (field6) `..._X.Y.Z`
                if '_' in desc:
                    desc_core, version = desc.rsplit('_', 1)
                else:
                    desc_core, version = desc, ''

                tokens = [t.lower() for t in desc_core.split('-') if t]
                token_set = set(tokens)

                has_func = want_func in token_set
                has_type = want_type in token_set
                has_rel = 'rel' in token_set

                if not (has_func and has_type):
                    continue
                if want_rel != has_rel:
                    continue

                ctok = core_token_from_tokens(tokens)
                score = core_score(ctok)
                vtuple = parse_version(version)
                # Instead of fullpath, store only the filename (fn)
                candidates.append((score, vtuple, fn.lower(), fn))

        if not candidates:
            raise FileNotFoundError(
                f"No matching UPF for element '{element}' in '{base_dir}' with func={want_func}, pptype={want_type}, fullrel={want_rel}")

        # Sort by: core score (desc), version (desc), filename (desc for consistency)
        candidates.sort(key=lambda t: (t[0], t[1], t[2]), reverse=True)
        return candidates[0][3]

    @property
    def id_mod(self):
        """Get suffix used in directory/file naming ('' or f'_{n}')."""
        return self._id_mod

    @id_mod.setter
    def id_mod(self, id_mod):
        """Set id modifier, update paths, and migrate files if needed."""
        if id_mod is None:
            self._id_mod = ''
            self.state_history['id_mod'] = id_mod
        else:
            self._id_mod = f'_{id_mod}'
            self.state_history['id_mod'] = id_mod
        self.__update_hdir()

    @property
    def cell(self) -> np.ndarray:
        """Return a copy of the 3×3 cell matrix (Å)."""
        return np.array(self._cell)

    @cell.setter
    def cell(self, cell: np.ndarray):
        """Set cell matrix (Å)."""
        self._cell = cell

    @property
    def apos(self) -> np.ndarray:
        """Return a copy of atomic positions (3×Nat, Å)."""
        return np.array(self._apos)

    @apos.setter
    def apos(self, apos: np.ndarray):
        """Set atomic positions (Å)."""
        self._apos = apos

    @property
    def r(self) -> np.ndarray:
        """Concatenated state [cell|apos] for convenience (3×(3+Nat))."""
        return np.concatenate((self._cell, self._apos), 1)

    @r.setter
    def r(self,r):
        """Set concatenated state; splits into cell and apos."""
        self._cell = r[:, :3]
        self._apos = r[:,3:]

    @property
    def r_prev(self) -> np.ndarray:
        """Previous step's concatenated state [cell|apos]."""
        return np.concatenate((self._cell_prev, self._apos_prev), 1)

    @r_prev.setter
    def r_prev(self, r_prev):
        """Set previous state; splits into cell_prev and apos_prev."""
        self._cell_prev = r_prev[:, :3]
        self._apos_prev = r_prev[:, 3:]

    @property
    def v(self) -> np.ndarray:
        """Return concatenated update vector v (3×(3+Nat))."""
        return np.concatenate((self._v_cell, self._v_apos), 1)

    @v.setter
    def v(self, v: np.ndarray):
        """Set update vector; splits into v_cell and v_apos."""
        self._v_cell = v[:, :3]
        self._v_apos = v[:, 3:]

    @property
    def f(self) -> np.ndarray:
        """Return concatenated forces F (3×(3+Nat))."""
        return np.concatenate((self._f_cell, self._f_apos), 1)

    @f.setter
    def f(self, f: np.ndarray):
        """Set forces; splits into f_cell and f_apos."""
        self._f_cell = f[:, :3]
        self._f_apos = f[:, 3:]

    @property
    def e(self):
        """Electronic total energy (eV)."""
        return self._energy

    @e.setter
    def e(self, e):
        """Set electronic total energy (eV)."""
        self._energy = e

    @property
    def h(self):
        """Enthalpy (eV)."""
        return self._enthalpy

    @h.setter
    def h(self, h):
        """Set enthalpy (eV)."""
        self._enthalpy = h

    @staticmethod
    def _kpts(cell: np.ndarray, kspacing: float) -> tuple[int, int, int]:
        """Return (n1,n2,n3) from reciprocal vector norms and target spacing (Å⁻¹)."""
        denom = cell[:, 0] @ np.cross(cell[:, 1], cell[:, 2])
        b1 = 2*np.pi * np.cross(cell[:, 1], cell[:, 2]) / denom
        b2 = 2*np.pi * np.cross(cell[:, 2], cell[:, 0]) / denom
        b3 = 2*np.pi * np.cross(cell[:, 0], cell[:, 1]) / denom
        n1 = round(np.linalg.norm(b1) / kspacing, 0)
        n2 = round(np.linalg.norm(b2) / kspacing, 0)
        n3 = round(np.linalg.norm(b3) / kspacing, 0)
        return int(n1), int(n2), int(n3)


def read_json(path):
    with open(path, 'r') as f:
        return json.load(f)

def write_json(path, data):
    with open(path, 'w') as f:
        json.dump(data, f, indent=4)