#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据模型模块

包含 PandasTableModel 等数据模型类。
"""

import pandas as pd
from PyQt5.QtCore import Qt, QAbstractTableModel


class PandasTableModel(QAbstractTableModel):
    """支持DataFrame的表格模型"""
    def __init__(self, data_frame: pd.DataFrame):
        super().__init__()
        self._data = data_frame.copy()
        if self._data.empty:
            self._data = pd.DataFrame()
        self._data.fillna("", inplace=True)
        self._columns = [str(col) for col in self._data.columns]

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

    def columnCount(self, parent=None):
        return len(self._columns)

    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid() or role != Qt.DisplayRole:
            return None

        value = self._data.iat[index.row(), index.column()]
        if pd.isna(value):
            return ""
        return str(value)

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if role != Qt.DisplayRole:
            return None
        if orientation == Qt.Horizontal:
            if section < len(self._columns):
                return self._columns[section]
            return str(section)
        return str(section + 1)

    def flags(self, index):
        if not index.isValid():
            return Qt.NoItemFlags
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable