"""Define schemas for DFPT, phonopy, and pheasy-derived phonon data."""

from __future__ import annotations

from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Annotated

import numpy as np
import yaml  # type: ignore[import-untyped]
from monty.io import zopen
from monty.os.path import zpath
from pydantic import BaseModel, BeforeValidator, Field, PlainSerializer, PrivateAttr
from typing_extensions import Literal, TypedDict

from emmet.core.band_theory import BandStructure, BandTheoryBase
from emmet.core.base import CalcMeta
from emmet.core.io.pymatgen import (
    BandStructureSymmLine,
    CompletePhononDos,
    Kpoint,
    Lattice,
    PhononBandStructureSymmLine,
)
from emmet.core.io.pymatgen import PhononDos as PhononDosObject
from emmet.core.io.pymatgen import Spin, Structure
from emmet.core.math import Matrix3D, Tensor4R, Vector3D
from emmet.core.polar import BornEffectiveCharges, DielectricDoc, IRDielectric
from emmet.core.structure import StructureMetadata
from emmet.core.types.enums import DocEnum, ValueEnum
from emmet.core.types.pymatgen_types.structure_adapter import StructureType
from emmet.core.types.typing import DateTimeType, FSPathType, IdentifierType
from emmet.core.utils import get_num_formula_units, type_override

if TYPE_CHECKING:
    from collections.abc import Sequence
    from typing import Any

    from typing_extensions import Self

DEFAULT_PHONON_FILES = {
    "structure": "POSCAR",
    "phonon_bandstructure": "band.yaml",  # chosen as Phonopy default
    "phonon_dos": "total_dos.dat",  # chosen as Phonopy default
    "projected_dos": "projected_dos.dat",  # chosen as Phonopy default
    "force_constants": "FORCE_CONSTANTS",  # chosen as Phonopy default
    "born": "born.npz",
    "epsilon_static": "epsilon_static.npz",
    "epsilon_static_and_born": "BORN",
    "phonopy_output": "phonopy.yaml",
}


class PhononMethod(ValueEnum):
    """Define common methods for computed phonon properties."""

    DFPT = "dfpt"
    PHONOPY = "phonopy"
    PHEASY = "pheasy"


class PhononDOS(BandTheoryBase):
    """Define schema of pymatgen phonon density of states."""

    frequencies: list[float] = Field(description="The phonon frequencies in THz.")
    densities: list[float] = Field(description="The phonon density of states.")
    projected_densities: list[list[float]] | None = Field(
        None, description="The projected phonon density of states."
    )

    @classmethod
    def from_pmg(cls, config: PhononDosObject | CompletePhononDos | dict) -> Any:
        """Correctly parse pymatgen-like objects."""
        remap = {
            "structure": "structure",
            "frequencies": "frequencies",
            "densities": "densities",
            "projected_densities": "pdos",
        }
        if isinstance(config, PhononDosObject | CompletePhononDos):
            cls_config = {k: getattr(config, v, None) for k, v in remap.items()}
            if isinstance(pdos_as_dict := cls_config.get("projected_densities"), dict):
                if not cls_config["structure"]:
                    raise ValueError(
                        "Cannot parse atom-projected phonon density of states without a structure."
                    )
                cls_config["projected_densities"] = [
                    list(pdos_as_dict.get(site, [])) for site in cls_config["structure"]
                ]
            return cls(**cls_config)
        return cls(**{k: config.get(v) for k, v in remap.items()})

    @cached_property
    def to_pmg(self) -> PhononDosObject | CompletePhononDos:
        """Get / cache corresponding pymatgen object."""
        dos = PhononDosObject(frequencies=self.frequencies, densities=self.densities)
        if self.projected_densities and self.structure:
            return CompletePhononDos(
                self.structure,
                dos,
                dict(zip(self.structure, self.projected_densities, strict=True)),
            )
        return dos

    @classmethod
    def from_phonopy(
        cls,
        phonon_dos_file: FSPathType,
        structure_file: FSPathType | None = None,
        projected_dos_file: FSPathType | None = None,
        **kwargs,
    ) -> Self:
        """Create a PhononDOS from phonopy .dat output.

        Ref. for data formats used by phonopy:
        - total_dos.dat :
        https://phonopy.github.io/phonopy/output-files.html#total-dos-dat-and-projected-dos-dat
        - projected_dos.dat :
        https://phonopy.github.io/phonopy/output-files.html#file-format-of-projected-dos-dat

        Parameters
        -----------
        phonon_dos_file (FSPathType) : path to total_dos.dat
        structure_file (FSPathType or None) : path to a POSCAR-like object,
            either the base computational cell (POSCAR) or the supercell
            used in the phonon calculation (SPOSCAR).
        projected_dos_file (FSPathType or None) : path to projected_dos.dat
        **kwargs : any other kwargs to pass into the PhononDOS constructor
        """
        phonopy_dos: dict[str, Any] = {
            k: []
            for k in (
                "frequencies",
                "densities",
            )
        }
        # can handle gzipped files
        total_dos = np.loadtxt(phonon_dos_file)
        phonopy_dos["frequencies"] = total_dos[:, 0].tolist()
        phonopy_dos["densities"] = total_dos[:, 1].tolist()

        if structure_file:
            phonopy_dos["structure"] = Structure.from_file(structure_file)

        if projected_dos_file:
            # First column are frequencies, can skip these in parsing since
            # they should be the same as in the total DOS.
            proj_dos = np.loadtxt(projected_dos_file)
            phonopy_dos["projected_densities"] = [
                proj_dos[:, i].tolist() for i in range(1, proj_dos.shape[1])
            ]

        return cls(**phonopy_dos, **kwargs)


class ShreddedEigendisplacements(TypedDict):
    real: list[list[list[tuple[float, float, float]]]]
    imag: list[list[list[tuple[float, float, float]]]]


