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.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

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,
    ) -> 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()
        self.dataset_edit.restore_settings(BASE_SETTINGS_KEY)

    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
