import numpy
import pandas as pd
from PyQt6 import QtCore
from PyQt6.QtCore import Qt, QModelIndex


def _get_str_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """Return a given dataframe with all values turned into strings.

    When the data given to the PandasModel class contains non-strings,
    an attached QTableView seems to respond rather slowly. This function
    turns all data in the DataFrame into strings, yielding a better experience.
    """
    return df.astype(str)


class PandasModel(QtCore.QAbstractTableModel):
    def __init__(self, parent=None):
        QtCore.QAbstractTableModel.__init__(self, parent)

        self._data = pd.DataFrame()
        self._data_str = pd.DataFrame()
        self._horizontalHeaders = None

    # Overloaded functions

    def rowCount(self, parent=None):
        return len(self._data_str.values)

    def columnCount(self, parent=None):
        return self._data_str.columns.size

    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
        if not index.isValid():
            return QtCore.QVariant()

        if (role != Qt.ItemDataRole.DisplayRole) and (
                role != Qt.ItemDataRole.EditRole):
            return QtCore.QVariant()

        if role == Qt.ItemDataRole.DisplayRole:
            item = self._data_str.iloc[index.row(), index.column()]
            return QtCore.QVariant(item)
        elif role == Qt.ItemDataRole.EditRole:
            item = self._data.iloc[index.row(), index.column()]
            if type(item) is numpy.float64:
                return QtCore.QVariant(float(item))
            else:
                return QtCore.QVariant(item)

    def headerData(self, section, orientation,
                   role=Qt.ItemDataRole.DisplayRole):

        if not ((orientation == Qt.Orientation.Horizontal)
                and (role == Qt.ItemDataRole.DisplayRole)):
            return super().headerData(section, orientation, role)

        return self._horizontalHeaders[section]

    # Other functions

    def set_dataframe(self, df):
        self._data = df
        self._data_str = _get_str_dataframe(df)
        self._horizontalHeaders = list(df.columns)

        self.layoutAboutToBeChanged.emit()
        self.layoutChanged.emit()

    def get_dataframe(self) -> pd.DataFrame:
        return self._data