_EIGENMODE_CTYPE = list[list[list[tuple[complex, complex, complex]]]]
_EIGENMODE_RTYPE = list[list[list[tuple[float, float, float]]]]


def _ser_eigenmode(eigenmode: _EIGENMODE_CTYPE) -> dict[str, _EIGENMODE_RTYPE] | None:
    if not eigenmode:
        return None

    eigv = np.array(eigenmode)
    return {k: getattr(eigv, k).tolist() for k in ("real", "imag")}


def _deser_eigenmode(
    dct: _EIGENMODE_CTYPE | dict[str, _EIGENMODE_RTYPE],
) -> _EIGENMODE_CTYPE | None:
    if isinstance(dct, dict):
        try:
            real = np.array(dct["real"])
            imag = np.array(dct["imag"])
            return real + 1.0j * imag  # type: ignore[return-value]
        except Exception:
            return None

    return dct


def _recover_labels_from_segments(
    phonopy_bandstructure: dict,
    structure_file: FSPathType | None,
    path_convention: str,
    tol: float = 1e-3,
) -> dict[str, list[float]]:
    """Reconstruct ``labels_dict`` for a phonopy band.yaml that lacks per-entry
    ``label`` keys by matching segment-boundary q-positions to the canonical
    high-symmetry kpoint labels of the requested path convention.

    Matching tries three transformations to find an equivalence between each
    boundary q-position and a canonical kpoint: identity, integer
    reciprocal-lattice translations, and the structure's point group
    operations applied to the boundary q-position. This is required because
    band.yaml q-positions are often related to canonical labels by
    point-group symmetry, not only by lattice translations.

    When two boundaries match the same canonical label at different
    q-positions (i.e. visit the same physical high-symmetry point along
    different symmetry-equivalent path segments), later occurrences are
    suffixed with primes (e.g. ``"X"`` then ``"X'"`` then ``"X''"``) so
    every boundary in the band.yaml receives a tick on the plot.

    Args:
    phonopy_bandstructure: parsed band.yaml dict; must contain ``phonon``,
        ``segment_nqpoint``, and either an embedded structure (``lattice`` +
        ``points``) or a separate ``structure_file``.
    structure_file: optional path to a structure file used in preference to
        the band.yaml's embedded structure data.
    path_convention: kpath convention name passed to
        ``pymatgen.symmetry.bandstructure.HighSymmKpath`` (e.g.
        ``"latimer_munro"``, ``"setyawan_curtarolo"``, ``"hinuma"``).
    tol: tolerance for matching a boundary q-position to a canonical kpoint
        (in fractional reciprocal-lattice coordinates).

    Returns:
        dict[str, list[float]]: ``{label: q-position}`` for every boundary
        that found a match within ``tol``. May be empty if no boundaries
        match. Errors from ``HighSymmKpath`` (e.g. missing ``seekpath``
        dependency for the ``"hinuma"`` convention) propagate.
    """
    from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
    from pymatgen.symmetry.bandstructure import HighSymmKpath

    phonon_entries = phonopy_bandstructure["phonon"]
    if not phonon_entries:
        return {}

    segment_nqpoint = phonopy_bandstructure.get("segment_nqpoint")
    if not segment_nqpoint:
        return {}

    if structure_file:
        struct = Structure.from_file(structure_file)
    else:
        latt_raw = phonopy_bandstructure.get("lattice")
        pts = phonopy_bandstructure.get("points")
        if latt_raw is None or not pts:
            return {}
        struct = Structure(
            lattice=Lattice(np.asarray(latt_raw)),
            species=[p["symbol"] for p in pts],
            coords=[p["coordinates"] for p in pts],
            coords_are_cartesian=False,
        )

    # Walk segment boundaries (first q-point + last q-point of each segment).
    boundary_idxs = [0]
    running = 0
    for s in segment_nqpoint:
        running += s
        boundary_idxs.append(running - 1)

    hsk = HighSymmKpath(struct, path_type=path_convention)
    canonical = hsk.kpath.get("kpoints", {})  # {label: frac_coords}
    if not canonical:
        return {}

    # Consider unit-cell translations when matching: phonopy may emit a
    # q-position related to a canonical kpoint by an integer lattice shift.
    shifts = [
        np.array(s)
        for s in (
            (0, 0, 0),
            (1, 0, 0),
            (-1, 0, 0),
            (0, 1, 0),
            (0, -1, 0),
            (0, 0, 1),
            (0, 0, -1),
        )
    ]

    # Point-group operations on fractional reciprocal coordinates. Boundary
    # q-positions in band.yaml are frequently equivalent to canonical labels
    # only after applying one of these (e.g. rhombohedral cells where the
    # path traverses symmetry-related copies of the same high-symmetry
    # point). Falls back to identity-only matching if the SpacegroupAnalyzer
    # fails for any reason (unusual cells, malformed structure, etc.).
    try:
        sym_ops = SpacegroupAnalyzer(struct).get_point_group_operations(cartesian=False)
    except Exception:
        sym_ops = []

    labels_dict: dict[str, list[float]] = {}
    for idx in boundary_idxs:
        q = np.asarray(phonon_entries[idx]["q-position"])
        best_label = None
        best_dist = float("inf")
        for label, kc in canonical.items():
            kc_arr = np.asarray(kc)
            # Try identity first, then each symmetry operation.
            for op in (None, *sym_ops):
                q_eff = q if op is None else op.operate(q)
                for shift in shifts:
                    d = float(np.linalg.norm(q_eff - (kc_arr + shift)))
                    if d < best_dist:
                        best_dist = d
                        best_label = label
        if best_label is None or best_dist > tol:
            continue

        # Disambiguate duplicate labels that come from symmetry-equivalent
        # but distinct path points. The first occurrence keeps the canonical
        # label; later occurrences are suffixed with primes.
        candidate = best_label
        suffix = 0
        while candidate in labels_dict and not np.allclose(
            labels_dict[candidate], q, atol=tol
        ):
            suffix += 1
            candidate = best_label + "'" * suffix
        labels_dict.setdefault(candidate, q.tolist())

    return labels_dict


