from typing import TYPE_CHECKING

import numpy as np
from nomad.datamodel.datamodel import JSON, ArchiveSection
from nomad.metainfo import Quantity, SubSection

if TYPE_CHECKING:
    from nomad.datamodel.datamodel import EntryArchive
    from structlog.stdlib import BoundLogger

from nomad_simulations.schema_packages.model_method import ModelMethod
from nomad_simulations.schema_packages.model_system import ModelSystem
from nomad_simulations.schema_packages.physical_property import PhysicalProperty
from nomad_simulations.schema_packages.properties import (
    AbsorptionSpectrum,
    ChemicalPotential,
    CrystalFieldSplitting,
    ElectronicBandGap,
    ElectronicBandStructure,
    ElectronicDensityOfStates,
    ElectronicEigenvalues,
    ElectronicGreensFunction,
    ElectronicSelfEnergy,
    FermiSurface,
    HoppingMatrix,
    HybridizationFunction,
    KineticEnergy,
    Occupancy,
    Permittivity,
    PotentialEnergy,
    QuasiparticleWeight,
    RadiusOfGyration,
    Temperature,
    TotalEnergy,
    TotalForce,
    XASSpectrum,
)

from .common import SimulationTime


# MK: I don't think this should live here.
# @all: where to move this?
class SCFSteps(ArchiveSection):
    """
    Data recorded at each step of a self-consistent DFT calculation.
    """

    energies_total = Quantity(
        shape=['*'],
        type=float,
        unit='joule',
        description="""
        Total energy at each SCF step.
        """,
    )

    delta_energies_total = Quantity(
        shape=['*'],
        type=float,
        unit='joule',
        description="""
        Absolute change of total energy at each SCF step.
        """,
    )

    delta_potential_rms = Quantity(
        shape=['*'],
        type=float,
        unit='joule',
        description="""
        Root mean square of change of potential energy at each SCF step.
        """,
    )

    delta_density_rms = Quantity(
        shape=['*'],
        type=float,
        unit='coulomb',
        description="""
        Root mean square of change of potential energy at each SCF step.
        """,
    )

    delta_wavefunction_rms = Quantity(
        shape=['*'],
        type=float,
        unit='dimensionless',
        description="""
        Root mean square of change of wavefunction coefficients at each SCF step.
        Dimensionless quantity representing convergence of orbital coefficients.
        """,
    )

    delta_force_abs = Quantity(
        shape=['*'],
        type=float,
        unit='newton',
        description="""
        Absolute change of forces at each SCF step.
        """,
    )

    durations = Quantity(
        shape=['*'],
        type=float,
        unit='s',
        description="""
        Time spent at each SCF step.
        """,
    )

    code_specific_quantities = Quantity(
        type=JSON,
        description="""
        Code specific quantities that are recorded during SCF convergence.
        """,
    )


