"""
YAIV | yaiv.phonon
=======================

This module provides tools to handle vibrational properties of crystals from
first-principles calculations.

It includes data structures and utilities to:
- Read `.dyn*` files generated by `ph.x`.
- Diagonalize phonon Hamiltonians to extract frequencies and modes.
- Compute commensurate supercells for charge density wave (CDW) distortions.
t- Build distorted atomic configurations based on soft phonon modes.
- Creation Born-Oppenheimer energy surfaces.

The main abstractions in this module are:

- `Dyn`: Represents a single dynamical matrix at a given q-point. Provides diagonalization and unit conversion.
- `CDW`: Builds a supercell compatible with multiple q-points and applies structural distortions.
- `BOES`: Constructs Born-Oppenheimer energy surfaces associated with charge-density wave (CDW) distortions.

Classes
-------
Dyn
    Container for a dynamical matrix at a specific q-point. Provides file I/O, diagonalization,
    and access to phonon frequencies and displacements.
    Provides:
    - from_file(): Class method to construct a `Dyn` object by reading a `.dyn*` file from Quantum ESPRESSO output.
    - diagonalize(): Diagonalizes the dynamical matrix, returning phonon frequencies and eigenvectors.

CDW
    Constructs and manipulates charge-density wave (CDW) distorted supercells from
    multiple q-point phonon modes.
    Provides:
    - from_files(): Class method to initialize a `CDW` object by reading `.dyn*` files at one or more q-points.
    - distort_crystal(): Applies the displacement pattern of selected phonon modes to the supercell.
    - sym_analysis(): Perform a symmetry analysis of commensurate CDW distortions across a grid of order parameters.

BOES
    Class for constructing a Born-Oppenheimer energy surface (BOES) associated with charge-density wave (CDW) distortions.
    Provides:
    - from_file(): Load a BOES object from a .pkl file.
    - save_as(): Save the entire BOES object as a .pkl file.
    - generate_structures_grid(): Generate distorted structures over a regular grid in order parameter space.
    - generate_structures_line(): Generate distorted structures along a 1D path in order parameter space.
    - save_jobs_pwi(): Export all generated distorted structures as QE input files (.pwi).
    - read_energies_pwo(): Grep total energies from Quantum ESPRESSO output files (.pwo).

Private Utilities
-----------------
_QEdyn2Realdyn(dyn_mat, masses)
    Convert a QE-format dynamical matrix (with √(mᵢ mⱼ) prefactors) to the true physical matrix
    whose eigenvalues are phonon frequencies squared (in s⁻² or (Ry/h)²).

_find_supercell(q_cryst)
    Given one or more q-points, computes the minimal supercell for which exp(2πi·q·R)
    is periodic over the lattice. Returns supercell size and grid of phase factors.

_format_input(CDW, amplitudes, modes, grid)
    Formats and normalizes user input for a CDW distortion calculation.

Examples
--------
>>> from yaiv.phonon import Dyn
>>> dyn = Dyn.from_file(q_cryst=[0.5, 0.0, 0.0], results_ph_path="ph_output/")
>>> dyn.diagonalize()
>>> dyn.freqs  # phonon frequencies in cm⁻¹

>>> from yaiv.phonon import CDW
>>> cdw = CDW.from_file(q_cryst=[[0,0,0],[1/2, 0.0, 0.0]], results_ph_path="ph_output/")
>>> distorted = cdw.distort_crystal(amplitudes=[1,1/5], modes=[0,0])

See Also
--------
yaiv.grep         : Greps structural and vibrational data from output files
yaiv.cell         : Atomic structure container and symmetry wrapper
yaiv.utils        : Unit management and basis transformation utilities
"""

from types import SimpleNamespace
import warnings

import numpy as np
import spglib as spg

from yaiv.defaults.config import ureg, defaults

from yaiv import utils as ut
from yaiv import grep
from yaiv import cell

__all__ = ["Dyn", "CDW", "BOES"]


