from datetime import datetime
from pathlib import Path
from typing import List

import numpy as np
import pandas
import pandas as pd
from qgis.gui import QgsCollapsibleGroupBox
from qgis.PyQt.QtWidgets import (
    QComboBox,
    QGridLayout,
    QLabel,
    QMessageBox,
    QTableWidget,
    QVBoxLayout,
    QWidget,
    QWizardPage,
)

from openlog.core import pint_utilities
from openlog.core.assay_interpolation import (
    ExtendedAssayColumn,
    ExtendedAssayInterpolation,
    GapDataResolution,
    OverlapDataResolution,
)
from openlog.datamodel.assay.detection_limit import AssayColumnDetectionLimit
from openlog.datamodel.assay.generic_assay import (
    AssayColumn,
    AssayDataExtent,
    AssayDefinition,
    AssayDomainType,
    AssaySeriesType,
    GenericAssay,
)
from openlog.datamodel.assay.spherical_assay import SphericalDefinition
from openlog.datamodel.assay.uncertainty import AssayColumnUncertainty
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.gui.import_assay.utils import sanitize_sql_name
from openlog.gui.import_assay.wizard.assay_selection import AssaySelectionPageWizard
from openlog.gui.utils.column_definition import ColumnDefinition
from openlog.gui.utils.delimited_text_import_widget import DelimitedTextImportWidget
from openlog.gui.utils.mdl_assay_column_categorie import CategoryPolicy

BASE_SETTINGS_KEY = "/OpenLog/gui/import/assay"


