import pandas
from qgis.gui import QgsCollapsibleGroupBox
from qgis.PyQt.QtWidgets import QCheckBox, QLabel, QVBoxLayout, QWidget, QWizardPage
from xplordb.datamodel.survey import Survey

from openlog.datamodel.assay.generic_assay import AssaySeriesType
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.gui.utils.column_definition import ColumnDefinition
from openlog.gui.utils.delimited_text_import_widget import DelimitedTextImportWidget

BASE_SETTINGS_KEY = "/OpenLog/gui/import/survey"


class SurveysImportPageWizard(QWizardPage):
    def __init__(
        self,
        parent: QWidget,
        openlog_connection: OpenLogConnection,
        restore_settings: bool = True,
    ) -> None:
        """
        QWizard to import surveys into xplordb from csv file

        Args:
            openlog_connection: OpenLogConnection used to import surveys
            parent : QWidget parent
        """
        super().__init__(parent)

        self._openlog_connection = openlog_connection
        self.setTitle("Survey import")

        self.HOLE_ID_COL = "HoleID"
        self.DIP_COL = "Dip"
        self.AZIMUTH_COL = "Azimuth"
        self.LENGTH_COL = "Length"

        label = QLabel("Select a .csv file or a QGIS layer in Source table section.")
        label.setWordWrap(True)

        layout = QVBoxLayout()
        layout.addWidget(label)

        # Add import option group box
        gb = QgsCollapsibleGroupBox("Import options", self)
        layout_gb = QVBoxLayout()
        gb.setLayout(layout_gb)
        self.dip_inversion_checkbox = QCheckBox("Invert Dips", self)
        layout_gb.addWidget(self.dip_inversion_checkbox)
        self.update_eoh_checkbox = QCheckBox("Update EOH", self)
        layout_gb.addWidget(self.update_eoh_checkbox)
        gb.setLayout(layout_gb)
        layout.addWidget(gb)

        self.dataset_edit = DelimitedTextImportWidget(self, None, "survey")
        self.dataset_edit.col_def_table_view.hideRow(
            self.dataset_edit.column_mapping_model.NO_DATA_ROW
        )
        # for displaying conflicts options
        self.dataset_edit.already_created = True
        self.dataset_edit.update_optionnal_widgets(None)
        # connect to red flag signal to enable/disable next button
        self.dataset_edit.red_flag_signal.connect(self.completeChanged.emit)

        self.dataset_edit.set_column_definition(
            [
                ColumnDefinition(
                    column=self.HOLE_ID_COL,
                    fixed=True,
                    series_type=AssaySeriesType.NOMINAL,
                ),
                ColumnDefinition(
                    column=self.DIP_COL,
                    unit="°",
                    fixed=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.AZIMUTH_COL,
                    unit="°",
                    fixed=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
                ColumnDefinition(
                    column=self.LENGTH_COL,
                    unit="m",
                    fixed=True,
                    series_type=AssaySeriesType.NUMERICAL,
                ),
            ]
        )
        layout.addWidget(self.dataset_edit)
        self.setLayout(layout)
        self.dataset_edit.add_column_button.hide()
        self.dataset_edit.remove_column_button.hide()

        # Add column conversion for dip
        self.dataset_edit.set_column_conversion({self.DIP_COL: self._invert_dip})
        self.dip_inversion_checkbox.clicked.connect(
            self.dataset_edit.update_result_table
        )
        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()
        res[
            "date_parameters"
        ] = self.dataset_edit.dateformat_selection_widget.get_parameters()

        # options
        res["invert_dip"] = self.dip_inversion_checkbox.isChecked()
        res["update_eoh"] = self.update_eoh_checkbox.isChecked()

        # columns
        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

        # 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 wizard with parameters.
        """
        if params.get("params") is None:
            return
        d = params.get("params")
        # 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")
            )

        # options
        if d.get("invert_dip"):
            self.dip_inversion_checkbox.setChecked(d.get("invert_dip"))
        if d.get("update_eoh"):
            self.update_eoh_checkbox.setChecked(d.get("update_eoh"))

        # 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)

        # 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.column_mapping_model.get_import_policies(
            sanitize_column_names=False
        )
        col_map = {
            self.HOLE_ID_COL: "hole_id",
            self.DIP_COL: "dip",
            self.AZIMUTH_COL: "azimuth",
            self.LENGTH_COL: "depth",
        }
        result = {}
        for key, item in policies.items():
            sanitized_key = col_map.get(key)
            result[sanitized_key] = "Ignore" if item is None else item

        return result

    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 _invert_dip(self, col: str, df: pandas.DataFrame) -> pandas.Series:
        """
        Conversion function for dip

        Args:
            col: column to be converted
            df: original dataframe

        Returns: new pandas.Series for column col

        """
        dip = df[col].astype(float)
        if self.dip_inversion_checkbox.isChecked():
            dip = -dip
        return dip

    def data_label(self) -> str:
        """
        Returns label to be used in confirmation dialog

        Returns: imported data label

        """
        return "Surveys"

    def data_count(self) -> int:
        """
        Returns expected imported data count to be displayed in confirmation dialog

        Returns: expected imported data count

        """
        df = self.dataset_edit.get_dataframe()
        return df.shape[0] if df is not None else 0

    def _update_collar_eoh(self, df: pandas.DataFrame) -> None:
        """
        If a survey has depth > eoh, then update EOH in collar table.
        Args:
            - df (pandas.DataFrame): DataFrame describing surveys to import
        """
        # get all collars
        collars = self._openlog_connection.get_read_iface().get_all_collars()
        collars_eoh = pandas.DataFrame(
            [(c.hole_id, c.eoh) for c in collars if c is not None],
            columns=[self.HOLE_ID_COL, "eoh"],
        )
        # get max depth by hole
        max_depth = (
            df.groupby(self.HOLE_ID_COL).agg({self.LENGTH_COL: "max"}).reset_index()
        )
        res = pandas.merge(max_depth, collars_eoh, on=self.HOLE_ID_COL)
        res = res[res[self.LENGTH_COL] > res["eoh"]]
        collars_to_be_updated = [
            self._openlog_connection.get_read_iface().get_collar(hole)
            for hole in res[self.HOLE_ID_COL]
        ]
        # set new EOH
        for collar, eoh in zip(collars_to_be_updated, res[self.LENGTH_COL]):
            setattr(collar, "eoh", eoh)
        # update database
        if len(collars_to_be_updated) > 0:
            self._openlog_connection.get_write_iface().update_collars(
                collars_to_be_updated
            )

    def import_data(self):
        """
        Import data into openlog database.

        OpenLogConnection.ImportData exception can be raised.

        """
        df = self.dataset_edit.get_dataframe()
        import_policies = self.get_import_policies()
        if self.update_eoh_checkbox.isChecked():
            self._update_collar_eoh(df)
        if df is not None:
            surveys = [
                Survey(
                    hole_id=r[self.HOLE_ID_COL],
                    data_set=self.field("dataset"),
                    loaded_by=self.field("person"),
                    dip=r[self.DIP_COL],
                    azimuth=r[self.AZIMUTH_COL],
                    depth=r[self.LENGTH_COL],
                )
                for index, r in df.iterrows()
            ]
            self._openlog_connection.get_write_iface().import_surveys(
                surveys, import_policies
            )

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

        Returns: True

        """
        self.dataset_edit.save_setting(BASE_SETTINGS_KEY)
        return self.dataset_edit.data_is_valid