def _QEdyn2Realdyn(
    dyn_mat: np.ndarray | ureg.Quantity, masses: np.ndarray | ureg.Quantity
) -> np.ndarray | ureg.Quantity:
    """
    Convert Quantum ESPRESSO dynamical matrix to the real dynamical matrix in physical units.

    In Quantum ESPRESSO, the dynamical matrix is written in units of [energy/h]^2 * [mass].
    To obtain the true dynamical matrix, whose eigenvalues are squared phonon frequencies (ν²) in
    units of [energy/h]^2, one must divide by √(m_i * m_j) for each atomic block.

    Parameters
    ----------
    dyn_mat : np.ndarray | ureg.Quantity
        The 3N×3N dynamical matrix as written by Quantum ESPRESSO.

    masses : np.ndarray | ureg.Quantity
        Array of atomic masses (length N), one per atom. Units should be consistent
        with those implied in `dyn_mat`. In QE, masses are typically given in units of 2 m_e.

    Returns
    -------
    np.ndarray or Quantity
        The real dynamical matrix with proper normalization, whose eigenvalues
        have units of squared frequencies, e.g., [Ry²/h²] or [s⁻²], depending on units.

    Raises
    ------
    ValueError
        If input shapes are inconsistent or mass array is not one-dimensional.
    """
    # Check unit consistency
    ut._check_unit_consistency([dyn_mat, masses], ["Dynamical matrix", "Atomic masses"])

    # Strip units if present
    if isinstance(dyn_mat, ureg.Quantity):
        units = dyn_mat.units / masses.units
        dyn_mat = dyn_mat.magnitude
        masses = masses.magnitude
    else:
        units = 1

    # Sanity checks
    masses = np.asarray(masses)
    dyn_mat = np.asarray(dyn_mat)

    if masses.ndim != 1:
        raise ValueError("Masses must be a 1D array (one mass per atom).")
    if dyn_mat.shape[0] != dyn_mat.shape[1]:
        raise ValueError("Dynamical matrix must be square.")
    if dyn_mat.shape[0] != 3 * len(masses):
        raise ValueError("Dynamical matrix shape is incompatible with number of atoms.")

    # Transform matrix
    N = len(masses)
    dyn = np.copy(dyn_mat)
    for n in range(N):
        for m in range(N):
            i = 3 * n
            j = 3 * m
            factor = np.sqrt(masses[n] * masses[m])
            dyn[i : i + 3, j : j + 3] /= factor

    return dyn * units


class Dyn:
    """
    A container for a single phonon dynamical matrix at a given q-point.

    This class provides functionality to read and manipulate dynamical matrices
    produced by Quantum ESPRESSO, convert them to real physical units, and diagonalize
    them to obtain phonon frequencies and polarization vectors.

    Attributes
    ----------
    q : pint.Quantity [2π/crystal]
        The q-point in reciprocal space (reduced coordinates).
    dyn : np.ndarray or Quantity
        The dynamical matrix, either in QE format (with mass prefactors) or physical units.
    Cell : cell.Cell
        The atomic structure associated with this q-point, including lattice, positions, elements.
    masses : Quantity [2*m_e]
        Array of atomic masses, usually in units of 2 electron masses as per QE convention.
    freqs : Quantity [cm⁻¹], optional
        Frequencies of the phonon modes after diagonalization.
    displacements : np.ndarray, optional
        Mass-normalized eigenvectors (displacements). Shape and units depend on output format.
    polarizations : np.ndarray, optional
        Raw polarization vectors before mass normalization.

    Methods
    -------
    from_file(...)
        Initialize the `Dyn` object from a Quantum ESPRESSO `.dyn*` file.
    diagonalize(...)
        Diagonalize a dynamical matrix and return phonon frequencies, polarization vectors and displacements.
    """

    def __init__(
        self,
        q: np.ndarray | ureg.Quantity = None,
        dyn: np.ndarray | ureg.Quantity = None,
        masses: np.ndarray | ureg.Quantity = None,
        Cell: cell.Cell = None,
    ):
        """
        Initialize the a `Dyn` object.

        Parameters
        ----------
        q : np.ndarray | ureg.Quantity
            The q-point in reciprocal space (reduced coordinates [2π/crystal]).
        dyn : np.ndarray | ureg.Quantity
            The dynamical matrix, either in QE format (with mass prefactors) or physical units.
        masses : np.ndarray | ureg.Quantity
            Array of atomic masses, usually in units of 2 electron masses as per QE convention.
        Cell : cell.Cell
            The atomic structure associated with this q-point, including lattice, positions, elements.
        """
        self.q = q
        self.dyn = dyn
        self.Cell = Cell
        self.masses = masses

    @classmethod
    def from_file(
        cls,
        q_cryst: np.ndarray | ureg.Quantity,
        results_ph_path: str,
        qe_format: bool = True,
    ):
        """
        Reads the Quantum ESPRESSO `.dyn*` file corresponding to a given q-point.

        This function locates the `.dyn*` file generated by `ph.x` that corresponds to a
        desired q-point (in reduced crystalline coordinates), extracts the corresponding dynamical
        matrix, and optionally converts it to the real physical dynamical matrix in units
        of 1 / [time]².

        Parameters
        ----------
        q_cryst : np.ndarray | ureg.Quantity
            The q-point of interest, expressed in reduced crystalline coordinates (fractions of reciprocal lattice vectors).
            If not a `Quantity`, it is assumed to be in `_2pi/crystal` units.

        results_ph_path : str
            Path to the directory containing the Quantum ESPRESSO `ph.x` output `.dyn*` files.

        qe_format : bool, optional
            If True (default), returns the raw QE dynamical matrix (includes sqrt(m_i m_j) mass factors).
            If False, converts the dynamical matrix to true physical form in 1 / [time]² units (Ry/h)^2.

        Returns
        -------
        Dyn
            A new Dynamical matrix instance.
        """
        system = grep.dyn_q(q_cryst, results_ph_path, qe_format)
        atoms = cell.Atoms(
            system.elements, system.positions.magnitude, cell=system.lattice.magnitude
        )
        return cls(
            q=system.q,
            dyn=system.dyn,
            Cell=cell.Cell(atoms=atoms),
            masses=system.masses,
        )

    def diagonalize(
        self,
        qe_format_in: bool = True,
        qe_format_out: bool = True,
    ):
        """
        Diagonalize a dynamical matrix and return phonon frequencies, polarization vectors and displacements.

        This function converts the input dynamical matrix to its physical form (if needed),
        diagonalizes it to obtain phonon frequencies (in cm⁻¹), and computes the corresponding
        normalized displacement vectors.

        Parameters
        ----------
        qe_format_in : bool, optional
            Whether `dyn_mat` is in QE format (default: True). If so, it is converted to the real
            physical format for diagonalization.

        qe_format_out : bool, optional
            Whether to return the normalized displacements and polarization vectors in QE format: an array of shape
            (3N, N_atoms, 3). If False, returns a plain (3N × 3N) array of normalized eigenvectors/displacements where
            polarizations[i] corresponds to eigenvalue freqs[i]**2.

        Returns
        -------
        New attributes to the Dyn object:
            - freqs : np.ndarray
                Array of 3N phonon frequencies in cm⁻¹. Real or negative imaginary parts indicate stable or unstable modes.
            - displacements : np.ndarray
                Normalized displacement vectors of the dynamical matrix. Format depends on `qe_format_out`.
            - polarizations : np.ndarray
                Raw polarization vectors (before mass-weighted normalization). Format depends on `qe_format_out`.

        Raises
        ------
        ValueError
            If eigenvalues are not purely real or imaginary (i.e., contain both parts).
        """
        # Check unit consistency
        dyn_mat, masses = self.dyn, self.masses
        ut._check_unit_consistency(
            [dyn_mat, masses], ["Dynamical matrix", "Atomic masses"]
        )

        # Strip units
        if isinstance(dyn_mat, ureg.Quantity):
            if dyn_mat.dimensionality == ureg("kg/s^2").dimensionality:
                dyn_mat = _QEdyn2Realdyn(dyn_mat, masses)
            dyn, units = dyn_mat.magnitude, (1 * np.sqrt(dyn_mat).units).to("c/cm")
            masses = masses.magnitude
        else:
            if qe_format_in:
                dyn_mat = _QEdyn2Realdyn(dyn_mat, masses)
            dyn, units = dyn_mat, 1

        # Diagonalize
        eigvals, pol = np.linalg.eig(dyn)

        # Convert eigenvalues (ω²) to frequencies in cm⁻¹
        freqs = np.zeros_like(eigvals, dtype=float)
        for i, val in enumerate(np.sqrt(eigvals)):
            if np.isclose(np.imag(val), 0, atol=1e-10):
                freqs[i] = np.real(val)
            elif np.isclose(np.real(val), 0, atol=1e-10):
                freqs[i] = -np.abs(np.imag(val))
            else:
                raise ValueError(
                    f"Complex eigenvalue with both real and imaginary part at index {i}: {val}"
                )

        # Normalize polarization vectors to obtain displacements
        dim = len(masses)
        displacements = np.copy(pol)
        for col in range(pol.shape[1]):
            for i in range(dim):
                displacements[3 * i : 3 * i + 3, col] /= np.sqrt(masses[i])
            displacements[:, col] /= np.linalg.norm(displacements[:, col])

        # Sort by increasing frequency
        order = np.argsort(freqs)
        freqs_sorted = freqs[order] * units
        displacements = displacements[:, order].T
        polarizations = pol[:, order].T

        if qe_format_out:
            displacements = displacements.reshape(len(masses) * 3, len(masses), 3)
            polarizations = polarizations.reshape(len(masses) * 3, len(masses), 3)

        self.freqs = freqs_sorted
        self.displacements = displacements
        self.polarizations = polarizations


