from __future__ import annotations

import panel as pn

from typing import List
from pydantic import Field, root_validator
from abipy.abio.inputs import AbinitInput
from abipy.panels.core import ply
from abipy.panels.viewers import JSONViewer
from abipy.abio.factories import ebands_input
from abipy.flowtk import TaskManager, Flow
from abipy.flowtk.flows import bandstructure_flow
from .base_models import MongoConnector, GfsFileDesc
from .flow_models import FlowModel, PresetQuery
from .gs_models import GsData, NscfData, RelaxData



class _BaseEbandsFlowModel(FlowModel):
    """
    This is the base class for models performing band structure calculations.
    It defines the three input files for the SCF/NSCF/DOS calculation
    as well as submodels storing the output results.
    """

    # These inputs can be generated either manually or from meta-parameters + factory functions.

    scf_input: AbinitInput = Field(None, description="Abinit Input file generated by AbiPy.")

    nscf_input: AbinitInput = Field(None, description="Abinit Input file generated by AbiPy.")

    dos_input: AbinitInput = Field(None, description="Abinit Input file generated by AbiPy.")

    ########
    # Output
    ########

    scf_data: GsData = Field(None, description="Results produced by the SCF run.")

    nscf_kpath_data: NscfData = Field(None, description="Results produced by the NSCF run with a k-path.")

    nscf_kmesh_data: NscfData = Field(None, description="Results produced by the NSCF run with a k-mesh.")

    with_gsr: bool = Field(False, description="Set it to True to save the GSR file in GridFS.")

    #with_den: bool = Field(False, description="Set it to True to save the DEN file in GridFS.")

    def postprocess_flow(self, flow: Flow, mng_connector: MongoConnector) -> None:
        """
        Analyze the flow and fills the model with output results.
        MongoConnector should be used only to insert files in GridFs as the final insertion is done by the caller.
        This function is called by the AbiPy Worker if the flow completed successfully.
        """
        with flow[0][0].open_gsr() as gsr:
            self.scf_data = GsData.from_gsr(gsr, mng_connector, self.with_gsr)

        with flow[0][1].open_gsr() as gsr:
            self.nscf_kpath_data = NscfData.from_gsr(gsr, mng_connector, self.with_gsr)

        if self.dos_input is not None:
            with flow[0][2].open_gsr() as gsr:
                self.nscf_kmesh_data = NscfData.from_gsr(gsr, mng_connector, self.with_gsr)

    def get_panel_view(self):
        """
        Return panel object with a view of the model.
        """
        title = self.in_structure_data.get_title()
        structure = self.in_structure_data.structure
        a, b, c = structure.lattice.abc
        alpha, beta, gamma = structure.lattice.angles
        header = f"""\
## {title}

- Lattice lengths: a = {a:.6f}, b = {b:.6f}, c = {c:.6f} Ang
- Lattice angles: α = {alpha:.3f}, β = {beta:.3f}, ɣ = {gamma:.3f} degrees
"""

        if self.is_completed:
            header += f"""\
- Pressure: {self.scf_data.pressure_gpa:.3f} GPa
- Max |Force|: {self.scf_data.max_force_ev_over_ang:.8f} eV/Ang
- Energy: {self.scf_data.energy:.4f} eV
- Energy per atom: {self.scf_data.energy_per_atom:.4f} eV
"""
            plotly_bands = ply(self.nscf_kpath_data.ebands.plotly(show=False))
        else:
            plotly_bands = pn.pane.Alert(f"Bands are not available because exec_status is `{self.flow_data.exec_status}`")

        return pn.Column(
            #self.in_structure_data.get_panel_view(),
            header,
            pn.Row(
                plotly_bands,
                pn.pane.HTML(self.scf_input._repr_html_()),
            ),
            "### MongoDB Document",
            JSONViewer(self.json(), depth=1),
            pn.layout.Divider(),
            sizing_mode="stretch_both",
        )

    @classmethod
    def get_preset_queries(cls) -> List[PresetQuery]:
        """
        Return list of dictionaries with the MongoDB queries typically used to filter documents for this model.
        Empty list if no suggestion is available.
        """
        return [
            PresetQuery.for_large_forces_or_high_pressure("scf_data", cls)
        ]

    #@classmethod
    #def mongo_aggregate_egaps(cls, collection: Collection) -> pd.DataFrame:
        #oids = cls.mng_find_completed_oids(collection)
        #    projection = [
        #        "in_structure_data
        #    }
        #
        #    fund_gap_projection = [
        #        "scf_data.fundamental_gap",
        #        "nscf_kpath_data.fundamental_gap",
        #        "nscf_kmesh_data.fundamental_gap",
        #    ]

        #    direct_gap_projection = {
        #        "scf_data.direct_gap",
        #        "nscf_kpath_data.direct_gap",
        #        "nscf_kmesh_data.direct_gap",
        #    }

        #    projection.extend(fundamental_gap_projection + direct_gap_projection)

        #    oids = cls.mng_find_completed_oids(collection)
        #    rows = []
        #    for oid in oids:
        #        doc = collection.find_one({"_id": oid}, projection)
        #        structure_data = AbipyDecoder().process_decoded(doc["in_structure_data"])
        #        #structure_data = cls.decode(doc["in_structure_data"])
        #        row = structure_data.dict4pandas()
        #        # Here I need a tool to access foo.bar instead of d["foo"]["bar"]
        #        row["fund_gap"] = min(doc[key] for key in fundamental_gap_projection)
        #        row["direct_gap"] = min(doc[key] for key in direct_gap_projection)
        #        rows.append(row)

        #    return pd.DataFrame(rows)


