""" Core definition of an Electronic Structure """

from __future__ import annotations

from collections import defaultdict
from datetime import datetime
from enum import Enum
from math import isnan
from typing import Dict, List, Optional, Type, TypeVar, Union

import numpy as np
from pydantic import BaseModel, Field
from pymatgen.analysis.magnetism.analyzer import (
    CollinearMagneticStructureAnalyzer,
    Ordering,
)
from pymatgen.core import Structure
from pymatgen.core.periodic_table import Element
from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine
from pymatgen.electronic_structure.core import OrbitalType, Spin
from pymatgen.electronic_structure.dos import CompleteDos
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.symmetry.bandstructure import HighSymmKpath
from typing_extensions import Literal

from emmet.core.material_property import PropertyDoc
from emmet.core.mpid import MPID
from emmet.core.settings import EmmetSettings

SETTINGS = EmmetSettings()


class BSPathType(Enum):
    setyawan_curtarolo = "setyawan_curtarolo"
    hinuma = "hinuma"
    latimer_munro = "latimer_munro"


class DOSProjectionType(Enum):
    total = "total"
    elemental = "elemental"
    orbital = "orbital"


class BSObjectDoc(BaseModel):
    """
    Band object document.
    """

    task_id: Optional[MPID] = Field(
        None,
        description="The source calculation (task) ID that this band structure comes from. "
        "This has the same form as a Materials Project ID.",
    )

    last_updated: datetime = Field(
        description="The timestamp when this calculation was last updated",
        default_factory=datetime.utcnow,
    )

    data: Optional[Union[Dict, BandStructureSymmLine]] = Field(
        None, description="The band structure object for the given calculation ID"
    )


class DOSObjectDoc(BaseModel):
    """
    DOS object document.
    """

    task_id: Optional[MPID] = Field(
        None,
        description="The source calculation (task) ID that this density of states comes from. "
        "This has the same form as a Materials Project ID.",
    )

    last_updated: datetime = Field(
        description="The timestamp when this calculation was last updated.",
        default_factory=datetime.utcnow,
    )

    data: Optional[CompleteDos] = Field(
        None, description="The density of states object for the given calculation ID."
    )


class ElectronicStructureBaseData(BaseModel):
    task_id: MPID = Field(
        ...,
        description="The source calculation (task) ID for the electronic structure data. "
        "This has the same form as a Materials Project ID.",
    )

    band_gap: float = Field(..., description="Band gap energy in eV.")

    cbm: Optional[Union[float, Dict]] = Field(
        None, description="Conduction band minimum data."
    )

    vbm: Optional[Union[float, Dict]] = Field(
        None, description="Valence band maximum data."
    )

    efermi: Optional[float] = Field(None, description="Fermi energy in eV.")


class ElectronicStructureSummary(ElectronicStructureBaseData):
    is_gap_direct: bool = Field(..., description="Whether the band gap is direct.")

    is_metal: bool = Field(..., description="Whether the material is a metal.")

    magnetic_ordering: Union[str, Ordering] = Field(
        ..., description="Magnetic ordering of the calculation."
    )


class BandStructureSummaryData(ElectronicStructureSummary):
    nbands: float = Field(..., description="Number of bands.")

    equivalent_labels: Dict = Field(
        ..., description="Equivalent k-point labels in other k-path conventions."
    )

    direct_gap: float = Field(..., description="Direct gap energy in eV.")


class DosSummaryData(ElectronicStructureBaseData):
    spin_polarization: Optional[float] = Field(
        None, description="Spin polarization at the fermi level."
    )


class BandstructureData(BaseModel):
    setyawan_curtarolo: Optional[BandStructureSummaryData] = Field(
        None,
        description="Band structure summary data using the Setyawan-Curtarolo path convention.",
    )

    hinuma: Optional[BandStructureSummaryData] = Field(
        None,
        description="Band structure summary data using the Hinuma et al. path convention.",
    )

    latimer_munro: Optional[BandStructureSummaryData] = Field(
        None,
        description="Band structure summary data using the Latimer-Munro path convention.",
    )


