src/ui/main_window.py - master - Walmet/CTF

src/ui/main_window.py
323 строки12 KB
Maksim Goncharov
MVP
14 дек 2025, 23:49
14 дек 2025, 23:49a4c2b1d
import sys
import pandas as pd
from PyQt5.QtWidgets import (
    QMainWindow, QApplication, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QTextEdit, QFileDialog, QSplitter, QLabel, QMessageBox, QAbstractItemView
)
from PyQt5.QtCore import Qt, QEvent

from ..components.document_handler import DocumentHandler
from ..components.data_extractor import DataExtractor

from PyQt5.QtWidgets import QTableWidget, QTableWidgetItem, QComboBox
from PyQt5.QtGui import QColor

COLOR_STATUS_0 = QColor(240, 240, 240)  # White (Нет ссылки)
COLOR_STATUS_1 = QColor(255, 255, 255)  # White (Распознавание уверено)
COLOR_STATUS_2 = QColor(220, 220, 255)  # Light Purple (Распознавание не уверено)
COLOR_STATUS_3 = QColor(193, 230, 193)  # Light Green (После редактирования)
COLOR_HIGHLIGHT = QColor('yellow')


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("MVP")
        self.setGeometry(100, 100, 1000, 600)

        self.document_handler = DocumentHandler()
        self.data_extractor = DataExtractor('./src/data/template1.csv')

        self.dataframes = {}
        self.current_left_dataframe = None
        self.current_right_dataframe = None
        self.current_metadata_dataframe = None
        self.template_path = None

        self.initui()

        self.right_table.installEventFilter(self)
        self.right_table.cellChanged.connect(self.update_dataframe_from_right_table)

        self.right_table.cellClicked.connect(self.right_cell_entered)

    def initui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        splitter = QSplitter(Qt.Horizontal)
        main_layout.addWidget(splitter)

        # --- Left Panel (Source Table) ---
        leftpanelwidget = QWidget()
        leftpanellayout = QVBoxLayout(leftpanelwidget)
        leftpanellayout.addWidget(QLabel("Loaded document:"))

        self.file_selector = QComboBox()
        self.file_selector.currentTextChanged.connect(self.on_file_selected)
        leftpanellayout.addWidget(self.file_selector)

        self.left_table = QTableWidget()
        self.left_table.setEditTriggers(QTableWidget.NoEditTriggers)
        leftpanellayout.addWidget(self.left_table)

        self.load_button = QPushButton("Load")
        self.load_button.clicked.connect(self.load_csv_files)
        leftpanellayout.addWidget(self.load_button)

        splitter.addWidget(leftpanelwidget)

        # --- Right Panel (Target Table) ---
        rightpanelwidget = QWidget()
        rightpanellayout = QVBoxLayout(rightpanelwidget)

        rightpanellayout.addWidget(QLabel("Transformed document:"))

        self.load_template_button = QPushButton("Set Template")
        self.load_template_button.clicked.connect(self.load_template_file)
        rightpanellayout.addWidget(self.load_template_button)

        self.right_table = QTableWidget()
        self.right_table.setEditTriggers(QTableWidget.DoubleClicked | QTableWidget.AnyKeyPressed)
        self.right_table.setMouseTracking(True)
        self.right_table.setWordWrap(True)
        self.right_table.resizeRowsToContents()

        rightpanellayout.addWidget(self.right_table)

        self.save_button = QPushButton("Save")
        self.save_button.clicked.connect(self.save_csv_file)
        self.save_button.setEnabled(False)
        rightpanellayout.addWidget(self.save_button)

        splitter.addWidget(rightpanelwidget)
        splitter.setSizes([self.width() // 2, self.width() // 2])

    def load_template_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select Template CSV File", "", "CSV Files (*.csv)")
        if file_path:
            self.template_path = file_path
        else:
            self.template_path = None

    def save_csv_file(self):
        if self.current_right_dataframe is None:
            QMessageBox.warning(self, "Save Error", "No data loaded to save.")
            return

        output_path, _ = QFileDialog.getSaveFileName(self, "Save CSV As", "output.csv", "CSV Files (*.csv)")
        if output_path:
            self.document_handler.save_dataframe_to_csv(self.current_right_dataframe, output_path)

    def load_csv_files(self):
        file_paths, _ = QFileDialog.getOpenFileNames(self, "Select CSV Files", "", "CSV Files (*.csv)")
        if not file_paths:
            return

        self.dataframes.clear()
        self.file_selector.clear()
        self.left_table.clear()
        self.right_table.clear()
        self.current_left_dataframe = None
        self.current_right_dataframe = None
        self.current_metadata_dataframe = None

        new_dataframes = {}

        for path in file_paths:
            df = self.document_handler.load_csv_document(path)
            if df is not None and not df.empty:
                filename = path.split('/')[-1]
                new_dataframes[filename] = df
            else:
                QMessageBox.warning(self, "Loading Error", f"Could not load or file is empty: {path}")

        self.dataframes = new_dataframes

        if self.dataframes:
            transformed_data = self.data_extractor.transform_dataframes(self.dataframes, self.template_path)

            if transformed_data is None or transformed_data[0] is None:
                QMessageBox.critical(self, "Error", "DataExtractor did not return valid DataFrames.")
                return

            self.current_right_dataframe, self.current_metadata_dataframe = transformed_data

            first_filename = next(iter(self.dataframes.keys()))
            self.current_left_dataframe = self.dataframes[first_filename]

            self.file_selector.addItems(self.dataframes.keys())
            self.display_dataframe(first_filename)
            self.save_button.setEnabled(True)
        else:
            self.save_button.setEnabled(False)

    def on_file_selected(self, filename):
        if filename in self.dataframes:
            self.display_dataframe(filename)

    def display_dataframe(self, filename):
        df_left = self.dataframes.get(filename)
        df_right = self.current_right_dataframe

        self.right_table.blockSignals(True)

        if df_left is not None:
            self.current_left_dataframe = df_left
            self.populate_table(self.left_table, self.current_left_dataframe)
            if df_right is not None:
                self.populate_table(self.right_table, df_right, self.current_metadata_dataframe)
            self.clear_left_highlight()
        else:
            self.left_table.clear()
            self.right_table.clear()

        self.right_table.blockSignals(False)

    def update_dataframe_from_right_table(self, row, column):
        """Updates the internal pandas DataFrame and sets status to 3 (Manually edited)."""
        df = self.current_right_dataframe
        meta_df = self.current_metadata_dataframe

        if df is None or meta_df is None:
            return

        item = self.right_table.item(row, column)
        if item is None:
            return

        new_value = item.text()

        self.right_table.blockSignals(True)

        try:
            df.iloc[row, column] = new_value
            meta_df.loc[(row, column), 'status'] = 3
            item.setBackground(COLOR_STATUS_3)
            self.right_table.resizeRowToContents(row)

        except Exception as e:
            print(f"Error updating DataFrame: {e}")
            item.setText(str(df.iloc[row, column]))

        finally:
            self.right_table.blockSignals(False)

    def populate_table(self, table: QTableWidget, df: pd.DataFrame, meta_df: pd.DataFrame = None):
        """Fill the QTableWidget with content and apply status coloring for the right table."""

        df_display = df.reset_index(drop=True)

        table.clear()
        table.setRowCount(df_display.shape[0])
        table.setColumnCount(df_display.shape[1])
        table.setHorizontalHeaderLabels(df_display.columns.tolist())

        table.resizeRowsToContents()

        for row in range(df_display.shape[0]):
            for col in range(df_display.shape[1]):
                item = QTableWidgetItem(str(df_display.iat[row, col]))
                color = QColor('white')

                item.setTextAlignment(Qt.AlignTop)
                if table == self.right_table or table == self.left_table:
                    item.setFlags(item.flags() | Qt.TextWordWrap)

                if meta_df is not None and table == self.right_table:
                    try:
                        status = meta_df.loc[(row, col), 'status']

                        if status == 0:
                            color = COLOR_STATUS_0
                        elif status == 1:
                            color = COLOR_STATUS_1
                        elif status == 2:
                            color = COLOR_STATUS_2
                        elif status == 3:
                            color = COLOR_STATUS_3
                    except KeyError:
                        pass

                item.setBackground(color)
                table.setItem(row, col, item)

        table.resizeRowsToContents()

    def right_cell_entered(self, row, column):
        """
        Handles highlighting the source cell based on metadata, triggered by cellClicked.
        """
        self.clear_left_highlight()

        meta_df = self.current_metadata_dataframe
        if meta_df is None or meta_df.empty:
            return

        try:
            meta_data = meta_df.loc[(row, column)]
            status = meta_data['status']

            if status == 0:
                return

            source_file = meta_data['source_file']
            source_row_idx = meta_data['source_row']
            source_col_idx = meta_data['source_col']

            if not source_file or source_row_idx is None or source_col_idx is None:
                return

        except KeyError:
            print(f"Metadata not found for cell ({row}, {column})")
            return

        try:
            df_key = next(k for k, v in self.dataframes.items() if k == source_file)
            df_source = self.dataframes[df_key]

            if self.file_selector.currentText() != df_key:
                self.file_selector.blockSignals(True)
                self.file_selector.setCurrentText(df_key)
                self.file_selector.blockSignals(False)
                self.current_left_dataframe = df_source
                self.populate_table(self.left_table, self.current_left_dataframe)

            if 0 <= source_row_idx < self.left_table.rowCount() and 0 <= source_col_idx < self.left_table.columnCount():
                item = self.left_table.item(source_row_idx, source_col_idx)
                if item:
                    item.setBackground(COLOR_HIGHLIGHT)

                    cell_rect = self.left_table.visualItemRect(item)
                    viewport_rect = self.left_table.viewport().rect()

                    is_row_out_of_view = cell_rect.top() < viewport_rect.top() or \
                                         cell_rect.bottom() > viewport_rect.bottom()

                    is_col_out_of_view = cell_rect.left() < viewport_rect.left() or \
                                         cell_rect.right() > viewport_rect.right()

                    if is_row_out_of_view or is_col_out_of_view:
                        self.left_table.scrollToItem(item, QAbstractItemView.PositionAtCenter)

        except StopIteration:
            print(f"Source file not found: {source_file}")

    def clear_left_highlight(self):
        """Clear all yellow background highlights in left_table."""
        if self.left_table.rowCount() == 0:
            return

        for row in range(self.left_table.rowCount()):
            for col in range(self.left_table.columnCount()):
                item = self.left_table.item(row, col)
                if item:
                    item.setBackground(QColor('white'))

    def eventFilter(self, source, event):
        """Intercepts events for self.right_table to handle the mouse leaving the widget."""
        if source == self.right_table and event.type() == QEvent.Leave:
            self.clear_left_highlight()

        return super().eventFilter(source, event)