class EbandsFlowModelWithParams(_BaseEbandsFlowModel):
    """
    This model defines the input arguments used to build a Flow for band structure calculations
    as well as the sub models used to store the final results.
    """

    ##################
    # Input parameters
    ##################
    kppa: int = Field(1000, description="Defines the sampling used for the SCF run. Defaults to 1000 if not given.")

    ndivsm: int = Field(2, description="Number of divisions used to sample the smallest segment of the k-path.")

    spin_mode: str = Field("unpolarized", description="Spin polarization")

    charge: float = Field(0.0, description="Electronic charge added to the unit cell.")

    smearing: str = Field("fermi_dirac:0.1 eV", description="Smearing technique.")

    dos_kppa: int = Field(None,
                          description="Scalar or List of integers with the number of k-points per atom " +
                                      "to be used for the computation of the DOS (None if DOS is not wanted")

    paral_kgb: int = Field(0, description="")

    # validators
    #_normalize_name = validator('name', allow_reuse=True)(normalize)
    #_normalize_name = reuse_validator('name')(normalize)

    def build_flow(self, workdir: str, manager: TaskManager) -> Flow:
        """
        Build an AbiPy Flow using the input data available in the model and return it.

        Args:
            workdir: Working directory provided by the caller.
            manager: |TaskManager| object.

        Return: |Flow| object.
        """
        pseudos = self.pseudos_specs.get_pseudos()

        multi = ebands_input(self.in_structure_data.structure, pseudos,
                             kppa=self.kppa, nscf_nband=None, ndivsm=self.ndivsm,
                             #ecut=6, pawecutdg=None,
                             scf_nband=None, accuracy="normal",
                             spin_mode=self.spin_mode,
                             smearing=self.smearing, charge=self.charge,
                             scf_algorithm=None, dos_kppa=self.dos_kppa,
                             )

        multi.set_vars(paral_kgb=self.paral_kgb)

        if self.dos_kppa is not None:
            self.scf_input, self.nscf_input, self.dos_input = multi.split_datasets()
        else:
            self.scf_input, self.nscf_input = multi.split_datasets()

        return bandstructure_flow(workdir, self.scf_input, self.nscf_input,
                                  dos_inputs=self.dos_input, manager=manager)


class EbandsFlowModelWithInputs(_BaseEbandsFlowModel):
    """
    More flexible model for band structure calculations that requires |AbinitInput| objects
    generated by the user before building the model.
    """

    @root_validator
    def check_inputs(cls, values):
        scf_input, nscf_input = values.get("scf_input"), values.get("nscf_input")
        if scf_input is None or nscf_input is None:
            raise ValueError(f"Building a {cls.__name__} model requires both scf_input and nscf_input")
        return values

    def build_flow(self, workdir: str, manager: TaskManager) -> Flow:
        """
        Build an AbiPy Flow using the input data available in the model and return it.

        Args:
            workdir: Working directory provided by the caller.
            manager: |TaskManager| object.

        Return: |Flow| object.
        """
        return bandstructure_flow(workdir, self.scf_input, self.nscf_input,
                                  dos_inputs=self.dos_input, manager=manager)