class AssayImportPageWizard(QWizardPage):
    def __init__(
        self,
        parent: QWidget,
        openlog_connection: OpenLogConnection,
        assay_selection_wizard: AssaySelectionPageWizard,
        person_dataset_import_wizard,
        restore_settings: bool = True,
    ) -> None:
        """
        QWizard to define assay import data

        Args:
            assay_selection_wizard : AssaySelectionPageWizard defining assay
            parent : QWidget parent
        """
        super().__init__(parent)
        self._openlog_connection = openlog_connection

        self._assay_selection_wizard = assay_selection_wizard
        self._person_dataset_import_wizard = person_dataset_import_wizard

        self.setTitle("Downhole data values import")
        label = QLabel("Select a .csv file or a QGIS layer in Source table section.")
        label.setWordWrap(True)
        self.label = label
        layout = QVBoxLayout()
        layout.addWidget(label)
        gb = QgsCollapsibleGroupBox("Import options", self)
        gb.setCollapsed(True)
        layout_gb = QGridLayout()
        gb.setLayout(layout_gb)

        # comboboxes for gap/overlap resolution
        self.gap_option_bg = QComboBox(self)
        for o in GapDataResolution:
            self.gap_option_bg.addItem(o.value)
        layout_gb.addWidget(QLabel("Gap resolution"), 0, 0)
        layout_gb.addWidget(self.gap_option_bg, 0, 1)

        self.overlap_option_bg = QComboBox(self)
        for o in OverlapDataResolution:
            self.overlap_option_bg.addItem(o.value)
        layout_gb.addWidget(QLabel("Overlap resolution"), 1, 0)
        layout_gb.addWidget(self.overlap_option_bg, 1, 1)

        gb.setLayout(layout_gb)
        gb.hide()
        self._gb = gb
        layout.addWidget(gb)

        # will be initialized in ._set_import_options(). Depending data extent
        self.dataset_edit = None
        self.setLayout(layout)

        self._HOLE_ID = "Hole ID*"
        self._X_COL = "X"
        self._X_END_COL = "X_END"

        self._set_import_options()
        if restore_settings:
            self.dataset_edit.restore_settings(BASE_SETTINGS_KEY)

    def get_import_parameters(self) -> dict:
        """
        Return defined import parameters.
        """
        res = {}
        res["header"] = self.dataset_edit.header_row.value()
        res["encoding"] = self.dataset_edit.encoding_cbx.currentText()
        res["delimiter"] = self.dataset_edit.delimiter_widget.get_delimiter()
        assay_def = self._assay_selection_wizard.get_assay_definition()
        res["variable"] = assay_def.display_name
        res["extent"] = assay_def.data_extent.value
        res["domain"] = assay_def.domain.value

        if assay_def.data_extent == AssayDataExtent.EXTENDED:
            res["gap_resolution"] = self.gap_option_bg.currentText()
            res["overlap_resolution"] = self.overlap_option_bg.currentText()

        res[
            "date_parameters"
        ] = self.dataset_edit.dateformat_selection_widget.get_parameters()

        # column definitions
        col_defs = self.dataset_edit.column_mapping_model.get_column_definition(
            only_with_mapping=False
        )
        columns = {
            col_def.column: {
                "mapping": col_def.mapping,
                "no_data": col_def.no_data_pattern,
            }
            for col_def in col_defs
        }
        res["columns"] = columns
        # uncertainties
        uncerts = (
            self.dataset_edit.mdl_assay_column_uncertainty.get_assay_column_uncertainty()
        )
        uncertainties = {
            k: {
                "upper_whisker_column": v.upper_whisker_column,
                "lower_whisker_column": v.lower_whisker_column,
                "upper_box_column": v.upper_box_column,
                "lower_box_column": v.lower_box_column,
            }
            for k, v in uncerts.items()
        }
        res["uncertainties"] = uncertainties
        # detection
        dets = (
            self.dataset_edit.mdl_assay_detection_limit.get_assay_column_detection_limit()
        )
        detections = {
            k: {
                "detection_min_col": v.detection_min_col,
                "detection_max_col": v.detection_max_col,
            }
            for k, v in dets.items()
        }
        res["detection_limits"] = detections
        # spherical
        sph = self.dataset_edit.mdl_spherical_data.get_spherical_assay()
        spherical = {
            k: {
                "dip_col": v.dip_col,
                "azimuth_col": v.azimuth_col,
                "polarity_col": v.polarity_col,
                "type_": v.type_,
            }
            for k, v in sph.items()
        }
        res["spherical"] = spherical
        # categories
        cat = self.dataset_edit.mdl_assay_column_category.get_assay_column_categories()
        categories = {
            k: {
                "category": v,
                "validation": self.dataset_edit.mdl_assay_column_category.get_assay_column_category_validation(
                    k
                ),
            }
            for k, v in cat.items()
        }
        res["categories"] = categories

        # import policies
        pols = self.dataset_edit.column_mapping_model.get_import_policies(
            sanitize_column_names=False
        )
        res["policies"] = pols

        return res

    def set_import_parameters(self, params: dict) -> None:
        """
        Restore all parameters from dict.
        """
        assay_def = self._assay_selection_wizard.get_assay_definition()
        if params.get("params") is None:
            return

        d = params.get("params")
        if d.get("variable"):
            profile_variable = d.get("variable")
            if profile_variable != assay_def.display_name:
                QMessageBox.warning(
                    self,
                    "Wrong profile",
                    f"You are currently importing data in {assay_def.display_name} table but you provide a profile for {profile_variable}",
                )
                return
        else:
            return

        if assay_def.data_extent == AssayDataExtent.EXTENDED:
            if d.get("gap_resolution"):
                self.gap_option_bg.setCurrentText(d.get("gap_resolution"))
            if d.get("overlap_resolution"):
                self.overlap_option_bg.setCurrentText(d.get("overlap_resolution"))

        # file format
        if d.get("header"):
            self.dataset_edit.header_row.setValue(d.get("header"))
        if d.get("encoding"):
            self.dataset_edit.encoding_cbx.setCurrentText(d.get("encoding"))
        if d.get("delimiter"):
            self.dataset_edit.delimiter_widget.set_delimiter(d.get("delimiter"))

        # date
        if d.get("date_parameters"):
            self.dataset_edit.dateformat_selection_widget.set_parameters(
                d.get("date_parameters")
            )

        # main columns
        col_defs = self.dataset_edit.column_mapping_model.get_column_definition(
            only_with_mapping=False
        )
        if d.get("columns"):
            mapping = d.get("columns")
            for col_def in col_defs:
                if mapping.get(col_def.column):
                    col_def.mapping = mapping.get(col_def.column).get("mapping")
                    col_def.no_data_pattern = mapping.get(col_def.column).get("no_data")

        self.dataset_edit.set_column_definition(col_defs)

        # uncertainties
        if d.get("uncertainties"):
            uncert = d.get("uncertainties")
            uncertainties = {
                k: AssayColumnUncertainty(
                    upper_whisker_column=v["upper_whisker_column"],
                    lower_whisker_column=v["lower_whisker_column"],
                    upper_box_column=v["upper_box_column"],
                    lower_box_column=v["lower_box_column"],
                )
                for k, v in uncert.items()
            }

            self.dataset_edit.mdl_assay_column_uncertainty.set_assay_column_uncertainty(
                uncertainties
            )

        # detection
        if d.get("detection_limits"):
            det = d.get("detection_limits")
            detections = {
                k: AssayColumnDetectionLimit(
                    detection_min_col=v["detection_min_col"],
                    detection_max_col=v["detection_max_col"],
                )
                for k, v in det.items()
            }

            self.dataset_edit.mdl_assay_detection_limit.set_assay_column_detection_limit(
                detections
            )

        # spherical
        if d.get("spherical"):
            sph = d.get("spherical")
            spherical = {
                k: SphericalDefinition(
                    dip_col=v["dip_col"],
                    azimuth_col=v["azimuth_col"],
                    polarity_col=v["polarity_col"],
                    type_=v["type_"],
                )
                for k, v in sph.items()
            }

            self.dataset_edit.mdl_spherical_data.set_spherical_assay(spherical)

        # categories
        if d.get("categories"):
            cat = d.get("categories")
            categories = [
                CategoryPolicy(
                    column=k, category=v["category"], validation=v["validation"]
                )
                for k, v in cat.items()
            ]
            self.dataset_edit.mdl_assay_column_category.set_assay_column_categories(
                categories
            )

        # conflicts policies
        if d.get("policies"):
            pol = d.get("policies")
            self.dataset_edit.column_mapping_model.set_import_policies(pol)

        self.dataset_edit.update_result_table()

    def get_import_policies(self) -> dict:
        """
        Return import policies.
        """
        policies = self.dataset_edit.get_import_policies()
        result = {}
        for key, item in policies.items():
            sanitized_key = sanitize_sql_name(key)
            result[sanitized_key] = "Ignore" if item is None else item

        return result

    def _set_import_options(self) -> None:
        """
        Set Import options depending of data extent.
        Set the right DelimitedTextImportWidget depending of data extent.
        Extended assay include gap/overlap solver + solver including in result preview.
        """
        assay_def = (
            self._assay_selection_wizard.assay_definition_widget.get_assay_definition_without_columns()
        )

        if self.dataset_edit:
            self.layout().removeWidget(self.dataset_edit)
            self.dataset_edit.deleteLater()
            self.dataset_edit = None

        self.dataset_edit = DelimitedTextImportWidget(
            self, self._openlog_connection, "assay"
        )
        # connect to red flag signal to enable/disable next button
        self.dataset_edit.red_flag_signal.connect(self.completeChanged.emit)
        self.layout().addWidget(self.dataset_edit)
        if assay_def.data_extent == AssayDataExtent.EXTENDED:
            self.dataset_edit.is_extended_assay = True
            self.gap_option_bg.currentIndexChanged.connect(
                self.dataset_edit.update_result_table
            )
            self.overlap_option_bg.currentIndexChanged.connect(
                self.dataset_edit.update_result_table
            )
            if self._gb.isHidden():
                self._gb.show()
        elif assay_def.data_extent == AssayDataExtent.DISCRETE:
            self.dataset_edit.is_extended_assay = False
            if not self._gb.isHidden():
                self._gb.hide()
        else:
            return

    def get_gap_data_resolution(self) -> GapDataResolution:
        for i, o in enumerate(GapDataResolution):
            if i == self.gap_option_bg.currentIndex():
                return o

    def get_overlap_data_resolution(self) -> OverlapDataResolution:
        for i, o in enumerate(OverlapDataResolution):
            if i == self.overlap_option_bg.currentIndex():
                return o

    def _remove_invalid_columns(self):
        """
        Remove wrongly defined columns before import.
        """
        # TODO : set removed columns in confirmation page
        # used only for spherical undefined data, for now
        df = self.dataset_edit.get_dataframe(spherical=True)
        assay_def = self.dataset_edit.get_column_definition()
        diff = [assay.column for assay in assay_def if assay.column not in df.columns]

        for column in diff:
            self.dataset_edit.column_mapping_model.remove_column(column)

    def isComplete(self) -> bool:
        """
        Override of QWizardPage.isComplete method to enable next button only if there is no red flag.
        """
        return not self.dataset_edit.red_flag

    def validatePage(self) -> bool:
        """
        Validate current page content (return always True since data is optional)

        Returns: True

        """
        valid = True
        self.dataset_edit.save_setting(BASE_SETTINGS_KEY)
        self._remove_invalid_columns()
        used_category = self.dataset_edit.get_used_category()
        empty_cat = [c for c in used_category if not c]
        if len(empty_cat) != 0:
            valid = False
            QMessageBox.warning(
                self,
                "No category defined",
                f"Define category name for all categorical columns.",
            )

        return valid

    def initializePage(self) -> None:
        """
        Initialize page before show.

        Update expected column from AssayDefinition

        """
        self.dataset_edit.set_button_layout_visible(
            self._assay_selection_wizard.assay_creation_needed()
        )
        self._set_assay_definition(self._assay_selection_wizard.get_assay_definition())

    def get_assay_definition(self) -> AssayDefinition:
        """
        Returns AssayDefinition from widget

        Returns: AssayDefinition

        """
        assay_definition = self._assay_selection_wizard.get_assay_definition()

        # Add created assay columns
        # Get all expected columns
        expected_columns = self.dataset_edit.get_column_definition()

        # Remove columns already defined in assay definition
        expected_columns = [
            col
            for col in expected_columns
            if col.column
            not in [x.column for x in self._get_columns_definition(assay_definition)]
        ]

        assay_column_uncertainty = self.dataset_edit.get_assay_column_uncertainty()

        # detection limit
        assay_column_detection_limit = (
            self.dataset_edit.get_assay_column_detection_limit()
        )

        for mapping in expected_columns:
            col = mapping.column
            # Define column uncertainty if available
            if col in assay_column_uncertainty:
                uncertainty = assay_column_uncertainty[col]
            else:
                uncertainty = AssayColumnUncertainty()
            # Define column detection limit if available
            if col in assay_column_detection_limit:
                detection_limit = assay_column_detection_limit[col]
            else:
                detection_limit = AssayColumnDetectionLimit()
            if mapping.series_type == AssaySeriesType.IMAGERY:
                image_format_col = "image_format"
            else:
                image_format_col = ""
            # col is display_name
            sanitized_col = sanitize_sql_name(col)
            # sanitize extra columns (uncertainties, detection limits)
            uncertainty.lower_box_column = sanitize_sql_name(
                uncertainty.lower_box_column
            )
            uncertainty.upper_box_column = sanitize_sql_name(
                uncertainty.upper_box_column
            )
            uncertainty.lower_whisker_column = sanitize_sql_name(
                uncertainty.lower_whisker_column
            )
            uncertainty.upper_whisker_column = sanitize_sql_name(
                uncertainty.upper_whisker_column
            )

            detection_limit.detection_max_col = sanitize_sql_name(
                detection_limit.detection_max_col
            )
            detection_limit.detection_min_col = sanitize_sql_name(
                detection_limit.detection_min_col
            )

            assay_definition.columns[sanitized_col] = AssayColumn(
                name=sanitized_col,
                display_name=col,
                series_type=mapping.series_type,
                unit=mapping.unit,
                uncertainty=uncertainty,
                detection_limit=detection_limit,
                category_name=mapping.category_name,
                image_format_col=image_format_col,
            )
        return assay_definition

    def get_assays(self) -> List[GenericAssay]:
        """
        Returns GenericAssay with imported data

        Returns: List[GenericAssay]

        """
        result = []
        df = self.dataset_edit.get_dataframe(spherical=True)

        if df is not None:
            # Create GenericAssay for each collar
            for hole_id, hole_id_df in df.groupby(self._HOLE_ID):
                assay = self._create_generic_assay_from_dataframe(hole_id, hole_id_df)
                assay.set_person(
                    self._person_dataset_import_wizard.person_cb.currentText()
                )
                assay.set_dataset(
                    self._person_dataset_import_wizard.dataset_cb.currentText()
                )
                assay.set_import_date(datetime.today())  # .strftime('%Y-%m-%d')
                result.append(assay)

        return result

    def get_used_categories_values(self) -> {str: [str]}:
        """
        Get values imported of each category

        Returns: {str: [str]} list of values for each used category

        """
        return self.dataset_edit.get_used_categories_values()

    def _create_generic_assay_from_dataframe(
        self, hole_id: str, hole_id_df: pd.DataFrame
    ) -> GenericAssay:
        """
        Create GenericAssay from input dataframe

        Args:
            hole_id: (str) hole id
            hole_id_df: (pd.DataFrame)

        Returns: GenericAssay with imported data

        """
        assay_definition = self.get_assay_definition()
        assay = GenericAssay(hole_id=str(hole_id), assay_definition=assay_definition)
        assay.import_x_values = self._import_x_values_from_dataframe(
            assay_definition, hole_id_df
        )
        # Add import y values for assay column
        for column, val in assay_definition.columns.items():
            import_val = hole_id_df.loc[:, val.display_name].to_numpy()

            if val.series_type == AssaySeriesType.IMAGERY:
                imported_file_path = self.dataset_edit.get_file_path()
                image_format_values = []
                image_binary_values = []
                for image_path_str in import_val:
                    # Read image
                    if image_path_str and Path(image_path_str).exists():
                        image_path = Path(image_path_str)
                        image_binary_values.append(image_path.read_bytes())
                        image_format_values.append(image_path.suffix)
                    else:
                        image_binary_values.append(None)
                        image_format_values.append(None)
                assay.import_y_values[column] = image_binary_values
                assay.import_y_values[val.image_format_col] = image_format_values
            else:
                # Data always stored in SI base unit. Convert data before import
                if val.unit and val.series_type != AssaySeriesType.POLAR:
                    import_val = pint_utilities.convert_to_base_unit(
                        val.unit, import_val
                    )
                assay.import_y_values[column] = import_val

        # Add import y values for uncertainty columns
        sanitized_column_names = [sanitize_sql_name(col) for col in hole_id_df.columns]
        for column in assay_definition.get_uncertainty_columns():
            # if column don't exist (e.g. uncertainties), create it with Nones
            if column not in sanitized_column_names:
                hole_id_df[column] = None
                sanitized_column_names = sanitized_column_names + [column]

            col_index = sanitized_column_names.index(column)
            assay.import_y_values[column] = hole_id_df.iloc[:, col_index].to_numpy(
                dtype=float
            )

        # Add import y values for detection limits. missing values are allowed
        for column in assay_definition.get_detection_limit_columns():
            # if column don't exist (e.g. uncertainties), create it with Nones
            if column not in sanitized_column_names:
                hole_id_df[column] = None
                sanitized_column_names = sanitized_column_names + [column]

            col_index = sanitized_column_names.index(column)

            array = (
                hole_id_df.iloc[:, col_index]
                .replace({"": np.nan})
                .to_numpy(dtype=float)
            )
            assay.import_y_values[column] = np.where(np.isnan(array), None, array)

        return assay

    def _import_x_values_from_dataframe(self, assay_definition, hole_id_df) -> np.array:
        """
        Create np.array for GenericAssay import x values

        Args:
            assay_definition: AssayDefinition
            hole_id_df: input pd.DataFrame

        Returns:
            np.array import x values

        """
        x_col_dtype = self._get_x_col_type(assay_definition)
        if assay_definition.data_extent == AssayDataExtent.EXTENDED:
            import_x_values = np.array(
                list(
                    zip(
                        hole_id_df.loc[:, self._X_COL].to_numpy(dtype=x_col_dtype),
                        hole_id_df.loc[:, self._X_END_COL].to_numpy(dtype=x_col_dtype),
                    )
                )
            )
        else:
            import_x_values = hole_id_df.loc[:, self._X_COL].to_numpy(dtype=x_col_dtype)
        return import_x_values

    @staticmethod
    def _get_x_col_type(assay_definition: AssayDefinition) -> type:
        """
        Return expected x column type from AssayDefinition

        Args:
            assay_definition: AssayDefinition

        Returns: expected x column type

        """
        if assay_definition.domain == AssayDomainType.TIME:
            x_col_dtype = datetime
        else:
            x_col_dtype = float
        return x_col_dtype

    def _set_assay_definition(self, assay_definition: AssayDefinition) -> None:
        """
        Define DelimitedTextImportWidget options from AssayDefinition

        Args:
            assay_definition: AssayDefinition
        """
        self._define_col_names(assay_definition)
        expected_columns = self._get_columns_definition(assay_definition)
        self.dataset_edit.set_column_definition(expected_columns)
        # update widgets for uncertainties and detection limit
        self.dataset_edit.already_created = (
            not self._assay_selection_wizard.assay_creation_needed()
        )
        self.dataset_edit.is_extended_assay = (
            assay_definition.data_extent == AssayDataExtent.EXTENDED
        )
        if assay_definition.data_extent == AssayDataExtent.EXTENDED:
            self._gb.show()
        else:
            self._gb.hide()
        self.dataset_edit.update_optionnal_widgets(assay_definition)

    def _define_col_names(self, assay_definition: AssayDefinition) -> None:
        """
        Define default column name from AssayDefinition

        Args:
            assay_definition: AssayDefinition
        """
        if assay_definition.domain == AssayDomainType.TIME:
            self._X_COL = "Time"
        else:
            self._X_COL = "Depth"

        if assay_definition.data_extent == AssayDataExtent.EXTENDED:
            self._X_END_COL = self._X_COL + " max.*"
            self._X_COL = self._X_COL + " min.*"
        else:
            self._X_COL = self._X_COL + "*"

    def _get_columns_definition(
        self, assay_definition: AssayDefinition
    ) -> [ColumnDefinition]:
        """
        Get expected columns from AssayDefinition

        Args:
            assay_definition:

        Returns: expected columns

        """
        x_unit = "m" if assay_definition.domain == AssayDomainType.DEPTH else "datetime"
        x_series_type = (
            AssaySeriesType.NUMERICAL
            if assay_definition.domain == AssayDomainType.DEPTH
            else AssaySeriesType.DATETIME
        )
        expected_columns = [
            ColumnDefinition(
                column=self._HOLE_ID, fixed=True, series_type=AssaySeriesType.NOMINAL
            ),
            ColumnDefinition(
                column=self._X_COL,
                unit=x_unit,
                fixed=True,
                series_type=x_series_type,
            ),
        ]
        if assay_definition.data_extent == AssayDataExtent.EXTENDED:
            expected_columns.append(
                ColumnDefinition(
                    column=self._X_END_COL,
                    unit=x_unit,
                    fixed=True,
                    series_type=x_series_type,
                )
            )
        for key, column in assay_definition.columns.items():
            expected_columns.append(
                ColumnDefinition(
                    column=column.display_name,
                    fixed=True,
                    unit=column.unit,
                    series_type=column.series_type,
                    category_name=column.category_name,
                )
            )
        return expected_columns
