#  Copyright (c) 2022 - 2023 D-BSSE, ETH Zurich.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#   limitations under the License.
#

import pandas as pd
from PySide6 import QtCore
from PySide6.QtCore import Qt


class PandasDataModel(QtCore.QAbstractTableModel):
    """
    Class to populate a table view with a pandas dataframe
    """

    def __init__(self, data, parent=None):
        QtCore.QAbstractTableModel.__init__(self, parent)
        self._data = data

    def clear(self):
        # Begin a batch operation to remove all rows
        self.beginResetModel()

        # Clear your data here
        self._data = pd.DataFrame()

        # End the batch operation, notifying the views
        self.endResetModel()

    def rowCount(self, parent=None):
        return self._data.shape[0]

    def columnCount(self, parent=None):
        return self._data.shape[1]

    def data(self, index, role=Qt.DisplayRole):
        if index.isValid():
            if role == Qt.DisplayRole:
                value = self._data.iloc[index.row(), index.column()]
                if self._data.columns[index.column()] in ["efo", "dwell"]:
                    return f"{int(value)}"
                if self._data.columns[index.column()] in ["x", "y", "z"]:
                    return f"{value:.1f}"
                if self._data.columns[index.column()] in ["cfr", "dcr"]:
                    return f"{value:.3f}"
                if isinstance(value, float):  # Check if the value is a float
                    return (
                        f"{value:.2f}"  # Format the float value with 2 decimal places
                    )
                return str(value)
        return None

    def headerData(self, col, orientation, role):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self._data.columns[col]
        return None

    def sort(self, column, order):
        self.layoutAboutToBeChanged.emit()
        column_name = self._data.columns[column]
        self._data.sort_values(
            column_name, ascending=(order == Qt.AscendingOrder), inplace=True
        )
        self.layoutChanged.emit()