def _find_supercell(
    q_cryst: np.ndarray | ureg.Quantity | list[np.ndarray] | list[ureg.Quantity],
) -> SimpleNamespace:
    """
    Compute the minimal supercell that is commensurate with a set of q-points in crystal units.

    This function finds the smallest supercell that accommodates the periodicity
    of the input q-points by ensuring that the phase factors `exp(2πi q·R)` close
    after a finite number of lattice translations. It also computes the corresponding
    phase factor grids over the supercell.

    Parameters
    ----------
    q_cryst : np.ndarray | ureg.Quantity | list[np.ndarray | ureg.Quantity]
        A single q-point or list of q-points, given in crystalline units
        (units of 2π/crystal). If provided with units, they must be `"_2pi/crystal"`.

    Returns
    -------
    SimpleNamespace
        A namespace with the following attributes:
        - size : np.ndarray
            A 3-element integer array defining the minimal commensurate supercell.
        - phase_factors : list[np.ndarray]
            List of 3D complex-valued arrays with the phase factors
            `exp(2πi q·R)` evaluated on the supercell grid for each q-point.
            Each has dimnesion equal to the supercell size [phase_factors.shape = size].
        - q : list[ureg.Quantity]
            The q-points, returned as quantities with units `_2pi/crystal`.

    Raises
    ------
    ValueError
        If q-points do not have units of 2π/crystal, or if no commensurate supercell
        is found in a direction within 100 lattice multiples.
    """
    from math import gcd

    # Normalize input to list of np.ndarray
    if len(np.shape(q_cryst)) == 1:
        q_cryst = [q_cryst]

    # Check units and convert to magnitudes
    ut._check_unit_consistency(q_cryst)
    if isinstance(q_cryst[0], ureg.Quantity):
        for q in q_cryst:
            if q.units != ureg("_2pi/crystal"):
                raise ValueError("Each q-point must have units of 2π/crystal.")
        q_cryst_mag = [q.magnitude for q in q_cryst]
    else:
        q_cryst_mag = q_cryst

    supersize = np.ones(3, dtype=int)

    for q in q_cryst_mag:
        p_super = np.zeros(3, dtype=int)
        for i, n in enumerate(q):
            for cell in range(1, 101):
                if (cell * n) % 1 == 0:
                    p_super[i] = cell
                    break
            else:
                raise ValueError(
                    f"Could not find commensurate cell for q = {q} in direction {i} "
                    "within 100 multiples. Check that q is rational."
                )
        # LCM between existing and partial supercell
        for i in range(3):
            supersize[i] = supersize[i] * p_super[i] // gcd(supersize[i], p_super[i])

    # Compute phase factors over supercell grid
    phase_factors = []
    for q in q_cryst_mag:
        phases_q = np.empty(supersize, dtype=complex)
        for i in range(supersize[0]):
            for j in range(supersize[1]):
                for k in range(supersize[2]):
                    R = np.array([i, j, k])
                    phases_q[i, j, k] = np.exp(2j * np.pi * np.dot(q, R))
        phase_factors.append(phases_q)

    q_cryst_with_units = [np.array(q) * ureg("_2pi/crystal") for q in q_cryst_mag]

    return SimpleNamespace(
        size=supersize,
        phase_factors=phase_factors,
        q=q_cryst_with_units,
    )


