Source code for solidipes_core_plugin.reports.widgets.validation

"""Widget to display validations on a DataContainer. Must be lazy-loaded to avoid loading Streamlit."""

import time
from typing import Literal, Union
from uuid import uuid4

import pandas as pd
import streamlit as st  # Cannot lazy import because of st.fragment decorator
from solidipes.loaders.data_container import DataContainer
from solidipes.validators.global_validation import DatasetProxy, get_global_validator
from solidipes.validators.ontology import (
    OntologyValidator,
    PydanticOntology,
    get_available_ontology_identifiers,
    get_ontology_identifier,
    set_ontology_identifier,
)

from solidipes_core_plugin.reports.widgets.solidipes_widget import SolidipesWidget

VALIDATION_RELOAD_SLEEP_TIME = 1.5




class ValidationWidget(SolidipesWidget):
    def __init__(
        self,
        data_container: Union[DataContainer, DatasetProxy],
        ontology_level: Literal["global", "data_container"],
        **kwargs,
    ):
        super().__init__(**kwargs)
        if self.layout is st:
            self.layout = st.container()
        self.data_container = data_container
        self.ontology_level = ontology_level

        with self.layout:
            self.display()



    @st.fragment
    def display(self):
        validation_results = self.data_container.validation_results
        dataset_is_valid = self.data_container.get_is_valid(validation_results=validation_results)
        title = ":white_check_mark:" if dataset_is_valid else ":no_entry_sign:"
        title += " Validations"

        with st.expander(title, expanded=(not dataset_is_valid or len(self.data_container.errors) > 0)):
            self.display_validations(validation_results)
            if self.ontology_level == "global":
                self.display_global_ontology_identifier()
            else:
                self.display_data_container_ontology_identifier()




    def display_validations(self, validation_results):
        data = []
        for validation_result in validation_results:
            validator = validation_result.validator
            data.append({
                "name": validator.name,
                "enabled": self.data_container.validator_enabled[validator.name],
                "mandatory": validator.mandatory,
                "description": validator.description,
                "valid": validation_result.valid,
                "manually_settable": validator.manually_settable,
                "errors": "\n\n".join(validation_result.errors),
            })
        df = pd.DataFrame(data)

        if "validations_editor_key" not in st.session_state:
            st.session_state["validations_editor_key"] = uuid4()

        edited_df = st.data_editor(
            df,
            hide_index=True,
            use_container_width=True,
            column_order=["enabled", "mandatory", "description", "valid", "manually_settable", "errors"],
            column_config={
                "enabled": st.column_config.CheckboxColumn("Enabled", required=False),
                "mandatory": "Mandatory",
                "description": "Description",
                "valid": "Valid",
                "manually_settable": "Manually settable",
                "errors": "Errors (double-click to expand)",
            },
            disabled=df.columns[2:],
            key=st.session_state["validations_editor_key"],
        )

        edited_lines = edited_df[df["enabled"] != edited_df["enabled"]]
        for i, row in edited_lines.iterrows():
            if row["enabled"]:
                self.data_container.enable_validator(row["name"])
                st.rerun(scope="fragment")

            else:
                try:
                    self.data_container.disable_validator(row["name"])

                except ValueError:
                    st.error(f'Cannot disable mandatory validator "{row["description"]}"')
                    st.session_state["validations_editor_key"] = uuid4()  # Force reload of data_editor
                    time.sleep(VALIDATION_RELOAD_SLEEP_TIME)

                finally:
                    st.rerun(scope="fragment")




    def display_global_ontology_identifier(self):
        ontology_identifier = get_ontology_identifier() or "solidipes.ontologies.solidipes"
        ontology_identifiers = get_available_ontology_identifiers()
        if ontology_identifier not in ontology_identifiers:
            ontology_identifiers.append(ontology_identifier)

        global_ontology_selection_key = "global_ontology_selection"

        with st.form(global_ontology_selection_key + "_form"):
            new_ontology_identifier = (
                st.selectbox(
                    label="Ontology identifier (module, or path in dataset)",
                    options=ontology_identifiers,
                    index=ontology_identifiers.index(ontology_identifier),
                    accept_new_options=True,
                )
                or ""
            ).strip()

            st.form_submit_button("Replace ontology")

        if new_ontology_identifier != ontology_identifier:
            set_ontology_identifier(new_ontology_identifier)
            st.rerun()

        # st.write(f"ontology_identifier: `{get_ontology_identifier()}`")
        # st.write(f"new_ontology_identifier: `{new_ontology_identifier}`")

        ontology_validator = get_global_validator(OntologyValidator)
        if isinstance(ontology_validator.ontology, PydanticOntology) and isinstance(
            module_load_status := ontology_validator.ontology.module_load_status, PydanticOntology.ModuleLoadError
        ):
            st.caption(module_load_status.error_message)




    def display_data_container_ontology_identifier(self):
        ontology_identifier = get_ontology_identifier() or "Default Solidipes ontology"
        st.write(f"Ontology: `{ontology_identifier}`")

        ontology_validator = get_global_validator(OntologyValidator)
        ontology_class = ontology_validator.ontology.get_file_class_name(self.data_container)
        ontology_classes = ontology_validator.ontology.get_file_compatible_class_names(self.data_container)
        if ontology_class not in ontology_classes:
            ontology_classes.append(ontology_class)

        ontology_class_selection_key = f"ontology_class_selection_{self.data_container.unique_identifier}"
        if ontology_class_selection_key not in st.session_state:
            st.session_state[ontology_class_selection_key] = ontology_class
        if st.session_state[ontology_class_selection_key] not in ontology_classes:
            ontology_classes.append(st.session_state[ontology_class_selection_key])

        new_ontology_class = st.selectbox(
            "Ontology class",
            options=ontology_classes,
            index=ontology_classes.index(st.session_state[ontology_class_selection_key]),
        )

        # st.write(f"obj.ontology_class: `{self.data_container.ontology_class}`")
        # st.write(f"get_file_class_name: `{ontology_class}`")
        # st.write(f"Selected class: `{new_ontology_class}`")

        if new_ontology_class != ontology_class:
            self.data_container.ontology_class = new_ontology_class
            st.rerun()