# TODO: Outputs should not be of type time, but the workflow should be instead?
class Outputs(SimulationTime):
    """
    Output properties of a simulation. This base class can be used for inheritance in any of the output properties
    defined in this schema.

    It contains references to the specific sections used to obtain the output properties, as well as
    information if the output `is_derived` from another output section or directly parsed from the simulation output files.
    """

    normalizer_level = 2

    model_system_ref = Quantity(
        type=ModelSystem,
        description="""
        Reference to the `ModelSystem` section in which the output physical properties were calculated.
        """,
    )

    model_method_ref = Quantity(
        type=ModelMethod,
        description="""
        Reference to the `ModelMethod` section containing the details of the mathematical
        model with which the output physical properties were calculated.
        """,
    )

    absorption_spectra = SubSection(sub_section=AbsorptionSpectrum.m_def, repeats=True)

    chemical_potentials = SubSection(sub_section=ChemicalPotential.m_def, repeats=True)

    crystal_field_splittings = SubSection(
        sub_section=CrystalFieldSplitting.m_def, repeats=True
    )

    electronic_band_gaps = SubSection(sub_section=ElectronicBandGap.m_def, repeats=True)

    electronic_band_structures = SubSection(
        sub_section=ElectronicBandStructure.m_def, repeats=True
    )

    electronic_dos = SubSection(
        sub_section=ElectronicDensityOfStates.m_def, repeats=True
    )

    electronic_eigenvalues = SubSection(
        sub_section=ElectronicEigenvalues.m_def, repeats=True
    )

    electronic_greens_functions = SubSection(
        sub_section=ElectronicGreensFunction.m_def, repeats=True
    )

    electronic_self_energies = SubSection(
        sub_section=ElectronicSelfEnergy.m_def, repeats=True
    )

    fermi_surfaces = SubSection(sub_section=FermiSurface.m_def, repeats=True)

    hopping_matrices = SubSection(sub_section=HoppingMatrix.m_def, repeats=True)

    hybridization_functions = SubSection(
        sub_section=HybridizationFunction.m_def, repeats=True
    )

    kinetic_energies = SubSection(sub_section=KineticEnergy.m_def, repeats=True)

    occupancies = SubSection(sub_section=Occupancy.m_def, repeats=True)

    permittivities = SubSection(sub_section=Permittivity.m_def, repeats=True)

    potential_energies = SubSection(sub_section=PotentialEnergy.m_def, repeats=True)

    quasiparticle_weights = SubSection(
        sub_section=QuasiparticleWeight.m_def, repeats=True
    )

    radii_of_gyration = SubSection(sub_section=RadiusOfGyration.m_def, repeats=True)

    temperatures = SubSection(sub_section=Temperature.m_def, repeats=True)

    total_forces = SubSection(sub_section=TotalForce.m_def, repeats=True)

    total_energies = SubSection(sub_section=TotalEnergy.m_def, repeats=True)

    xas_spectra = SubSection(sub_section=XASSpectrum.m_def, repeats=True)

    scf_steps = SubSection(sub_section=SCFSteps.m_def, repeats=False)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def extract_spin_polarized_property(
        self, property_name: str
    ) -> list[PhysicalProperty]:
        """
        Extracts the spin-polarized properties if present from the property name and returns them as a list of two elements in
        which each element refers to each `spin_channel`. If the return list is empty, it means that the simulation is not
        spin-polarized (i.e., `spin_channel` is not defined).

        Args:
            property_name (str): The name of the property to be extracted.

        Returns:
            (list[PhysicalProperty]): The list of spin-polarized properties.
        """
        spin_polarized_properties = []
        properties = getattr(self, property_name)
        for prop in properties:
            if prop.spin_channel is None:
                continue
            spin_polarized_properties.append(prop)
        return spin_polarized_properties

    def set_model_system_ref(self) -> ModelSystem | None:
        """
        Provide a suggested `ModelSystem` that corresponds to the output collection.

        Returns:
            (Optional[ModelSystem]): corresponding `ModelSystem`. `None` if there is no such section.
        """
        if self.m_parent is not None:
            model_systems = self.m_parent.model_system
            outputs = self.m_parent.outputs
            if (
                isinstance(model_systems, list)
                and isinstance(outputs, list)
                and len(model_systems) == len(outputs)
                and len(model_systems) > 0
            ):
                return model_systems[self.m_parent_index]
        return None

    def set_model_method_ref(self) -> ModelMethod | None:
        """
        Set the reference to the last `ModelMethod` if this is not set in the output. This is only
        valid if there is only one `ModelMethod` in the parent section.

        Returns:
            (Optional[ModelMethod]): The reference to the last `ModelMethod`.
        """
        if self.m_parent is not None:
            model_methods = self.m_parent.model_method
            if model_methods is not None and len(model_methods) == 1:
                return model_methods[-1]
        return None

    def _compute_energy_deltas(self, logger: 'BoundLogger'):
        """
        Compute delta_energies_total from consecutive total_energies values.

        Returns array of absolute energy differences between consecutive steps.
        """
        try:
            if self.total_energies is None or len(self.total_energies) < 2:
                return None

            # Extract energy values from each step
            energy_values = [e.value for e in self.total_energies]

            # Compute differences manually to preserve Pint units
            deltas = []
            for i in range(1, len(energy_values)):
                delta = np.abs(energy_values[i] - energy_values[i - 1])
                deltas.append(delta)

            # Convert list to array-like structure
            if len(deltas) > 0 and hasattr(deltas[0], 'magnitude'):
                # Pint quantities - stack magnitudes and add units
                magnitudes = [d.magnitude for d in deltas]
                return np.array(magnitudes) * deltas[0].units
            else:
                return np.array(deltas)

        except (AttributeError, IndexError, TypeError, ValueError) as e:
            logger.debug(f'Could not compute delta_energies_total: {e}')
            return None

    def _compute_force_norms(self, logger: 'BoundLogger'):
        """
        Compute delta_force_abs from total_forces as force norms.

        Returns array of force norms (L2 norm of 3D force vectors).
        """
        try:
            if self.total_forces is None or len(self.total_forces) == 0:
                return None

            # Get force values (Pint Quantity with shape [n_atoms, 3])
            force_values = self.total_forces[-1].value

            # Compute force norms (preserves Pint units)
            force_norms = ((force_values**2).sum(axis=1)) ** 0.5

            return force_norms
        except (AttributeError, IndexError, TypeError) as e:
            logger.debug(f'Could not compute delta_force_abs: {e}')
            return None

    def normalize(self, archive: 'EntryArchive', logger: 'BoundLogger') -> None:
        super().normalize(archive, logger)

        # Set ref to the last `ModelSystem` if this is not set in the output
        if self.model_system_ref is None:
            self.model_system_ref = self.set_model_system_ref()

        # Set ref to the last `ModelMethod` if this is not set in the output
        if self.model_method_ref is None:
            self.model_method_ref = self.set_model_method_ref()

        # Populate missing SCF delta quantities from available data
        if self.scf_steps is not None:
            # Define delta computations: (delta_field, compute_function)
            delta_computations = [
                ('delta_energies_total', self._compute_energy_deltas),
                ('delta_force_abs', self._compute_force_norms),
            ]

            for delta_field, compute_func in delta_computations:
                if getattr(self.scf_steps, delta_field, None) is None:
                    computed_value = compute_func(logger)
                    if computed_value is not None:
                        setattr(self.scf_steps, delta_field, computed_value)


class WorkflowOutputs(Outputs):
    """
    This section contains output properties that depend on a single system, but were
    calculated as part of a workflow (e.g., the energies from a geometry optimization),
    and thus may include step information.
    """

    step = Quantity(
        type=np.int32,
        description="""
        The step number with respect to the workflow.
        """,
    )

    # TODO add this in when we link to nomad-simulations-workflow schema
    # ? check potential circular imports problems when the nomad-simulations-workflow schema is transferred here
    # workflow_ref = Quantity(
    #     type=SimulationWorkflow,
    #     description="""
    #     Reference to the `SelfConsistency` section that defines the numerical settings to converge the
    #     output property.
    #     """,
    # )


class TrajectoryOutputs(WorkflowOutputs):
    """
    This section contains output properties that depend on a single system, but were
    calculated as part of a trajectory (e.g., temperatures from a molecular dynamics
    simulation), and thus may include time information.
    """

    time = Quantity(
        type=np.float64,
        unit='ps',
        description="""
        The elapsed simulated physical time since the start of the trajectory.
        """,
    )