class DosData(BaseModel):
    total: Optional[Dict[Union[Spin, str], DosSummaryData]] = Field(
        None, description="Total DOS summary data."
    )

    elemental: Optional[
        Dict[
            Element,
            Dict[
                Union[Literal["total", "s", "p", "d", "f"], OrbitalType],
                Dict[Union[Literal["1", "-1"], Spin], DosSummaryData],
            ],
        ]
    ] = Field(
        None,
        description="Band structure summary data using the Hinuma et al. path convention.",
    )

    orbital: Optional[
        Dict[
            Union[Literal["total", "s", "p", "d", "f"], OrbitalType],
            Dict[Union[Literal["1", "-1"], Spin], DosSummaryData],
        ]
    ] = Field(
        None,
        description="Band structure summary data using the Latimer-Munro path convention.",
    )

    magnetic_ordering: Optional[Union[Ordering, str]] = Field(
        None, description="Magnetic ordering of the calculation."
    )


T = TypeVar("T", bound="ElectronicStructureDoc")


class ElectronicStructureDoc(PropertyDoc, ElectronicStructureSummary):
    """
    Definition for a core Electronic Structure Document
    """

    property_name: str = "electronic_structure"

    bandstructure: Optional[BandstructureData] = Field(
        None, description="Band structure data for the material."
    )

    dos: Optional[DosData] = Field(
        None, description="Density of states data for the material."
    )

    last_updated: datetime = Field(
        description="Timestamp for when this document was last updated.",
        default_factory=datetime.utcnow,
    )

    @classmethod
    def from_bsdos(  # type: ignore[override]
        cls: Type[T],
        material_id: MPID,
        dos: Dict[MPID, CompleteDos],
        is_gap_direct: bool,
        is_metal: bool,
        origins: List[dict] = [],
        structures: Optional[Dict[MPID, Structure]] = None,
        setyawan_curtarolo: Optional[Dict[MPID, BandStructureSymmLine]] = None,
        hinuma: Optional[Dict[MPID, BandStructureSymmLine]] = None,
        latimer_munro: Optional[Dict[MPID, BandStructureSymmLine]] = None,
        **kwargs,
    ) -> T:
        """
        Builds a electronic structure document using band structure and density of states data.

        Args:
            material_id (MPID): A material ID.
            dos (Dict[MPID, CompleteDos]): Dictionary mapping a calculation (task) ID to a CompleteDos object.
            is_gap_direct (bool): Direct gap indicator included at root level of document.
            is_metal (bool): Metallic indicator included at root level of document.
            structures (Dict[MPID, Structure]) = Dictionary mapping a calculation (task) ID to the structures used
                as inputs. This is to ensures correct magnetic moment information is included.
            setyawan_curtarolo (Dict[MPID, BandStructureSymmLine]): Dictionary mapping a calculation (task) ID to a
                BandStructureSymmLine object from a calculation run using the Setyawan-Curtarolo k-path convention.
            hinuma (Dict[MPID, BandStructureSymmLine]): Dictionary mapping a calculation (task) ID to a
                BandStructureSymmLine object from a calculation run using the Hinuma et al. k-path convention.
            latimer_munro (Dict[MPID, BandStructureSymmLine]): Dictionary mapping a calculation (task) ID to a
                BandStructureSymmLine object from a calculation run using the Latimer-Munro k-path convention.
            origins (List[dict]): Optional origins information for final doc

        """

        # -- Process density of states data

        dos_task, dos_obj = list(dos.items())[0]

        orbitals = [OrbitalType.s, OrbitalType.p, OrbitalType.d]
        spins = list(dos_obj.densities.keys())

        ele_dos = dos_obj.get_element_dos()
        tot_orb_dos = dos_obj.get_spd_dos()

        elements = ele_dos.keys()

        dos_efermi = dos_obj.efermi

        is_gap_direct = is_gap_direct
        is_metal = is_metal

        structure = dos_obj.structure

        if structures is not None and structures[dos_task]:
            structure = structures[dos_task]

        dos_mag_ordering = CollinearMagneticStructureAnalyzer(structure).ordering

        dos_data = {
            "total": defaultdict(dict),
            "elemental": {element: defaultdict(dict) for element in elements},
            "orbital": defaultdict(dict),
            "magnetic_ordering": dos_mag_ordering,
        }

        for spin in spins:
            # - Process total DOS data
            band_gap = dos_obj.get_gap(spin=spin)
            (cbm, vbm) = dos_obj.get_cbm_vbm(spin=spin)

            try:
                spin_polarization = dos_obj.spin_polarization
                if spin_polarization is None or isnan(spin_polarization):
                    spin_polarization = None
            except KeyError:
                spin_polarization = None

            dos_data["total"][spin] = DosSummaryData(  # type: ignore[index]
                task_id=dos_task,
                band_gap=band_gap,
                cbm=cbm,
                vbm=vbm,
                efermi=dos_efermi,
                spin_polarization=spin_polarization,
            )

            # - Process total orbital projection data
            for orbital in orbitals:
                band_gap = tot_orb_dos[orbital].get_gap(spin=spin)

                (cbm, vbm) = tot_orb_dos[orbital].get_cbm_vbm(spin=spin)

                spin_polarization = None

                dos_data["orbital"][orbital][spin] = DosSummaryData(  # type: ignore[index]
                    task_id=dos_task,
                    band_gap=band_gap,
                    cbm=cbm,
                    vbm=vbm,
                    efermi=dos_efermi,
                    spin_polarization=spin_polarization,
                )

        # - Process element and element orbital projection data
        for ele in ele_dos:
            orb_dos = dos_obj.get_element_spd_dos(ele)

            for orbital in ["total"] + list(orb_dos.keys()):  # type: ignore[assignment]
                if orbital == "total":
                    proj_dos = ele_dos
                    label = ele
                else:
                    proj_dos = orb_dos
                    label = orbital

                for spin in spins:
                    band_gap = proj_dos[label].get_gap(spin=spin)
                    (cbm, vbm) = proj_dos[label].get_cbm_vbm(spin=spin)

                    spin_polarization = None

                    dos_data["elemental"][ele][orbital][spin] = DosSummaryData(  # type: ignore[index]
                        task_id=dos_task,
                        band_gap=band_gap,
                        cbm=cbm,
                        vbm=vbm,
                        efermi=dos_efermi,
                        spin_polarization=spin_polarization,
                    )

        #  -- Process band structure data
        bs_data = {  # type: ignore
            "setyawan_curtarolo": setyawan_curtarolo,
            "hinuma": hinuma,
            "latimer_munro": latimer_munro,
        }

        for bs_type, bs_input in bs_data.items():
            if bs_input is not None:
                bs_task, bs = list(bs_input.items())[0]

                if structures is not None and structures[bs_task]:
                    bs_mag_ordering = CollinearMagneticStructureAnalyzer(
                        structures[bs_task]
                    ).ordering
                else:
                    bs_mag_ordering = CollinearMagneticStructureAnalyzer(
                        bs.structure  # type: ignore[arg-type]
                    ).ordering

                gap_dict = bs.get_band_gap()
                is_metal = bs.is_metal()
                direct_gap = bs.get_direct_band_gap()

                if is_metal:
                    band_gap = 0.0
                    cbm = None  # type: ignore[assignment]
                    vbm = None  # type: ignore[assignment]
                    is_gap_direct = False
                else:
                    band_gap = gap_dict["energy"]
                    cbm = bs.get_cbm()  # type: ignore[assignment]
                    vbm = bs.get_vbm()  # type: ignore[assignment]
                    is_gap_direct = gap_dict["direct"]

                bs_efermi = bs.efermi
                nbands = bs.nb_bands

                # - Get equivalent labels between different conventions
                hskp = HighSymmKpath(
                    bs.structure,
                    path_type="all",
                    symprec=0.1,
                    angle_tolerance=5,
                    atol=1e-5,
                )
                equivalent_labels = hskp.equiv_labels

                if bs_type == "latimer_munro":
                    gen_labels = set(
                        [
                            label
                            for label in equivalent_labels["latimer_munro"][
                                "setyawan_curtarolo"
                            ]
                        ]
                    )
                    kpath_labels = set(
                        [
                            kpoint.label
                            for kpoint in bs.kpoints
                            if kpoint.label is not None
                        ]
                    )

                    if not gen_labels.issubset(kpath_labels):
                        new_structure = SpacegroupAnalyzer(
                            bs.structure  # type: ignore[arg-type]
                        ).get_primitive_standard_structure(
                            international_monoclinic=False
                        )

                        hskp = HighSymmKpath(
                            new_structure,
                            path_type="all",
                            symprec=SETTINGS.SYMPREC,
                            angle_tolerance=SETTINGS.ANGLE_TOL,
                            atol=1e-5,
                        )
                        equivalent_labels = hskp.equiv_labels

                bs_data[bs_type] = BandStructureSummaryData(  # type: ignore
                    task_id=bs_task,
                    band_gap=band_gap,
                    direct_gap=direct_gap,
                    cbm=cbm,
                    vbm=vbm,
                    is_gap_direct=is_gap_direct,
                    is_metal=is_metal,
                    efermi=bs_efermi,
                    nbands=nbands,
                    equivalent_labels=equivalent_labels,
                    magnetic_ordering=bs_mag_ordering,
                )

        bs_entry = BandstructureData(**bs_data)  # type: ignore
        dos_entry = DosData(**dos_data)  # type: ignore[arg-type]

        # Obtain summary data

        bs_gap = (
            bs_entry.setyawan_curtarolo.band_gap
            if bs_entry.setyawan_curtarolo is not None
            else None
        )
        dos_cbm, dos_vbm = dos_obj.get_cbm_vbm()
        dos_gap = max(dos_cbm - dos_vbm, 0.0)

        new_origin_last_updated = None
        new_origin_task_id = None

        if bs_gap is not None and bs_gap <= dos_gap + 0.2:
            summary_task = bs_entry.setyawan_curtarolo.task_id  # type: ignore
            summary_band_gap = bs_gap
            summary_cbm = (
                bs_entry.setyawan_curtarolo.cbm.get("energy", None)  # type: ignore
                if bs_entry.setyawan_curtarolo.cbm is not None  # type: ignore
                else None
            )
            summary_vbm = (
                bs_entry.setyawan_curtarolo.vbm.get("energy", None)  # type: ignore
                if bs_entry.setyawan_curtarolo.cbm is not None  # type: ignore
                else None
            )  # type: ignore
            summary_efermi = bs_entry.setyawan_curtarolo.efermi  # type: ignore
            is_gap_direct = bs_entry.setyawan_curtarolo.is_gap_direct  # type: ignore
            is_metal = bs_entry.setyawan_curtarolo.is_metal  # type: ignore
            summary_magnetic_ordering = bs_entry.setyawan_curtarolo.magnetic_ordering  # type: ignore

            for origin in origins:
                if origin["name"] == "setyawan_curtarolo":
                    new_origin_last_updated = origin["last_updated"]
                    new_origin_task_id = origin["task_id"]

        else:
            summary_task = dos_entry.model_dump()["total"][Spin.up]["task_id"]
            summary_band_gap = dos_gap
            summary_cbm = dos_cbm
            summary_vbm = dos_vbm
            summary_efermi = dos_efermi
            summary_magnetic_ordering = dos_mag_ordering
            is_metal = True if np.isclose(dos_gap, 0.0, atol=0.01, rtol=0) else False

            for origin in origins:
                if origin["name"] == "dos":
                    new_origin_last_updated = origin["last_updated"]
                    new_origin_task_id = origin["task_id"]

        if new_origin_task_id is not None:
            for origin in origins:
                if origin["name"] == "electronic_structure":
                    origin["last_updated"] = new_origin_last_updated
                    origin["task_id"] = new_origin_task_id

        return cls.from_structure(
            material_id=MPID(material_id),
            task_id=summary_task,
            meta_structure=structure,
            band_gap=summary_band_gap,
            cbm=summary_cbm,
            vbm=summary_vbm,
            efermi=summary_efermi,
            is_gap_direct=is_gap_direct,
            is_metal=is_metal,
            magnetic_ordering=summary_magnetic_ordering,
            bandstructure=bs_entry,
            dos=dos_entry,
            **kwargs,
        )