@type_override({"eigendisplacements": ShreddedEigendisplacements})
class PhononBS(BandStructure):
    """Define schema of pymatgen phonon band structure."""

    has_nac: bool = Field(
        False,
        description="Whether the calculation includes non-analytical corrections at Gamma.",
    )

    frequencies: list[list[float]] = Field(
        description="The eigen-frequencies, with the first index representing the band, and the second the k-point.",
        validation_alias="bands",
    )

    eigendisplacements: Annotated[
        _EIGENMODE_CTYPE | None,
        BeforeValidator(_deser_eigenmode),
        PlainSerializer(_ser_eigenmode),
    ] = Field(None, description="Phonon eigendisplacements in Cartesian coordinates.")

    _primitive_structure: Structure | None = PrivateAttr(None)

    @classmethod
    def from_pmg(cls, config: PhononBandStructureSymmLine | dict) -> Any:
        """Ensure fields are correctly populated."""
        config = (
            config.as_dict()
            if isinstance(config, PhononBandStructureSymmLine)
            else config
        )

        # legacy data contains abipy structure objects
        if (struct := config.get("structure")) and not isinstance(struct, Structure):
            config["structure"] = Structure.from_dict(struct)

        # remap legacy fields
        for k, v in {
            "lattice_rec": "reciprocal_lattice",
            "bands": "frequencies",
        }.items():
            if config.get(k):
                config[v] = config.pop(k)

        if isinstance(config["reciprocal_lattice"], dict):
            config["reciprocal_lattice"] = config["reciprocal_lattice"].get("matrix")

        return cls(**config)

    @property
    def primitive_structure(self) -> Structure | None:
        """Cache primitive structure for use in computing entropy, heat capacity, etc."""
        if self.structure and not self._primitive_structure:
            self._primitive_structure = self.structure.get_primitive_structure()
        return self._primitive_structure

    @cached_property
    def to_pmg(self) -> PhononBandStructureSymmLine:
        """Get / cache corresponding pymatgen object."""
        rlatt = Lattice(self.reciprocal_lattice)

        eigendisplacements = (
            None
            if self.eigendisplacements is None
            else np.array(self.eigendisplacements)
        )

        return PhononBandStructureSymmLine(
            [Kpoint(q, lattice=rlatt).frac_coords for q in self.qpoints],  # type: ignore[misc]
            np.array(self.frequencies),
            rlatt,
            has_nac=self.has_nac,
            eigendisplacements=eigendisplacements,
            structure=self.structure,
            labels_dict={
                k: Kpoint(v, lattice=rlatt).frac_coords
                for k, v in (self.labels_dict or {}).items()
            },
            coords_are_cartesian=False,
        )

    @property
    def _to_pmg_es_bs(self) -> BandStructureSymmLine:
        """Convenience method for crystal toolkit to create electronic-structure style band structure object."""
        rlatt = Lattice(self.reciprocal_lattice)
        return BandStructureSymmLine(
            [Kpoint(q, lattice=rlatt).frac_coords for q in self.qpoints],  # type: ignore[arg-type]
            {Spin.up: np.array(self.frequencies)},  # type: ignore[dict-item]
            rlatt,
            1e-6,  # There is no Fermi level in a Phonon DOS (these are bosons definitionally) but we want to plot the bands from zero.
            {
                k: Kpoint(v, lattice=rlatt).frac_coords  # type: ignore[misc]
                for k, v in (self.labels_dict or {}).items()
            },
            coords_are_cartesian=False,
            structure=self.structure,
        )

    @classmethod
    def from_phonopy(
        cls,
        phonon_bandstructure_file: FSPathType,
        structure_file: FSPathType | None = None,
        path_convention: str | None = None,
        label_match_tol: float = 1e-3,
        **kwargs,
    ) -> Self:
        """Create a PhononBS from phonopy .yaml output.

        Args:
        phonon_bandstructure_file (FSPathType) : path to a band.yaml-like file.
        structure_file (FSPathType or None) : path to a POSCAR-like structure file.
            If provided, takes precedence over the structure embedded in the
            band.yaml.
        path_convention (str or None) : name of the high-symmetry path
            convention used to generate the band.yaml (e.g.
            ``"latimer_munro"``, ``"setyawan_curtarolo"``, ``"hinuma"``).
            When provided AND the band.yaml does not already carry per-entry
            ``label`` keys, `from_phonopy` reconstructs ``labels_dict`` by
            matching segment-boundary q-positions against
            `pymatgen.symmetry.bandstructure.HighSymmKpath` labels for the
            requested convention. If the lookup fails for any reason (e.g.
            missing optional `seekpath` dep for ``"hinuma"``, or no match
            within `label_match_tol`), label recovery is skipped silently.
        label_match_tol (float) : tolerance (in fractional reciprocal-lattice
            coordinates) for matching boundary q-positions to canonical
            high-symmetry kpoint labels. Defaults to 1e-3.
        **kwargs : other kwargs to pass to the class constructor
        """
        with zopen(phonon_bandstructure_file, "rt") as f:
            phonopy_bandstructure = yaml.safe_load(f.read())

        phonon_entries = phonopy_bandstructure["phonon"]
        qpoints = [entry["q-position"] for entry in phonon_entries]

        # Phonopy's band.yaml stores frequencies per q-point (one entry per
        # q-point, each containing a list of bands). The `PhononBS` schema
        # expects the transpose: first index = band, second = q-point.
        n_bands = len(phonon_entries[0]["band"]) if phonon_entries else 0
        frequencies = [
            [
                phonon_entries[q_idx]["band"][b_idx]["frequency"]
                for q_idx in range(len(phonon_entries))
            ]
            for b_idx in range(n_bands)
        ]

        # Prefer per-entry labels when present (phonopy with BAND_LABELS set).
        labels_dict = {
            entry["label"]: entry["q-position"]
            for entry in phonon_entries
            if entry.get("label") is not None
        }

        # Fall back to reconstructing labels via HighSymmKpath when the
        # caller tells us which path convention was used to generate the
        # band.yaml. Pheasy-produced band.yamls lack per-entry labels but
        # do include segment boundary metadata, which is enough to look up
        # canonical high-symmetry kpoint labels.
        if not labels_dict and path_convention:
            labels_dict = (
                _recover_labels_from_segments(
                    phonopy_bandstructure,
                    structure_file=structure_file,
                    path_convention=path_convention,
                    tol=label_match_tol,
                )
                or {}
            )

        phonopy_bandstructure.update(
            qpoints=qpoints,
            frequencies=frequencies,
        )
        if labels_dict:
            phonopy_bandstructure.setdefault("labels_dict", labels_dict)

        if structure_file:
            phonopy_bandstructure["structure"] = Structure.from_file(structure_file)
        return cls(**phonopy_bandstructure, **kwargs)


