# -*- coding: utf-8 -*-

from rich.progress import track

from simmate.database.core import table_column
from simmate.database.mixins import Structure, ThirdPartyData
from simmate.toolkit import Structure as ToolkitStructure


class JarvisStructure(ThirdPartyData, Structure):
    """
    Crystal structures from the [JARVIS](https://jarvis.nist.gov/) database.

    Currently, this table only stores strucure and reported energy above hull.
    The calculated energy is not reported, so the Thermodynamics mixin is not used.
    """

    class Meta:
        db_table = "jarvis__structures"

    # -------------------------------------------------------------------------

    html_display_name = "JARVIS"
    html_description_short = (
        "A collection of materials data for high-throughput discovery from "
        "NIST's JARVIS project. This dataset focuses on 2D and 3D materials "
        "with high-fidelity calculations of electronic, optical, and "
        "mechanical properties."
    )

    html_entries_template = "jarvis/structures/table.html"
    html_entry_template = "jarvis/structures/view.html"

    # -------------------------------------------------------------------------

    external_website = "https://jarvis.nist.gov/"
    source_doi = "https://doi.org/10.1038/s41524-020-00440-1"
    is_redistribution_allowed = True

    @property
    def external_link(self) -> str:
        """
        URL to this structure in the JARVIS website.
        """
        # All JARVIS structures have their data mapped to a URL in the same way
        # ex: https://www.ctcms.nist.gov/~knc6/static/JARVIS-DFT/JVASP-1234.xml
        # we store the id as "jvasp-123" so we need to convert this to uppercase
        id = self.id.upper()
        return f"https://www.ctcms.nist.gov/~knc6/static/JARVIS-DFT/{id}.xml"

    # -------------------------------------------------------------------------

    remote_archive_link = "https://archives.simmate.org/JarvisStructure-2026-03-20.zip"
    archive_fields = ["energy_above_hull"]

    # -------------------------------------------------------------------------

    id = table_column.CharField(max_length=25, primary_key=True)
    """
    The id used to represent the structure (ex: "jvasp-12345")
    """

    # TODO: contact their team to ask about reporting energy instead. That way
    # we can use the Thermodynamics mixin instead of manually listing this.
    energy_above_hull = table_column.FloatField(blank=True, null=True)
    """
    The energy above hull, as reported by the JARVIS database (no units given)
    """

    # -------------------------------------------------------------------------

    @classmethod
    def load_source_data(cls):
        """
        This method pulls JARVIS data into the Simmate database.

        JARVIS has a python package "jarvis-tools" that let's us pull some of
        their database dumps. For instructions on how to do this, they provided
        [this link](https://colab.research.google.com/github/knc6/jarvis-tools-notebooks/blob/master/jarvis-tools-notebooks/Get_JARVIS_DFT_final_structures_in_ASE_or_Pymatgen_format.ipynb)

        Alternatively, we could manually download
        [their database json files](https://jarvis-materials-design.github.io/dbdocs/thedownloads/). We specifically look at the "3D-materials curated data".
        """

        # Jarvis is not a dependency of simmate, so make sure you install it
        # before using this method
        try:
            from jarvis.db.figshare import data as jarvis_helper
        except:
            raise ModuleNotFoundError(
                "You must install jarvis with `conda install -c conda-forge jarvis-tools`"
            )

        # Load all of the 3D data from JARVIS. This gives a list of dictionaries
        data = jarvis_helper("dft_3d")
        # TODO: In the future, we can add other datasets like the 2D dataset.

        # Now iterate through all the data -- which is a list of dictionaries.
        # We convert the data into a pymatgen object and sanitize it before
        # saving to the Simmate database
        db_objects = []
        for entry in track(data):
            # The structure is in the atoms field as a dictionary. We pull
            # this data out and convert it to a pymatgen Structure object
            try:
                structure = ToolkitStructure(
                    lattice=entry["atoms"]["lattice_mat"],
                    species=entry["atoms"]["elements"],
                    coords=entry["atoms"]["coords"],
                    coords_are_cartesian=entry["atoms"]["cartesian"],
                )

                # now convert the entry to a database object
                structure_db = cls.from_toolkit(
                    id=entry["jid"].lower(),
                    structure=structure,
                    energy_above_hull=(
                        entry["ehull"] if entry["ehull"] != "na" else None
                    ),
                )
            except Exception:
                structure_db = cls.from_toolkit(
                    id=entry["jid"].lower(),
                    structure=None,
                    energy_above_hull=(
                        entry["ehull"] if entry["ehull"] != "na" else None
                    ),
                    is_invalid_structure=True,
                )

            db_objects.append(structure_db)

        # and save it to our database
        cls.objects.bulk_create(
            db_objects,
            batch_size=1_000,
            ignore_conflicts=True,
        )
