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
from openlog.toolbelt.translator import PlgTranslator

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.tr = PlgTranslator().tr
        self._openlog_connection = openlog_connection
        self.setTitle(self.tr("Survey import"))

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

        label = QLabel(self.tr("Select a .csv file for surveys import"))
        label.setWordWrap(True)

        layout = QVBoxLayout()
        layout.addWidget(label)

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

        self.dataset_edit = DelimitedTextImportWidget(self)
        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 _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 self.tr("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 import_data(self):
        """
        Import data into openlog database.

        OpenLogConnection.ImportData exception can be raised.

        """
        df = self.dataset_edit.get_dataframe()
        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