class SumRuleChecks(BaseModel):
    """Container class for defining sum rule checks."""

    asr: float | None = Field(
        None, description="The violation of the acoustic sum rule."
    )
    cnsr: float | None = Field(
        None, description="The violation of the charge neutral sum rule."
    )


class PhononComputationalSettings(BaseModel):
    """Collection to store computational settings for the phonon computation."""

    # could be optional and implemented at a later stage?
    npoints_band: int | None = Field(
        None, description="number of points for band structure computation"
    )
    kpath_scheme: str | None = Field(None, description="indicates the kpath scheme")
    kpoint_density_dos: int | None = Field(
        None,
        description="number of points for computation of free energies and densities of states",
    )


class ThermalDisplacementData(BaseModel):
    """Collection to store information on the thermal displacement matrices."""

    freq_min_thermal_displacements: float | None = Field(
        None,
        description="cutoff frequency in THz to avoid numerical issues in the "
        "computation of the thermal displacement parameters",
    )
    thermal_displacement_matrix_cif: list[list[Matrix3D]] | None = Field(
        None, description="field including thermal displacement matrices in CIF format"
    )
    thermal_displacement_matrix: list[list[Matrix3D]] | None = Field(
        None,
        description="field including thermal displacement matrices in Cartesian "
        "coordinate system",
    )
    temperatures_thermal_displacements: list[float] | None = Field(
        None,
        description="temperatures at which the thermal displacement matrices"
        "have been computed",
    )


class PhononBSDOSTask(StructureMetadata):
    """Phonon band structures and density of states data."""

    identifier: IdentifierType | None = Field(
        None, description="The identifier of this phonon analysis task."
    )

    phonon_method: PhononMethod | None = Field(
        None, description="The method used to calculate phonon properties."
    )

    epsilon_static: Matrix3D | None = Field(
        None, description="The high-frequency dielectric constant."
    )

    epsilon_electronic: Matrix3D | None = Field(
        None,
        description="The electronic contribution to the high-frequency dielectric constant.",
    )

    born: list[Matrix3D] | None = Field(
        None,
        description="Born charges, only for symmetrically inequivalent atoms",
    )

    last_updated: DateTimeType = Field(
        description="Timestamp for the most recent calculation for this Material document.",
    )

    sum_rules_breaking: SumRuleChecks | None = Field(
        None,
        description="Deviations from sum rules.",
    )

    structure: StructureType | None = Field(
        None, description="Structure used in the calculation."
    )

    total_dft_energy: float | None = Field(
        None, description="total DFT energy in eV/atom."
    )

    volume_per_formula_unit: float | None = Field(
        None, description="volume per formula unit in Angstrom**3."
    )

    formula_units: int | None = Field(None, description="Formula units per cell.")

    supercell_matrix: Matrix3D | None = Field(
        None, description="matrix describing the supercell."
    )
    primitive_matrix: Matrix3D | None = Field(
        None, description="matrix describing relationship to primitive cell."
    )

    code: str | None = Field(
        None, description="String describing the code for the computation."
    )

    post_process_settings: PhononComputationalSettings | None = Field(
        None,
        description="Field including settings for the post processing code, e.g., phonopy.",
    )

    thermal_displacement_data: ThermalDisplacementData | None = Field(
        None,
        description="Includes all data of the computation of the thermal displacements",
    )

    calc_meta: list[CalcMeta] | None = Field(
        None,
        description="Metadata for individual calculations used to build this document.",
    )

    @cached_property
    def charge_neutral_sum_rule(self) -> Matrix3D | None:
        r"""Sum of Born effective charges over sites should be zero.

        See Eq. 46 of https://doi.org/10.1103/PhysRevB.55.10355:
        $$
        \sum_{i=1}^N Z_{i\alpha\beta}
        $$
        where $Z_{i\alpha\beta}$ is the N x 3 x 3 Born effective
        charge tensor, $i = 1,2,...,N$ ranges up to the number of
        atoms $N$, and $\alpha, \beta = x, y, z$ are Cartesian
        directions.

        This returns a 3 x 3 matrix. The maximum deviation
        from the sum rule is checked by `check_sum_rule_deviations`.
        """
        if self.born:
            bec = np.array(self.born)
            return tuple(tuple(row) for row in np.sum(bec, axis=0).tolist())
        return None

    @property
    def check_sum_rule_deviations(self) -> SumRuleChecks:
        """Report deviations from sum rules."""
        if not self.sum_rules_breaking:
            self.sum_rules_breaking = SumRuleChecks(
                **{
                    k: (
                        np.abs(getattr(self, attr)).max()
                        if getattr(self, attr) is not None
                        else None
                    )
                    for k, attr in {
                        "asr": "acoustic_sum_rule",
                        "cnsr": "charge_neutral_sum_rule",
                    }.items()
                }
            )
        return self.sum_rules_breaking


