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

        self.dataset_edit.restore_settings(BASE_SETTINGS_KEY)

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

    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
