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

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


class ExtendedAssayDelimitedTextImportWidget(DelimitedTextImportWidget):
    def __init__(
        self, parent=None, openlog_connection: OpenLogConnection = None
    ) -> None:
        """
        Override DelimitedTextImportWidget get_dataframe method to get interpolated
        assay (extended) data with gap and overlap

        Args:
            parent: QWidget parent
        """
        self.UPDATE_STATUS_COL = "update_status"
        super(ExtendedAssayDelimitedTextImportWidget, self).__init__(
            parent, openlog_connection
        )
        self.removed_rows = 0

    def get_dataframe(
        self, nb_row: int = None, spherical: bool = False
    ) -> pandas.DataFrame:
        """
        Override DelimitedTextImportWidget to interpolate gap and overlap data

        Args:
            nb_row: read only a specific number of row (default None : all rows are read)

        Returns:
            panda dataframe with expected column and interpolated data

        """
        df = super().get_dataframe(nb_row, spherical)

        if df is not None:
            value_cols = [
                col
                for col in df.columns
                if col
                not in [
                    self.parent()._HOLE_ID,
                    self.parent()._X_COL,
                    self.parent()._X_END_COL,
                ]
            ]
            interpolation = ExtendedAssayInterpolation(
                gap_resolution=self.parent().get_gap_data_resolution(),
                overlap_resolution=self.parent().get_overlap_data_resolution(),
                columns=ExtendedAssayColumn(
                    hole_id_col=self.parent()._HOLE_ID,
                    from_col=self.parent()._X_COL,
                    to_col=self.parent()._X_END_COL,
                    val_cols=value_cols,
                    update_status_col=self.UPDATE_STATUS_COL,
                ),
            )
            try:
                df = interpolation.interpolated_dataframe(df)
                self.data_is_valid = True
            except ExtendedAssayInterpolation.InvalidColumnData:
                self.data_is_valid = False

            # delete rows with 0 thickness
            nrow = df.shape[0]
            df = df[df[self.parent()._X_COL] != df[self.parent()._X_END_COL]]
            self.removed_rows = nrow - df.shape[0]
        return df

    def _update_table_content(
        self, table_widget: QTableWidget, array: pandas.DataFrame
    ) -> None:
        """
        Override DelimitedTextImportWidget to display status of interpolation

        Args:
            table_widget: QTableWidget used of display
            array: Dataframe to display in QTableWidget
        """
        super()._update_table_content(table_widget, array)

        if self.UPDATE_STATUS_COL in array.columns:
            # Update item style
            for j in range(0, table_widget.columnCount()):
                i = 0
                for index, r in array.iterrows():
                    if r[self.UPDATE_STATUS_COL]:
                        font = table_widget.item(i, j).font()
                        font.setBold(True)
                        table_widget.item(i, j).setFont(font)
                    i = i + 1

            # Don't show update status col (last col). Must be reinitialized at each call because columns could be added
            n_col = table_widget.columnCount()
            for index in range(n_col):
                if index == n_col - 1:
                    hide = True
                else:
                    hide = False
                table_widget.setColumnHidden(index, hide)


class AssayImportPageWizard(QWizardPage):
    def __init__(
        self,
        parent: QWidget,
        openlog_connection: OpenLogConnection,
        assay_selection_wizard: AssaySelectionPageWizard,
    ) -> 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.tr = PlgTranslator().tr
        self._assay_selection_wizard = assay_selection_wizard

        self.setTitle(self.tr("Downhole data values import"))
        label = QLabel(self.tr("Select a .csv file for lithology import"))
        label.setWordWrap(True)
        layout = QVBoxLayout()
        layout.addWidget(label)
        gb = QgsCollapsibleGroupBox(self.tr("Import options"), self)
        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(self.tr(o.value))
        layout_gb.addWidget(QLabel(self.tr("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(self.tr(o.value))
        layout_gb.addWidget(QLabel(self.tr("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 _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

        if assay_def.data_extent == AssayDataExtent.EXTENDED:

            self.dataset_edit = ExtendedAssayDelimitedTextImportWidget(
                self, self._openlog_connection
            )
            self.layout().addWidget(self.dataset_edit)
            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 = DelimitedTextImportWidget(
                self, self._openlog_connection
            )
            self.layout().addWidget(self.dataset_edit)
            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 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,
                self.tr("No category defined"),
                self.tr(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()

        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()
            if mapping.series_type == AssaySeriesType.IMAGERY:
                image_format_col = "image_format"
            else:
                image_format_col = ""
            assay_definition.columns[col] = AssayColumn(
                name=col,
                series_type=mapping.series_type,
                unit=mapping.unit,
                uncertainty=uncertainty,
                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)
                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[:, column].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
        for column in assay_definition.get_uncertainty_columns():
            assay.import_y_values[column] = hole_id_df.loc[:, column].to_numpy(
                dtype=float
            )
        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)

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

    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=key,
                    fixed=True,
                    unit=column.unit,
                    series_type=column.series_type,
                    category_name=column.category_name,
                )
            )
        return expected_columns