class PhononBSDOSDoc(PhononBSDOSTask):
    """Built data version of PhononBSDOSTask."""

    phonon_bandstructure: PhononBS | None = Field(
        None,
        description="Phonon band structure object.",
    )

    phonon_dos: PhononDOS | None = Field(
        None,
        description="Phonon density of states object.",
    )

    # needed, e.g. to compute Grueneisen parameter etc
    force_constants: list[list[Matrix3D]] | None = Field(
        None, description="Force constants between every pair of atoms in the structure"
    )

    @classmethod
    def migrate_fields(cls, **config) -> Any:
        """Migrate legacy input fields."""

        # This block is for older DFPT data
        for k, v in {
            "ph_dos": "phonon_dos",
            "ph_bs": "phonon_bandstructure",
            "e_total": "epsilon_static",
            "e_electronic": "epsilon_electronic",
            "becs": "born",
        }.items():
            if config.get(k):
                config[v] = config.pop(k)

        # migrate pymatgen objects
        for k, emmet_cls in {
            "phonon_dos": PhononDOS,
            "phonon_bandstructure": PhononBS,
        }.items():
            if (old_obj := config.get(k)) and not isinstance(old_obj, emmet_cls):
                try:
                    new_obj = emmet_cls.from_pmg(old_obj)  # type: ignore[attr-defined]
                except Exception:
                    # purposely no except here - want this to hard fail if we
                    # can't parse the input type
                    new_obj = emmet_cls(**old_obj)
                config[k] = new_obj

        if (ph_bs := config.get("phonon_bandstructure")) and not config.get(
            "structure"
        ):
            config["structure"] = ph_bs.structure

        # This block is for the migration from atomate2 --> emmet-core schemas
        if config.get("structure"):
            if isinstance(config["structure"], dict):
                config["structure"] = Structure.from_dict(config["structure"])

            old_formula_units = config.pop("formula_units", 1)
            if toten := config.pop("total_dft_energy", None):
                config["total_dft_energy"] = (
                    toten * old_formula_units / config["structure"].num_sites
                )

            config["formula_units"] = get_num_formula_units(
                config["structure"].composition
            )

        if (fc := config.get("force_constants")) and isinstance(fc, dict):
            config["force_constants"] = fc["force_constants"]

        if comp_sett := config.get("phonopy_settings"):
            config["post_process_settings"] = comp_sett

        calc_meta_migrate = {
            "uuids": "_uuid",
            "jobdirs": "_job_dir",
        }
        calc_meta_remap = {"uuids": "identifier", "jobdirs": "dir_name"}
        if any(config.get(k) for k in calc_meta_migrate):
            calc_meta: dict[str, dict[str, str]] = {}
            for field_name, splitter in calc_meta_migrate.items():
                if vals := config.get(field_name):
                    if not isinstance(vals, dict):
                        vals = vals.model_dump()
                    for k, v in vals.items():
                        job_name = k.split(splitter)[0]
                        if isinstance(v, list):
                            for idx, sub_v in enumerate(v):
                                if (sub_name := f"{job_name}-{idx}") not in calc_meta:
                                    calc_meta[sub_name] = {}
                                calc_meta[sub_name][calc_meta_remap[field_name]] = sub_v
                        else:
                            if job_name not in calc_meta:
                                calc_meta[job_name] = {}
                            calc_meta[job_name][calc_meta_remap[field_name]] = v

            config["calc_meta"] = [
                CalcMeta(
                    name=job_name,
                    **meta,
                )
                for job_name, meta in calc_meta.items()
            ]

        return cls.from_structure(
            meta_structure=config.pop("meta_structure", config.get("structure")),
            **config,
        )

    @classmethod
    def from_phonopy_pheasy_files(
        cls,
        structure_file: FSPathType,
        phonon_bandstructure_file: FSPathType | None = None,
        phonon_dos_file: FSPathType | None = None,
        projected_dos_file: FSPathType | None = None,
        force_constants_file: FSPathType | None = None,
        born_file: FSPathType | None = None,
        epsilon_static_file: FSPathType | None = None,
        epsilon_static_and_born_file: FSPathType | None = None,
        phonopy_output_file: FSPathType | None = None,
        **kwargs,
    ) -> Self:
        """
        Create a PhononBSDOSTask from a list of explicit Phonopy/Pheasy file paths.
        """

        try:
            import phonopy
        except ImportError:
            phonopy = None  # type: ignore[assignment]

        cls_config: dict[str, Any] = {
            "structure": Structure.from_file(structure_file),
        }
        if "poscar" in str(structure_file).lower():
            cls_config["code"] = "vasp"

        if phonon_bandstructure_file:
            cls_config["phonon_bandstructure"] = PhononBS.from_phonopy(
                phonon_bandstructure_file,
                structure=cls_config["structure"],
            )

        if phonon_dos_file:
            cls_config["phonon_dos"] = PhononDOS.from_phonopy(
                phonon_dos_file,
                projected_dos_file=projected_dos_file,
                structure=cls_config["structure"],
            )

        if force_constants_file:
            # read FORCE_CONSTANTS manually
            force_constant_matrix: np.ndarray
            idxs: tuple[int | None, int | None] = (
                None,
                None,
            )
            irow = 0
            with zopen(force_constants_file, "rt") as f:
                for idx, line in enumerate(f.read().splitlines()):
                    vals = line.strip().split()
                    if idx == 0:
                        force_constant_matrix = np.zeros(
                            (int(vals[0]), int(vals[1]), 3, 3)
                        )
                    elif len(vals) == 2:
                        # idxs written like this for mypy
                        idxs = (
                            int(vals[0]) - 1,
                            int(vals[1]) - 1,
                        )
                        irow = 0
                    elif len(vals) == 3:
                        force_constant_matrix[idxs[0], idxs[1], irow] = [
                            float(v) for v in vals
                        ]
                        irow += 1
            cls_config["force_constants"] = force_constant_matrix.tolist()

        if phonopy_output_file:
            with zopen(phonopy_output_file, "rt") as f:
                phonopy_output = yaml.safe_load(f.read())
            for k in ("primitive_matrix", "supercell_matrix"):
                cls_config[k] = phonopy_output.get(k)

        if all(
            fs is not None and ".npz" in Path(fs).name
            for fs in (born_file, epsilon_static_file)
        ):
            cls_config["born"] = np.load(born_file)  # type: ignore[arg-type]
            cls_config["epsilon_static"] = np.load(epsilon_static_file)  # type: ignore[arg-type]

        elif (
            phonopy_output_file
            and epsilon_static_and_born_file
            and "BORN" in Path(epsilon_static_and_born_file).name
        ):
            if phonopy is None:
                raise ImportError("You must `pip install phonopy` to parse BORN.")

            phonopy_calc = phonopy.load(phonopy_output_file)
            with zopen(epsilon_static_and_born_file, "rt") as f:
                born_data = phonopy.file_IO.parse_BORN_from_strings(
                    f.read(), phonopy_calc.unitcell
                )

            cls_config.update(
                {
                    "epsilon_static": born_data["dielectric"],
                    "born": born_data["born"],
                }
            )

        return cls.from_structure(cls_config["structure"], **cls_config, **kwargs)

    @classmethod
    def from_phonopy_pheasy_directory(
        cls,
        phonon_dir: Path | str,
        **kwargs,
    ) -> Self:
        """Create a PhononBSDOSDoc from a Phonopy/Pheasy directory.

        Parameters
        -----------
        phonon_dir : str or Path
        **kwargs to pass to `PhononBSDOSDoc.from_phonopy_pheasy_files`.
            The user can override the default file paths for select files
            using this.
        """
        phonon_path = Path(phonon_dir).resolve()
        file_paths = {}
        for k, file_name in DEFAULT_PHONON_FILES.items():
            file_path = phonon_path / kwargs.pop(f"{k}_file", file_name)
            if (file_path := Path(zpath(file_path))).exists():
                file_paths[f"{k}_file"] = file_path
        return cls.from_phonopy_pheasy_files(**file_paths, **kwargs)

    @cached_property
    def has_imaginary_modes(self) -> bool | None:
        tol: float = 1e-5
        if self.phonon_bandstructure:
            return self.phonon_bandstructure.to_pmg.has_imaginary_freq(tol=tol)
        elif self.phonon_dos:
            return bool(
                np.any(
                    np.array(self.phonon_dos.densities)[
                        np.array(self.phonon_dos.frequencies) < tol
                    ]
                    > tol
                )
            )
        return None

    def _get_thermo_from_dos(
        self,
        quantity: Literal["entropy", "cv", "internal_energy", "helmholtz_free_energy"],
        temperature: float,
        normalization: Literal["atoms", "formula_units"] | None = "formula_units",
    ) -> float | None:
        """Get a thermodynamic property from the phonon DOS.

        quantity : "entropy", "cv", "internal_energy", or "helmholtz_free_energy"
            If "entropy" or "cv", the units before normalization are J/(K * mol).
            If "internal_energy" or "helmholtz_free_energy", the units before normalization are J/mol.
        temperature : float
        normalization : "atoms", "formula_units", or None
            Whether to normalize by the number of atoms in the cell ("atoms"),
            the number of formula units ("formula_units", default), or not (None).
        """

        if not self.phonon_dos:
            return None

        if normalization and not self.structure:
            raise ValueError(
                "Cannot normalize thermodynamic quantities without a structure."
            )

        norm_fac = 1.0
        if normalization == "atoms":
            norm_fac = 1.0 / self.structure.num_sites  # type: ignore[union-attr]
        elif normalization == "formula_units":
            norm_fac = 1.0 / self.formula_units  # type: ignore[operator]
        elif normalization:
            raise ValueError(f"Unknown {normalization=} convention.")

        return (
            getattr(self.phonon_dos.to_pmg, quantity)(
                temp=temperature,
            )
            * norm_fac
        )

    def entropy(
        self,
        temperature: float,
        normalization: Literal["atoms", "formula_units"] | None = "formula_units",
    ) -> float | None:
        """Compute the entropy in J/(K * mol * formula units).

        temperature : float
        normalization : "atoms", "formula_units", or None
            Whether to normalize by the number of atoms in the cell ("atoms"),
            the number of formula units ("formula_units"), or not (None).
        """
        return self._get_thermo_from_dos(
            "entropy", temperature, normalization=normalization
        )

    def heat_capacity(
        self,
        temperature: float,
        normalization: Literal["atoms", "formula_units"] | None = "formula_units",
    ) -> float | None:
        """Compute the heat capacity in J/(K * mol * formula units).

        temperature : float
        normalization : "atoms", "formula_units", or None
            Whether to normalize by the number of atoms in the cell ("atoms"),
            the number of formula units ("formula_units"), or not (None).
        """
        return self._get_thermo_from_dos("cv", temperature, normalization=normalization)

    def internal_energy(
        self,
        temperature: float,
        normalization: Literal["atoms", "formula_units"] | None = "formula_units",
    ) -> float | None:
        """Compute the internal energy in J/(mol * formula units).

        temperature : float
        normalization : "atoms", "formula_units", or None
            Whether to normalize by the number of atoms in the cell ("atoms"),
            the number of formula units ("formula_units"), or not (None).
        """
        return self._get_thermo_from_dos(
            "internal_energy", temperature, normalization=normalization
        )

    def free_energy(
        self,
        temperature: float,
        normalization: Literal["atoms", "formula_units"] | None = "formula_units",
    ) -> float | None:
        """Compute the Helmholtz free energy in J/(mol * formula units).

        temperature : float
        normalization : "atoms", "formula_units", or None
            Whether to normalize by the number of atoms in the cell ("atoms"),
            the number of formula units ("formula_units"), or not (None).
        """
        return self._get_thermo_from_dos(
            "helmholtz_free_energy", temperature, normalization=normalization
        )

    def compute_thermo_quantities(
        self,
        temperatures: Sequence[float],
        normalization: Literal["atoms", "formula_units"] | None = "formula_units",
    ) -> dict[str, Sequence[float | None]]:
        """Compute all thermodynamic quantities as a convenience method."""

        quantities = {
            "entropy": "entropy",
            "heat_capacity": "cv",
            "internal_energy": "internal_energy",
            "free_energy": "helmholtz_free_energy",
        }

        thermo_props: dict[str, Sequence[float | None]] = {
            k: [
                self._get_thermo_from_dos(v, temp, normalization=normalization)  # type: ignore[arg-type]
                for temp in temperatures
            ]
            for k, v in quantities.items()
        }
        thermo_props["temperature"] = temperatures
        return thermo_props

    @cached_property
    def acoustic_sum_rule(self) -> np.ndarray | None:
        r"""Determine violations of the acoustic sum rule.

        The acoustic sum rule for the interatomic force constant (IFC) tensor
        is (Eq. 44 of https://doi.org/10.1103/PhysRevB.55.10355):
        $$
        \sum_{j=1}^N C_{ij\alpha\beta} = 0.
        $$
        The IFCs are an N x N x 3 x 3 tensor, where N = number of atoms
        in the structure (`structure.num_sites`).
        $i, j = 1,2,3,...,N$, and $\alpha,\beta = x,y,z$ are the Cartesian
        directions.

        The maximum deviation from the sum rule is checked
        by `check_sum_rule_deviations`.

        For an ASR-corrected force-constant set, this is numerically zero.
        """
        if self.force_constants:
            return np.asarray(self.force_constants).sum(axis=1)
        return None