class CDW:
    """
    Class representing a charge-density wave (CDW) structure built from a set
    of phonon dynamical matrices at different q-points.

    This class:
    - Stores the dynamical matrices.
    - Determines the commensurate supercell that accommodates the q-points.
    - Builds the corresponding supercell.
    - Applies distortions associated with the soft phonon modes.

    Attributes
    ----------
    dyn_matrices : list
        List of `Dyn` objects, each containing a dynamical matrix and metadata.
    Cell : Cell
        The primitive unit cell object, taken from the first `Dyn` entry.
    q : list[np.ndarray]
        List of q-points (in 2π/crystal units) corresponding to the dynamical matrices.
    SuperCell : Cell
        The supercell constructed from the primitive cell and commensurate size.
    masses : Quantity [2*m_e]
        Array of atomic masses, usually in units of 2 electron masses as per QE convention.
        This are read from the first dynamical matrix.
    _supercell : SimpleNamespace
        Contains the commensurate supercell size and phase factors.

    Methods
    -------
    from_files(...)
        Initialize a `CDW` object by reading .dyn files from a directory.
    distort_crystal(...)
        Apply a charge density wave distortion to the supercell.
    sym_analysis(...)
        Perform a symmetry analysis of commensurate CDW distortions across a grid of order parameters.
    """

    def __init__(self, dyn_matrices: list):
        """
        Initialize the `CDW` object from a list of `Dyn` objects.

        Parameters
        ----------
        dyn_matrices : list of Dyn
            List of dynamical matrices at different q-points.

        Raises
        ------
        ValueError
            If q-points do not have units of 2π/crystal.
        """
        self.dyn_matrices = dyn_matrices
        self.Cell = self.dyn_matrices[0].Cell

        # Diagonalize and extract q-points
        q_cryst = []
        for dyn in self.dyn_matrices:
            dyn.diagonalize()
            q_cryst.append(dyn.q)
        ut._check_unit_consistency(q_cryst)
        if isinstance(q_cryst[0], ureg.Quantity):
            for q in q_cryst:
                if q.units != ureg("_2pi/crystal"):
                    raise ValueError("Each q-point must have units of 2π/crystal.")
            q_cryst_mag = [q.magnitude for q in q_cryst]
        else:
            q_cryst_mag = q_cryst
        self.q = q_cryst_mag * ureg("_2pi/crystal")

        # Determine supercell and construct it
        self._supercell = _find_supercell(self.q)
        self.SuperCell = self.Cell.get_supercell(self._supercell.size)

    @classmethod
    def from_files(cls, q_cryst, results_ph_path):
        """
        Construct a CDW object by reading .dyn files from a directory.

        Parameters
        ----------
        q_cryst : np.ndarray or list[np.ndarray]
            One or more q-points to read (in 2π/crystal units).
        results_ph_path : str
            Path to folder containing QE phonon outputs (.dyn*).

        Returns
        -------
        CDW
            The constructed CDW object.
        """
        if len(np.shape(q_cryst)) == 1:
            q_cryst = [q_cryst]
        dyn_matrices = [Dyn.from_file(q, results_ph_path) for q in q_cryst]
        return cls(dyn_matrices)

    @property
    def masses(self):
        """
        Property that returns the masses of the first stored dynamical matrix.
        Masses should be the same for all matrices.
        """
        return self.dyn_matrices[0].masses

    def distort_crystal(
        self,
        amplitudes: list[ureg.Quantity] | list[complex] = defaults.CDW_amplitude,
        modes: list[int] = None,
    ):
        """
        Apply a charge density wave distortion to the supercell.

        Parameters
        ----------
        amplitudes : list[ureg.Quantity] | list[complex], optional
            List of complex amplitudes (including phase) for each q-mode. Default being
            a value of 0.5 Å for each q-point (np.ones(len(self.q)))*0.5 Å.
        modes : list[int], optional
            List of mode indices to be activated at each q-point. Default being the
            lowest energy mode for each q point.

        Returns
        -------
        Cell
            A new `Cell` object representing the distorted supercell.

        Notes
        -----
        Final displacement is found by applying the given amplitudes to the individual
        normalized displacements:
            `Real( Σ (amplitudes * displacement))`, where `norm (displacement) = 1`
            If no units are defined the amplitudes are considered to be in angstroms.
        """
        _formatted = _format_input(self, amplitudes=amplitudes, modes=modes)
        amplitudes, modes = _formatted.amplitudes, _formatted.modes

        supercell_ind = 0
        cell_size = len(self.Cell[1])
        # Positions are given in angstroms
        positions = np.copy(self.SuperCell.atoms.positions)
        amplitudes = np.array(amplitudes)

        for i in range(self._supercell.size[0]):
            for j in range(self._supercell.size[1]):
                for k in range(self._supercell.size[2]):
                    for q_ind in range(len(self.q)):
                        start = supercell_ind * cell_size
                        end = (supercell_ind + 1) * cell_size
                        phase = self._supercell.phase_factors[q_ind][i, j, k]
                        displacements = self.dyn_matrices[q_ind].displacements[
                            modes[q_ind]
                        ]
                        positions[start:end] += np.real(
                            displacements * phase * amplitudes[q_ind]
                        )
                    supercell_ind += 1
        # Construct new Atoms object and return as Cell
        distorted = cell.Atoms(
            self.SuperCell.atoms.numbers, positions, cell=self.SuperCell.atoms.cell
        )
        return cell.Cell(atoms=distorted)

    def sym_analysis(
        self,
        grid: list[int] = None,
        modes: list[int] = None,
        amplitudes: complex | list[complex] = defaults.CDW_amplitude,
        symprec: float = defaults.symprec,
    ):
        """
        Perform a symmetry analysis of commensurate CDW distortions across a grid of order parameters.

        This function generates all symmetry-inequivalent distortions resulting from a combination
        of phonon modes and computes the resulting space groups after displacing atoms along the
        corresponding eigenvectors with various amplitudes.

        Parameters
        ----------
        grid : list[int], optional
            Grid resolution for scanning order parameters in each mode direction. Each entry N_i sets
            the number of points in [-1, 1] for the i-th mode. Defaults to [3, 3, ..., 3].

        modes : list[int], optional
            List of mode indices (starting at 0) to be used for the displacement. If None, the lowest
            frequency modes will be selected.

        amplitudes : list[ureg.Quantity] | list[complex], optional
            Amplitud factors applied to each normalized distortion vector. Final displacement is given by:
            `Real( Σ (grid_point * amplitudes * displacement))`, where `norm(displacement) = 1`.
            Default is 0.5 Å for each mode.

        symprec : float, optional
            Tolerance used by spglib to determine space group symmetry.
            Default is `yaiv.defaults.config.defaults.symprec`.

        Returns
        -------
        SimpleNamespace
            Contains:
            - amplitudes : ureg.Quantity[np.ndarray]
                Array of complex vectors defining the grid of amplitude values, [lenght] units.
            - order_parameters : ureg.Quantity[np.ndarray]
                Array of complex vectors defining the grid of order_parameter values [lenght * sqrt(mass)] units.
            - space_groups : list[str]
                List of space group symbols (e.g. 'P1', 'R-3m') returned by spglib.
            - cell_sizes : np.ndarray
                Ratio of the number of atoms in the distorted supercell to the original cell.
            - structures : list[cell.Cell]
                List of distorted crystal structures.

        Notes
        -----
        - For even-sized grids, the origin (0) will not be included.
        - This function relies on `spglib` to determine space groups and primitive cells.
        """

        _formatted = _format_input(self, grid=grid, modes=modes, amplitudes=amplitudes)
        grid, modes, max_amplitudes = (
            _formatted.grid,
            _formatted.modes,
            _formatted.amplitudes,
        )

        # Generate amplitudes
        amplitudes = ut.grid_generator(grid) * max_amplitudes
        num_in_atoms = len(self.Cell[1])
        SGs, sizes, structures = [], [], []

        # Distort
        for amp in amplitudes:
            distorted = self.distort_crystal(amplitudes=amp, modes=modes)
            SG = spg.get_spacegroup(distorted, symprec=symprec)
            num_out_atoms = len(spg.find_primitive(distorted, symprec=symprec)[1])
            s = int(num_out_atoms / num_in_atoms)
            SGs.append(SG)
            sizes.append(s)
            structures.append(distorted)

        # Obtain order parameters
        amplitudes *= ureg.ang
        displacements = []
        for i in range(len(self.q)):
            displacements.append(self.dyn_matrices[i].displacements[modes[i]])

        order_parameters = ut.amplitude2order_parameter(
            amplitudes=amplitudes,
            masses=self.masses,
            displacements=displacements,
        )

        return SimpleNamespace(
            amplitudes=amplitudes,
            order_parameters=order_parameters,
            space_groups=SGs,
            cell_sizes=np.array(sizes),
            structures=structures,
        )


