from typing import TYPE_CHECKING

import numpy as np
from nomad.metainfo import MEnum, Quantity

if TYPE_CHECKING:
    from nomad.datamodel.context import Context
    from nomad.datamodel.datamodel import EntryArchive
    from nomad.metainfo import Section
    from structlog.stdlib import BoundLogger

from nomad_simulations.schema_packages.data_types import positive_int
from nomad_simulations.schema_packages.physical_property import PhysicalProperty
from nomad_simulations.schema_packages.properties.energies import BaseEnergy

######################################
# fundamental thermodynamic properties
######################################


class Pressure(PhysicalProperty):
    """
    The force exerted per unit area by gas particles as they collide with the walls of
    their container.
    """

    # iri = 'http://fairmat-nfdi.eu/taxonomy/Pressure' # ! Does not yet exist in taxonomy

    value = Quantity(
        type=np.float64,
        unit='pascal',
        description="""
        """,
    )


class Volume(PhysicalProperty):
    """
    the amount of three-dimensional space that a substance or material occupies.
    """

    #! Above description suggested for taxonomy
    # TODO check back on definition after first taxonomy version

    iri = 'http://fairmat-nfdi.eu/taxonomy/Volume'

    value = Quantity(
        type=np.float64,
        unit='m ** 3',
        description="""
        """,
    )


class Temperature(PhysicalProperty):
    """
    a measure of the average kinetic energy of the particles in a system.
    """

    value = Quantity(
        type=np.float64,
        unit='kelvin',
        description="""
        """,
    )


class Heat(BaseEnergy):
    """
    The transfer of thermal energy **into** a system.
    """


class Work(BaseEnergy):
    """
    The energy transferred to a system by means of force applied over a distance.
    """


class InternalEnergy(BaseEnergy):
    """
    The total energy contained within a system, encompassing both kinetic and potential
    energies of the particles. The change in `InternalEnergy` for some thermodynamic
    process may be expressed as the `Heat` minus the `Work`.
    """


class Enthalpy(BaseEnergy):
    """
    The total heat content of a system, defined as 'InternalEnergy' + 'Pressure' * 'Volume'.
    """


class Entropy(PhysicalProperty):
    """
    A measure of the disorder or randomness in a system.

    From a thermodynamic perspective, `Entropy` is a measure of the system's energy
    dispersal at a specific temperature, and can be interpreted as the unavailability of
    a system's thermal energy for conversion into mechanical work. For a reversible
    process, the change in `Entropy` is given mathematically by an integral over the
    infinitesimal `Heat` (i.e., thermal energy transfered into the system) divided by the
    `Temperature`.

    From a statistical mechanics viewpoint, entropy quantifies the number of microscopic
    configurations (microstates) that correspond to a thermodynamic system's macroscopic
    state, as given by the Boltzmann equation for entropy.
    """

    value = Quantity(
        type=np.float64,
        unit='joule / kelvin',
        description="""
        """,
    )


class GibbsFreeEnergy(BaseEnergy):
    """
    The energy available to do work in a system at constant temperature and pressure,
    given by `Enthalpy` - `Temperature` * `Entropy`.
    """


class HelmholtzFreeEnergy(BaseEnergy):
    """
    The energy available to do work in a system at constant volume and temperature,
    given by `InternalEnergy` - `Temperature` * `Entropy`.
    """


class ChemicalPotential(BaseEnergy):
    """
    Free energy cost of adding or extracting a particle from a thermodynamic system.

    At finite temperature, the chemical potential determines the equilibrium condition
    for particle exchange between different phases or subsystems. It can be defined
    as the partial derivative of the internal energy with respect to particle number
    at constant entropy and volume, or equivalently as the partial derivative of
    the Gibbs free energy with respect to particle number at constant temperature
    and pressure.
    """

    iri = 'http://fairmat-nfdi.eu/taxonomy/ChemicalPotential'

    temperature = Quantity(
        type=np.float64,
        unit='kelvin',
        description="""
        Temperature at which the chemical potential is calculated.
        Essential for finite-temperature calculations.
        """,
    )

    particle_number = Quantity(
        type=np.float64,
        description="""
        Number of particles (or particle density) for which the chemical potential applies.
        Can represent electron number, atom number, or other relevant particle count.
        """,
    )

    fermi_energy = Quantity(
        type=np.float64,
        unit='joule',
        description="""
        Fermi energy at T=0K, used as reference for finite-temperature chemical potential.
        At T=0, the chemical potential equals the Fermi energy.
        """,
    )

    type = Quantity(
        type=str,
        description="""
        Type of chemical potential calculation. Examples: 'electronic', 'atomic',
        'ionic', 'molecular'. Helps identify what kind of particles this applies to.
        """,
    )