class PhononWarnings(DocEnum):
    ASR = "ASR break", "acoustic sum rule max breaking is larger than 30 cm^-1."
    CNSR = "CNSR break", "charge neutrality sum rule max breaking is larger than 0.2."
    NEG_FREQ = (
        "has negative frequencies",
        "phonon band structure has negative "
        "frequencies anywhere in the Brillouin zone.",
    )
    SMALL_Q_NEG_FREQ = (
        "has small q negative frequencies",
        "the phonon band structure has negative frequencies,"
        " but these are small and very close to the Gamma point "
        "(usually related to numerical errors).",
    )


vec2D = list[float, float]  # type: ignore[type-arg]
vec3D = list[float, float, float]  # type: ignore[type-arg]


class PhononWebsiteStruct(TypedDict):
    atom_numbers: list[int]
    atom_pos_car: list[vec3D]
    atom_pos_red: list[vec3D]
    atom_types: list[str]
    distances: list[float]
    eigenvalues: list[list[float]]
    formula: str
    # fails arrow conversion
    # highsym_qpts: list[list[int, str]]
    lattice: list[vec3D]
    line_breaks: list[tuple[int, int]]
    name: str
    natoms: int
    qpoints: list[vec3D]
    repetitions: list[int, int, int]  # type: ignore[type-arg]
    vectors: list[list[list[list[vec2D, vec2D, vec2D]]]]  # type: ignore[type-arg]