def _format_input(
    CDW: CDW,
    amplitudes: list[complex] = None,
    modes: list[int] = None,
    grid: list[int] = None,
):
    """
    Formats and normalizes user input for a CDW distortion calculation.

    Ensures that `amplitudes`, `modes`, and `grid` are provided in consistent list formats
    and that amplitudes are in angstroms.
    Defaults are assigned based on the number of q-vectors in the CDW object.

    Parameters
    ----------
    CDW : CDW
        CDW object containing the q-vectors and structure information.
    amplitudes : list[complex] or scalar, optional
        Complex amplitudes for each mode. Defaults to `[0.5, ..., 0.5] Å` (real).
    modes : list[int] or int, optional
        List of phonon mode indices corresponding to each q-vector. Defaults to lowest-frequency modes (0).
        If a single integer is given, it is wrapped in a list.
    grid : list[int] or int, optional
        Number of points per direction in the grid of order parameters. Defaults to `[3, ..., 3]`.
        If a scalar is given, it is broadcasted.

    Returns
    -------
    SimpleNamespace
        A namespace with normalized fields:
        - amplitudes : list[complex]
        - modes : list[int]
        - grid : list[int]
        - global_amplitude : list[complex]
    """
    DIM = len(CDW.q)

    # Order parameter (convert to angstrom and stripe units)
    if amplitudes is None:
        amplitudes = np.ones(DIM) * defaults.CDW_amplitude
    elif isinstance(amplitudes, ureg.Quantity):
        amplitudes = amplitudes.to("ang").magnitude
    if isinstance(amplitudes, (complex, float, int)):
        amplitudes = [amplitudes] * DIM

    # Modes
    if modes is None:
        modes = np.zeros(DIM, dtype=int)
    elif isinstance(modes, int):
        modes = [modes]

    # Grid
    if grid is None:
        grid = np.ones(DIM, dtype=int) * 3
    elif isinstance(grid, int):
        grid = [grid]

    return SimpleNamespace(amplitudes=amplitudes, modes=modes, grid=grid)