class HeatCapacity(PhysicalProperty):
    """
    Amount of heat to be supplied to a material to produce a unit change in its temperature.
    """

    value = Quantity(
        type=np.float64,
        unit='joule / kelvin',
        description="""
        """,
    )


################################
# other thermodynamic properties
################################


class VirialTensor(BaseEnergy):
    """
    A measure of the distribution of internal forces and the overall stress within
    a system of particles. Mathematically, the virial tensor is defined as minus the sum
    of the dot product between the position and force vectors for each particle.
    The `VirialTensor` can be related to the non-ideal pressure of the system through
    the virial theorem.
    """

    rank = [3, 3]


class MassDensity(PhysicalProperty):
    """
    Mass per unit volume of a material.
    """

    value = Quantity(
        type=np.float64,
        unit='kg / m ** 3',
        description="""
        """,
    )


# ? fit better elsewhere
class Hessian(PhysicalProperty):
    """
    A square matrix of second-order partial derivatives of a potential energy function,
    describing the local curvature of the energy surface.
    """

    n_hessian_dim = Quantity(
        type=positive_int(),
        description="""
        Matrix dimension (number of degrees of freedom) of the square Hessian in the
        coordinate basis used by the parser. For Cartesian atomic Hessians this is
        typically 3 * N_atoms; constrained/filtered coordinates should use the
        remaining degrees of freedom.
        """,
    )

    value = Quantity(
        type=np.float64,
        unit='joule / m ** 2',
        shape=['n_hessian_dim', 'n_hessian_dim'],
        description="""
        """,
    )

    n_negative_eigenvalues = Quantity(
        type=positive_int(),
        description="""
        Number of negative Hessian eigenvalues (imaginary vibrational frequencies).
        A value of 0 indicates a local minimum, 1 a first-order saddle point, and
        >1 a higher-order saddle point. Leave unset if the Hessian was evaluated
        away from a stationary point.
        """,
    )

    eigenvalues = Quantity(
        type=np.float64,
        unit='joule / m ** 2',
        shape=['*'],
        description="""
        Eigenvalues of the Hessian. Sorted during normalization with positive values
        descending, followed by zeros, then negative values ascending (most negative
        first). Very low-magnitude modes in solid-state phonon calculations
        (e.g., <100 cm-1) or when using RI/DF approximations often reflect numerical
        artifacts rather than true instabilities.
        """,
    )

    stationary_point_type = Quantity(
        type=MEnum(
            'minimum',
            'maximum',
            'saddle_point',
            'non_stationary',
            'unavailable',
        ),
        description="""
        Stationary-point classification (requires zero gradient) based on Hessian
        eigenvalue signs. Use 'saddle_point' for any stationary point with one or
        more negative eigenvalues (a transition state corresponds to exactly one).
        Use 'maximum' when all eigenvalues are negative (negative-definite Hessian).
        Use 'non_stationary' if the Hessian was evaluated where the gradient is
        non-zero and no stationary point classification applies. Use 'unavailable'
        when no classification could be determined from the data.
        """,
    )

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

        if self.eigenvalues is not None:
            vals = np.array(self.eigenvalues)
            positives = np.sort(vals[vals > 0])[::-1]
            zeros = vals[np.isclose(vals, 0)]
            negatives = np.sort(vals[vals < 0])
            self.eigenvalues = np.concatenate([positives, zeros, negatives])
