from typing import TYPE_CHECKING

import numpy as np
from nomad.metainfo import 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.physical_property import PhysicalProperty


class HoppingMatrix(PhysicalProperty):
    """
    Transition probability between two atomic orbitals in a tight-binding model.

    Entity references for orbitals should use the `ElectronicState` navigation section.
    Individual quantum state objects (e.g., `SphericalSymmetryState`) are referenced within the basis_orbitals list of ElectronicState.
    """

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

    n_orbitals = Quantity(
        type=np.int32,
        description="""
        Number of orbitals in the tight-binding model. The `entity_ref` reference is used to refer to
        the `ElectronicState` section, which navigates to the relevant basis orbitals (e.g., `SphericalSymmetryState`).
        """,
    )

    degeneracy_factors = Quantity(
        type=np.int32,
        shape=['*'],
        description="""
        Degeneracy of each Wigner-Seitz point.
        """,
    )

    value = Quantity(
        type=np.complex128,
        unit='joule',
        description="""
        Value of the hopping matrix in joules. The elements are complex numbers defined for each Wigner-Seitz point and
        each pair of orbitals. Note this contains also the onsite values, i.e., it includes the Wigner-Seitz point (0, 0, 0), hence the `CrystalFieldSplitting` values.
        """,
    )

    # TODO add normalization to extract DOS, band structure, etc, properties from `HoppingMatrix`


class CrystalFieldSplitting(PhysicalProperty):
    """
    Energy difference between the degenerated orbitals of an ion in a crystal field environment.
    """

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

    n_orbitals = Quantity(
        type=np.int32,
        description="""
        Number of orbitals in the tight-binding model. The `entity_ref` reference is used to refer to
        the `ElectronicState` section, which navigates to the relevant basis orbitals (e.g., `SphericalSymmetryState`).
        """,
    )

    value = Quantity(
        type=np.float64,
        unit='joule',
        description="""
        Value of the crystal field splittings in joules. This is the intra-orbital local contribution, i.e., the same orbital
        at the same Wigner-Seitz point (0, 0, 0).
        """,
    )