class BOES:
    """
    Class for constructing a Born-Oppenheimer energy surface (BOES) associated with charge-density wave (CDW) distortions.

    This class uses a CDW object to generate symmetry-inequivalent crystal structures
    by displacing atoms along selected phonon eigenmodes. It supports scanning multidimensional
    order parameter grids or paths in the distortion space and exporting input files for QE calculations.

    Attributes
    ----------
    CDW : CDW
        A `CDW` object containing dynamical matrices and crystal structure.
    structures : list[cell.Cell]
        Distorted structures generated by grid or line scans.
    amplitudes : ureg.Quantity[np.ndarray]
        Array of complex vectors defining the amplitude values applied to normalized displacements, [lenght] units.
    order_parameters : ureg.Quantity[np.ndarray]
        Array of complex vectors defining the order_parameter values [lenght * sqrt(mass)] units.
    modes : list[int]
        Indices of phonon modes used in the distortions.
    space_groups : list[str]
        List of space group symbols (e.g. 'P1', 'R-3m') returned by spglib.
    energies : SimpleNamespace | list[float]
        List of BOES energies, either total Free energies or a decomposition of the different terms.

    Methods
    -------
    from_file(...)
        Initialize the `BOES` object from a `.pkl` file.
    save_as(...)
        Save the entire `BOES` object, including the `CDW` instance and all structures in a `.pkl` file.
    generate_structures_grid(...)
        Generate distorted structures over a regular grid in order parameter space.
    generate_structures_line(...)
        Generate distorted structures along a 1D path in the amplitude (order parameter) space.
    save_jobs_pwi(...)
        Export all generated distorted structures as QE input `.pwi` files.
    read_energies_pwo(...)
        Grep total energies from Quantum ESPRESSO output `.pwo` files for each distorted structure.
    """

    def __init__(self, CDW: CDW = None):
        """
        Initialize a `BOES` object from a `CDW` instance.

        Parameters
        ----------
        CDW : CDW
            The `CDW` object defining the parent structure and q-point modes.
        """
        self.CDW = CDW
        # Initialize to None the rest of attributes
        self.structures = self.amplitudes = self.order_parameters = self.modes = (
            self.space_groups
        ) = self.energies = None

    @classmethod
    def from_file(cls, filename: str):
        """
        Load a BOES object from a .pkl file.

        Parameters
        ----------
        filename : str
            Path to file created with `save_as`.

        Returns
        -------
        BOES
            Fully reconstructed BOES object.
        """
        import pickle

        with open(filename, "rb") as f:
            obj = pickle.load(f)
        return obj

    def save_as(self, filename: str = "BOES.pkl"):
        """
        Save the entire BOES object, including the CDW instance and all structures in a .pkl file.

        Parameters
        ----------
        filename : str, Optional
            Path to the output file. '.pkl' will be added if missing. Default `BOES.pkl`.
        """
        import pickle

        if not filename.endswith(".pkl"):
            filename += ".pkl"
        with open(filename, "wb") as f:
            pickle.dump(self, f)

    def generate_structures_grid(
        self,
        grid: list[int] = None,
        modes: list[int] = None,
        amplitudes: complex | list[complex] = defaults.CDW_amplitude,
        symprec: float = defaults.symprec,
    ):
        """
        Generate distorted structures over a regular grid in order parameter space.

        Parameters
        ----------
        grid : list[int], optional
            Grid resolution for scanning order parameters in each mode direction. Each entry N_i sets
            the number of points in [-1, 1] for the i-th mode. Defaults to [3, 3, ..., 3].

        modes : list[int], optional
            List of mode indices (starting at 0) to be used for the displacement. If None, the lowest
            frequency modes will be selected.

        amplitudes : list[ureg.Quantity] | list[complex], optional
            Amplitud factors applied to each normalized distortion vector. Final displacement is given by:
            `Real( Σ (grid_point * amplitudes * displacement))`, where `norm(displacement) = 1`.
            Default is 0.5 Å for each mode.

        symprec : float, optional
            Tolerance used by spglib to determine space group symmetry.
            Default is `yaiv.defaults.config.defaults.symprec`.
        """
        _formatted = _format_input(
            self.CDW, grid=grid, modes=modes, amplitudes=amplitudes
        )
        grid, modes, amplitudes = (
            _formatted.grid,
            _formatted.modes,
            _formatted.amplitudes,
        )

        print(
            "Amplitudes in a",
            grid,
            "grid  with max amplitudes of",
            amplitudes,
            "Å, for modes",
            modes,
            ", in q points:",
        )
        print(self.CDW.q)
        print()
        print("The general displacement at each config is given by:")
        # print("Real ( Σ (OP * amplitudes * displacements)) * Å")
        print("Real( Σ (amplitudes * displacement)) Å`, where `norm(displacement) = 1`")
        print()
        print("---------------------------------------------------")
        print(
            'Final applied "maximum" displacement (Å): (Σ amplitudes * displacements ) ='
        )
        distorted = self.CDW.distort_crystal(amplitudes=amplitudes, modes=modes)
        displacement = (
            self.CDW.Cell.atoms.positions
            - distorted.atoms.positions[: len(self.CDW.Cell[1])]
        )
        print(np.around(displacement, decimals=5))

        out = self.CDW.sym_analysis(
            grid=grid, modes=modes, amplitudes=amplitudes, symprec=symprec
        )
        self.structures = out.structures
        self.amplitudes = out.amplitudes
        self.modes = modes
        self.space_groups = out.space_groups
        self.order_parameters = out.order_parameters

    def generate_structures_line(
        self,
        amplitude_i: list[int] = -defaults.CDW_amplitude,
        amplitude_f: list[int] = defaults.CDW_amplitude,
        modes: list[int] = None,
        steps: int = 31,
        symprec: float = defaults.symprec,
    ):
        """
        Generate distorted structures along a 1D path in the amplitude (order parameter) space.

        Parameters
        ----------
        amplitude_i : list[ureg.Quantity] | list[complex]
            Initial point in the amplitude space. Default units are Å.
        amplitude_f : list[ureg.Quantity] | list[complex]
            Final point in the amplitude space. Default units are Å.
        modes : list[int]
            Mode indices used to generate the displacements.
        steps : int
            Number of interpolation steps between initial and final distortion.
        symprec : float, optional
            Tolerance used by spglib to determine space group symmetry.
            Default is `yaiv.defaults.config.defaults.symprec`.

        Notes
        -----
        Final displacement is given by:
            `Real( Σ amplitudes * displacement)`, where `norm(displacement) = 1`.
        """
        modes = _format_input(self.CDW, modes=modes).modes
        amplitude_i = _format_input(self.CDW, amplitudes=amplitude_i).amplitudes
        amplitude_f = _format_input(self.CDW, amplitudes=amplitude_f).amplitudes

        print(
            "Amplitude varying from",
            amplitude_i,
            "Å to",
            amplitude_f,
            "Å in",
            steps,
            "steps, for modes",
            modes,
            ", in q points:",
        )
        print(self.CDW.q)
        print()
        print("The general displacement at each config is given by:")
        print("Real ( Σ amplitudes * displacements) Å")
        print()
        print("---------------------------------------------------")
        print(
            'Final applied "maximum" displacement (Å): (Σ max(amplitude) * displacements ) ='
        )
        if np.linalg.norm(amplitude_f) >= np.linalg.norm(amplitude_i):
            amplitude_max = amplitude_f
        else:
            amplitude_max = amplitude_i

        distorted = self.CDW.distort_crystal(amplitudes=amplitude_max, modes=modes)
        displacement = (
            self.CDW.Cell.atoms.positions
            - distorted.atoms.positions[: len(self.CDW.Cell[1])]
        )
        print(np.around(displacement, decimals=5))
        amplitudes = np.linspace(amplitude_i, amplitude_f, steps)
        structures, SGs = [], []
        for amp in amplitudes:
            distorted = self.CDW.distort_crystal(amplitudes=amp, modes=modes)
            SG = spg.get_spacegroup(distorted, symprec=symprec)
            SGs.append(SG)
            structures.append(distorted)

        # Obtain order parameters
        amplitudes *= ureg.ang
        displacements = []
        for i in range(len(self.CDW.q)):
            displacements.append(self.CDW.dyn_matrices[i].displacements[modes[i]])

        order_parameters = ut.amplitude2order_parameter(
            amplitudes=amplitudes,
            masses=self.CDW.masses,
            displacements=displacements,
        )

        self.structures = structures
        self.amplitudes = amplitudes
        self.modes = modes
        self.space_groups = SGs
        self.order_parameters = order_parameters

    def save_jobs_pwi(
        self,
        dest_folder: str,
        template: str = None,
        kgrid: tuple | list = None,
        primitive: bool = False,
        automatic_kgrid: bool = False,
        symprec: float = defaults.symprec,
    ):
        """
        Export all generated distorted structures as QE input files (.pwi).

        Parameters
        ----------
        dest_folder : str
            Path to the folder where input files will be saved.
        template : str
            Path to a QE input file to use as a template.
        kgrid : list, optional
            Desiered number of kgrid [N1,N2,N3].
            Defaults to the template (if given) or `qe_defaults`.
        primitive : bool, optional
            Whether to convert distorted structures to primitive cells.
            Defaults to False.
        automatic_kgrid : bool, optional
            If True, it inversely scales the k-point grid with the lattice. The scale is
            computed from the volume of the undistorted crystal to the provided kgrid.
            Defaults to False.
        symprec : float
            Symmetry tolerance used by spglib to define the primitive cell.
            Default is `yaiv.defaults.config.defaults.symprec`.
        """
        import os
        from yaiv.defaults.config import qe_defaults

        if not os.path.exists(dest_folder):
            os.makedirs(dest_folder)
            print("Folder created...")

        if kgrid is None:
            kgrid = qe_defaults.kpts

        print("Generating input files...")
        # Get the scaling factor scale = kpoints * volume
        if automatic_kgrid:
            KPPRA = np.prod(kgrid) * len(self.CDW.Cell.spglib[1])
        for i, structure in enumerate(self.structures):
            if primitive:
                structure.spglib = spg.find_primitive(structure, symprec=symprec)
                structure.atoms = cell.spglib2ase(structure.spglib)
            if automatic_kgrid:
                kgrid = ut.auto_kgrid(
                    lattice=structure.spglib[0],
                    n_atoms=len(structure.spglib[1]),
                    kppra=KPPRA,
                )
            filename = dest_folder + "/" + str(i) + ".pwi"
            structure.write_espresso_in(filename, template=template, kgrid=kgrid)
        print("Done.")

    def read_energies_pwo(self, dest_folder: str, decomposition: bool = True):
        """
        Grep total energies from Quantum ESPRESSO output files (.pwo) for each distorted structure.

        This method reads energy values from files named as `0.pwo`, `1.pwo`, ..., located in
        the specified folder, and stores the resulting list in `self.energies`.

        Parameters
        ----------
        dest_folder : str
            Path to the folder containing the .pwo output files corresponding to each distorted structure.
        decomposition : bool, Optional
            If True, extract both total Free energy and component energies (e.g., electronic, ion-ion).
            If False, only the total Free energy energy is returned. Default is True.
        """
        energies = []
        for i in range(len(self.structures)):
            filename = dest_folder + "/" + str(i) + ".pwo"
            energies.append(grep.total_energy(filename, decomposition=decomposition))

        if isinstance(energies[0], SimpleNamespace):
            units = energies[0].F.units
            data = {}
            for atr in energies[0].__dict__:
                data = dict(
                    data,
                    **{
                        atr: [X.__getattribute__(atr).magnitude for X in energies]
                        * units
                    },
                )
            energies = SimpleNamespace(**data)
        else:
            energies = [X.magnitude for X in energies] * energies[0].units
        self.energies = energies