class PhononWebsiteBS(BaseModel):
    """
    Document with a serialized version of the phonon band structure suitable
    for the phononwebsite (http://henriquemiranda.github.io/phononwebsite/).
    """

    material_id: str = Field(
        ...,
        description="The ID of this material, used as a universal reference across property documents."
        "This comes in the form: mp-******",
    )

    doc_type: Literal["phononwebsite"] = Field(
        "phononwebsite",
        description="The type of the document: a phonon band structure for the phononwebsite.",
    )

    phononwebsite: PhononWebsiteStruct | None = Field(
        None,
        description="Phononwebsite dictionary to plot the animated " "phonon modes.",
    )

    last_updated: DateTimeType = Field(
        description="Timestamp for the most recent calculation update for this property",
    )

    created_at: DateTimeType = Field(
        description="Timestamp for when this material document was first created",
    )


class Ddb(BaseModel):
    """
    Document with a the string version of the DDB file produced by abinit.
    """

    material_id: str = Field(
        ...,
        description="The ID of this material, used as a universal reference across property documents."
        "This comes in the form: mp-******",
    )

    doc_type: Literal["ddb"] = Field(
        "ddb", description="The type of the document: a DDB file."
    )

    ddb: str | None = Field(None, description="The string of the DDB file.")

    last_updated: DateTimeType = Field(
        description="Timestamp for the most recent calculation update for this property",
    )

    created_at: DateTimeType = Field(
        description="Timestamp for when this material document was first created",
    )


class ThermodynamicProperties(BaseModel):
    """
    Definition of the thermodynamic properties extracted from the phonon frequencies.
    """

    temperatures: list[float] = Field(
        ...,
        description="The list of temperatures at which the thermodynamic properties "
        "are calculated",
    )

    cv: list[float] = Field(
        ...,
        description="The values of the constant-volume specific heat.",
        alias="heat_capacity",
    )

    entropy: list[float] = Field(
        ..., description="The values of the vibrational entropy."
    )


class VibrationalEnergy(BaseModel):
    """
    Definition of the vibrational contribution to the energy as function of
    the temperature.
    """

    temperatures: list[float] = Field(
        ...,
        description="The list of temperatures at which the thermodynamic properties "
        "are calculated",
    )

    internal_energy: list[float] = Field(
        ..., description="The values of the phonon contribution to the internal energy."
    )

    helmholtz_free_energy: list[float] = Field(
        ..., description="The values of the Helmholtz free energy."
    )

    zero_point_energy: float = Field(
        ..., description="The value of the zero point energy."
    )


