from PyQt5.QtCore import Qt, QAbstractTableModel


class TableModel(QAbstractTableModel):
    def __init__(self, data):
        super(TableModel, self).__init__()
        self._data = data

    def data(self, index, role):
        if role == Qt.DisplayRole:
            value = self._data.iloc[index.row(), index.column()]
            # if role == Qt.TextAlignmentRole:
            #     value = self._data.iloc[index.row(), index.column()]
            #     if isinstance(value, int) or isinstance(value, float):
            #         return Qt.AlignVCenter + Qt.AlignRight
            return str(value)
            # if isinstance(value, float):
            #     # Render float to 2 dp
            #     return "%.2f" % value
            # if isinstance(value, str):
            #     # Render strings with quotes
            #     return '"%s"' % value

    def rowCount(self, index):
        # The length of the outer list.
        return self._data.shape[0]

    def columnCount(self, index):
        # The following takes the first sub-list, and returns
        # the length (only works if all rows are an equal length)
        return self._data.shape[1]

    def headerData(self, section, orientation, role):
        # section is the index of the column/row.
        if role == Qt.DisplayRole:
            if orientation == Qt.Horizontal:
                return str(self._data.columns[section])
            if orientation == Qt.Vertical:
                return str(self._data.index[section])