class Phonon(StructureMetadata):
    """
    Definition for a document with data produced by a phonon calculation.
    """

    material_id: str = Field(
        ...,
        description="The ID of this material, used as a universal reference across property documents."
        "This comes in the form: mp-******",
    )

    structure: StructureType = Field(
        ..., description="The relaxed structure for the phonon calculation."
    )

    asr_break: float | None = Field(
        None, description="The maximum breaking of the acoustic sum rule (ASR)."
    )

    warnings: list[PhononWarnings] | None = Field(
        None, description="List of warnings associated to the phonon calculation."
    )

    dielectric: DielectricDoc | None = Field(
        None, description="Dielectric properties obtained during a phonon calculations."
    )

    becs: BornEffectiveCharges | None = Field(
        None, description="Born effective charges obtained for a phonon calculation."
    )

    ir_spectra: IRDielectric | None = Field(None, description="The IRDielectricTensor.")

    thermodynamic: ThermodynamicProperties | None = Field(
        None,
        description="The thermodynamic properties extracted from the phonon "
        "frequencies.",
    )

    vibrational_energy: VibrationalEnergy | None = Field(
        None, description="The vibrational contributions to the total energy."
    )

    last_updated: DateTimeType = Field(
        description="Timestamp for when this document was last updated",
    )

    created_at: DateTimeType = Field(
        description="Timestamp for when this material document was first created",
    )


class TypedBaseInputDict(TypedDict):
    charge: float
    chksymbreak: int
    ecut: float
    fband: float
    kptopt: int
    nband: int
    nbdbuf: int
    ngkpt: list[int]
    nshiftk: int
    nspden: int
    nspinor: int
    nsppol: int
    nstep: int
    pawecutdg: float
    shiftk: list[list[float]]
    tolvrs: float


class TypedDdBaseDict(TypedBaseInputDict):
    chkprim: int
    nqpt: int
    qpt: list[int]
    rfdir: list[int]
    rfelfd: int


class TypedDdeInputDict(TypedDdBaseDict):
    prtwf: int


class TypedDdkInputDict(TypedDdBaseDict):
    iscf: float


class TypedPhononInputDict(TypedDdeInputDict):
    rfatpol: list[int]


class TypedPseudopotentialsDict(TypedDict):
    md5: list[str]
    name: list[str]


class TypedAbinitInputVars(TypedDict):
    dde_input: TypedDdeInputDict
    ddk_input: TypedDdkInputDict
    ecut: float
    gs_input: TypedBaseInputDict
    ngkpt: list[int]
    ngqpt: list[int]
    occopt: int
    phonon_input: TypedPhononInputDict
    pseudopotentials: TypedPseudopotentialsDict
    shiftk: list[list[float]]
    tsmear: int
    wfq_input: TypedBaseInputDict  # have to guess here without an example entry


class AbinitPhonon(Phonon):
    """
    Definition for a document with data produced from a phonon calculation
    with Abinit.
    """

    abinit_input_vars: TypedAbinitInputVars | None = Field(
        None,
        description="Dict representation of the inputs used to obtain the phonon"
        "properties and the main general options (e.g. number of "
        "k-points, number of q-points).",
    )


class SoundVelocity(BaseModel):
    """
    Definition for a document with the sound velocities of the acoustic modes
    close to Gamma, as obtained from a phonon calculation.
    """

    material_id: str = Field(
        ...,
        description="The ID of this material, in the form: mp-******",
    )

    structure: StructureType = Field(
        ..., description="The relaxed structure for the phonon calculation."
    )

    directions: list[Vector3D] = Field(
        ...,
        description="Q-points identifying the directions for the calculation"
        "of the speed of sound. In fractional coordinates.",
    )

    labels: list[str | None] = Field(..., description="labels of the directions.")

    sound_velocities: list[Vector3D] = Field(
        ...,
        description="Values of the sound velocities in SI units.",
    )

    mode_types: list[tuple[str | None, str | None, str | None]] = Field(
        ...,
        description="The types of the modes ('transversal', 'longitudinal'). "
        "None if not correctly identified.",
    )

    last_updated: DateTimeType = Field(
        description="Timestamp for when this document was last updated",
    )

    created_at: DateTimeType = Field(
        description="Timestamp for when this material document was first created",
    )


class ThermalDisplacement(BaseModel):
    """
    Definition of a Document for the generalized density of states and
    mean square displacements related to phonon oscillations.
    """

    material_id: str = Field(
        ...,
        description="The ID of this material, used as a universal reference across property documents."
        "This comes in the form: mp-******",
    )

    last_updated: DateTimeType = Field(
        description="Timestamp for the most recent calculation update for this property",
    )

    created_at: DateTimeType = Field(
        description="Timestamp for when this material document was first created",
    )

    nsites: int = Field(
        ...,
        description="The number of sites in the structure.",
    )

    nomega: int = Field(
        ...,
        description="The number of frequencies.",
    )

    ntemp: int = Field(
        ...,
        description="The number of temperatures for which the displacements are calculated",
    )

    temperatures: list[float] = Field(
        ...,
        description="The list of temperatures at which the thermodynamic properties "
        "are calculated",
    )

    frequencies: list[float] = Field(
        ..., description="The list of frequencies for the generalized DOS"
    )

    gdos_aijw: Tensor4R = Field(
        ...,
        description=" Generalized DOS in Cartesian coords, with shape (nsites, 3, 3, nomega)",
    )

    amu: dict[str, float] = Field(
        ..., description="Dictionary of the atomic masses in atomic units."
    )

    structure: StructureType = Field(
        ..., description="The relaxed structure for the phonon calculation."
    )

    ucif_t: Tensor4R = Field(
        ...,
        description="Mean squared displacement U tensors as a function of T for T in tmesh in CIF format."
        "With shape (natom, 3, 3, ntemp) ",
    )
    ucif_string_t300k: str = Field(
        ...,
        description="Mean squared displacement U tensors at T=300K in CIF string format.",
    )
