#!/usr/bin/env python # -*- coding: utf-8 -*- """ Step5_5 面板 - 水质指数计算 """ import os from pathlib import Path from typing import Dict, List, Union import pandas as pd from PyQt5.QtWidgets import ( QWidget, QVBoxLayout, QGroupBox, QFormLayout, QGridLayout, QHBoxLayout, QLabel, QLineEdit, QComboBox, QCheckBox, QPushButton, QMessageBox, ) from PyQt5.QtCore import Qt from src.gui.components.custom_widgets import FileSelectWidget from src.gui.styles import ModernStylesheet class Step5_5Panel(QWidget): """步骤5.5：水质指数计算""" def __init__(self, parent=None): super().__init__(parent) self.index_checkboxes: Dict[str, QCheckBox] = {} self.csv_columns = [] # 存储CSV文件列名 self.init_ui() def init_ui(self): main_layout = QVBoxLayout() # 标题 # 数据文件选择 data_group = QGroupBox("数据文件") data_layout = QVBoxLayout() # 训练数据CSV文件选择 self.training_data_widget = FileSelectWidget("训练数据CSV文件:", "CSV Files (*.csv)") data_layout.addWidget(self.training_data_widget) # 公式CSV文件选择 self.formula_csv_widget = FileSelectWidget("公式CSV文件:", "CSV Files (*.csv)") data_layout.addWidget(self.formula_csv_widget) # 刷新公式按钮 refresh_layout = QHBoxLayout() self.refresh_button = QPushButton("刷新公式列表") self.refresh_button.clicked.connect(self.refresh_formulas) refresh_layout.addWidget(self.refresh_button) refresh_layout.addStretch() data_layout.addLayout(refresh_layout) data_group.setLayout(data_layout) main_layout.addWidget(data_group) # 公式选择区域 self.formula_group = QGroupBox("选择要计算的公式") formula_outer_layout = QVBoxLayout() # 按钮控制区域 button_layout = QHBoxLayout() self.select_all_btn = QPushButton("全选") self.select_all_btn.clicked.connect(self.select_all_formulas) self.deselect_all_btn = QPushButton("清空") self.deselect_all_btn.clicked.connect(self.deselect_all_formulas) button_layout.addWidget(self.select_all_btn) button_layout.addWidget(self.deselect_all_btn) button_layout.addStretch() formula_outer_layout.addLayout(button_layout) # 公式勾选框网格布局 self.formula_layout = QGridLayout() formula_outer_layout.addLayout(self.formula_layout) self.formula_group.setLayout(formula_outer_layout) main_layout.addWidget(self.formula_group) # 输出文件设置 output_group = QGroupBox("输出设置") output_layout = QVBoxLayout() self.output_file_widget = FileSelectWidget( "输出文件:", "CSV Files (*.csv)", mode="save" ) output_layout.addWidget(self.output_file_widget) output_group.setLayout(output_layout) main_layout.addWidget(output_group) # 启用选项 self.enable_checkbox = QCheckBox("启用此步骤") self.enable_checkbox.setChecked(True) main_layout.addWidget(self.enable_checkbox) # 独立运行按钮 self.run_button = QPushButton("独立运行此步骤") self.run_button.setStyleSheet(ModernStylesheet.get_button_stylesheet('success')) self.run_button.clicked.connect(self.run_step) main_layout.addWidget(self.run_button) # 公式编辑区域 formula_edit_group = QGroupBox("添加自定义公式") formula_edit_layout = QFormLayout() self.formula_name_edit = QLineEdit() # 公式类别下拉选择框 self.formula_category_combo = QComboBox() self.formula_category_combo.addItems([ "chlorophyll_a", "Phycocyanin (BGA_PC)", "Total Nitrogen (TN)", "Total Phosphorus (TP)", "Orthophosphate", "COD", "BOD", "TOC", "Dissolved Oxygen (DO)", "E. coli", "Total Coliforms", "Turbidity", "Total Suspended Solids (TSS)", "Color", "pH", "Temperature", "Conductivity", "Total Dissolved Solids (TDS)" ]) self.formula_category_combo.setEditable(True) # 允许用户输入自定义类别 self.formula_expression_edit = QLineEdit() self.formula_reference_edit = QLineEdit() formula_edit_layout.addRow("公式名称:", self.formula_name_edit) formula_edit_layout.addRow("公式类别:", self.formula_category_combo) formula_edit_layout.addRow("公式表达式:", self.formula_expression_edit) formula_edit_layout.addRow("参考文献:", self.formula_reference_edit) add_button = QPushButton("添加公式") add_button.clicked.connect(self.add_custom_formula) formula_edit_layout.addRow(add_button) formula_edit_group.setLayout(formula_edit_layout) main_layout.addWidget(formula_edit_group) main_layout.addStretch() self.setLayout(main_layout) # 自动加载内置公式文件 formula_csv_path = ( Path(__file__).resolve().parent.parent / "model" / "waterindex.csv" ) if formula_csv_path.is_file(): self.formula_csv_widget.set_path(str(formula_csv_path)) self.refresh_formulas() def refresh_formulas(self): """刷新公式列表""" formula_csv_path = self.formula_csv_widget.get_path() if not formula_csv_path or not os.path.exists(formula_csv_path): QMessageBox.warning(self, "警告", "请先选择有效的公式CSV文件") return try: # 清除现有的勾选框 for checkbox in self.index_checkboxes.values(): self.formula_layout.removeWidget(checkbox) checkbox.deleteLater() self.index_checkboxes.clear() # 读取公式CSV文件 df = pd.read_csv(formula_csv_path) if df.empty or 'Formula_Name' not in df.columns: QMessageBox.warning(self, "警告", "公式CSV文件格式不正确") return # 获取所有公式名称（跳过第一行） formula_names = df['Formula_Name'].tolist()[1:] # 创建3列布局的勾选框 row, col = 0, 0 for formula_name in formula_names: if pd.isna(formula_name) or not formula_name.strip(): continue checkbox = QCheckBox(formula_name.strip()) checkbox.setChecked(True) self.index_checkboxes[formula_name.strip()] = checkbox self.formula_layout.addWidget(checkbox, row, col) col += 1 if col >= 3: # 每行3列 col = 0 row += 1 except Exception as e: QMessageBox.critical(self, "错误", f"读取公式文件失败: {str(e)}") def add_custom_formula(self): """添加自定义公式到公式CSV文件""" formula_csv_path = self.formula_csv_widget.get_path() if not formula_csv_path: QMessageBox.warning(self, "警告", "请先选择公式CSV文件") return formula_name = self.formula_name_edit.text().strip() formula_category = self.formula_category_combo.currentText().strip() formula_expression = self.formula_expression_edit.text().strip() formula_reference = self.formula_reference_edit.text().strip() if not all([formula_name, formula_category, formula_expression]): QMessageBox.warning(self, "警告", "请填写公式名称、类别和表达式") return try: # 读取现有公式文件或创建新文件 if os.path.exists(formula_csv_path): df = pd.read_csv(formula_csv_path) else: df = pd.DataFrame(columns=['Formula_Name', 'Category', 'Formula', 'Reference']) # 添加新公式 new_row = pd.DataFrame({ 'Formula_Name': [formula_name], 'Category': [formula_category], 'Formula': [formula_expression], 'Reference': [formula_reference] }) df = pd.concat([df, new_row], ignore_index=True) # 保存文件 df.to_csv(formula_csv_path, index=False, encoding='utf-8') # 清空输入框 self.formula_name_edit.clear() self.formula_category_combo.setCurrentIndex(0) # 重置到第一个选项 self.formula_expression_edit.clear() self.formula_reference_edit.clear() # 刷新公式列表 self.refresh_formulas() QMessageBox.information(self, "成功", "公式添加成功") except Exception as e: QMessageBox.critical(self, "错误", f"添加公式失败: {str(e)}") def get_config(self) -> Dict[str, Union[List[str], str, bool]]: """获取配置""" selected = [ name for name, checkbox in self.index_checkboxes.items() if checkbox.isChecked() ] output_path = self.output_file_widget.get_path() return { 'training_spectra_path': self.training_data_widget.get_path() or None, 'formula_csv_file': self.formula_csv_widget.get_path() or None, 'formula_names': selected, 'output_file': output_path or None, 'enabled': self.enable_checkbox.isChecked() } def set_config(self, config): """设置配置""" if 'training_spectra_path' in config: self.training_data_widget.set_path(config['training_spectra_path']) if 'formula_csv_file' in config: self.formula_csv_widget.set_path(config['formula_csv_file']) self.refresh_formulas() if 'formula_names' in config: selected_formulas = set(config['formula_names']) for name, checkbox in self.index_checkboxes.items(): checkbox.setChecked(name in selected_formulas) if 'output_file' in config and config['output_file']: self.output_file_widget.set_path(config['output_file']) elif 'output_filename' in config and config['output_filename']: self.output_file_widget.set_path(config['output_filename']) if 'enabled' in config: self.enable_checkbox.setChecked(config['enabled']) def update_from_config(self, work_dir=None, pipeline=None): """从全局配置自动填充训练数据和输出路径 Args: work_dir: 工作目录路径 pipeline: Pipeline 实例（未使用，保留接口兼容性） """ if work_dir: self.work_dir = work_dir elif hasattr(self, 'work_dir') and self.work_dir: pass else: self.work_dir = None # 1. 自动填入训练数据路径（从 Step5 的输出中获取） # 优先级：直接 widget > pipeline.step_outputs 回退 main_window = self.window() if hasattr(main_window, 'step5_panel'): # 优先直接从 Step5 的输出 widget 读取（已运行的最新输出） step5_output = main_window.step5_panel.output_file.get_path() if step5_output: # 若为相对路径，使用 work_dir 合成为绝对路径 if not os.path.isabs(step5_output): step5_output = os.path.join(self.work_dir or '', step5_output).replace('\\', '/') self.training_data_widget.set_path(step5_output) else: # 退而求其次，使用 Step5 的输入 CSV step5_csv = main_window.step5_panel.csv_file.get_path() if step5_csv: # 若为相对路径，使用 work_dir 合成为绝对路径 if not os.path.isabs(step5_csv): step5_csv = os.path.join(self.work_dir or '', step5_csv).replace('\\', '/') self.training_data_widget.set_path(step5_csv) # 如果上述都没找到，尝试从 pipeline.step_outputs 回退 if not self.training_data_widget.get_path() and pipeline and hasattr(pipeline, 'step_outputs'): step5_outputs = getattr(pipeline, 'step_outputs', {}).get('step5', {}) training_path = step5_outputs.get('training_spectra') if training_path: self.training_data_widget.set_path(training_path) # 2. 自动填入输出文件的绝对路径 if self.work_dir: output_abs = os.path.join(self.work_dir, "6_water_quality_indices", "training_spectra_indices.csv").replace('\\', '/') self.output_file_widget.set_path(output_abs) def is_enabled(self) -> bool: return self.enable_checkbox.isChecked() def select_all_formulas(self): """全选所有公式""" for checkbox in self.index_checkboxes.values(): checkbox.setChecked(True) def deselect_all_formulas(self): """清空所有公式""" for checkbox in self.index_checkboxes.values(): checkbox.setChecked(False) def run_step(self): """独立运行步骤5.5：计算水质指数。 动态根据输入 CSV 文件名生成输出文件名，自动填入 output_file_widget。 例如：training_spectra.csv → training_spectra_indices.csv sampling_spectra.csv → sampling_spectra_indices.csv """ # 验证输入 training_csv_path = self.training_data_widget.get_path() formula_csv_path = self.formula_csv_widget.get_path() if not training_csv_path: QMessageBox.warning(self, "输入验证失败", "请选择训练数据CSV文件") return if not formula_csv_path: QMessageBox.warning(self, "输入验证失败", "请选择公式CSV文件") return if not os.path.exists(training_csv_path): QMessageBox.warning(self, "输入验证失败", "训练数据CSV文件不存在") return if not os.path.exists(formula_csv_path): QMessageBox.warning(self, "输入验证失败", "公式CSV文件不存在") return # 动态生成输出文件：自动拼接 _indices 后缀 input_name = Path(training_csv_path).stem dynamic_output = f"{input_name}_indices.csv" # 合成完整绝对路径（优先使用 work_dir，其次从 training_csv_path 推导） work_dir = getattr(self, 'work_dir', None) if work_dir: dynamic_output = os.path.join( work_dir, "6_water_quality_indices", dynamic_output ).replace('\\', '/') self.output_file_widget.set_path(dynamic_output) # 获取配置 config = self.get_config() # 调用GUI的run_single_step方法 parent = self.parent() while parent and not hasattr(parent, 'run_single_step'): parent = parent.parent() if parent and hasattr(parent, 'run_single_step'): parent.run_single_step('step5_5', {'step5_5': config}) else: QMessageBox.critical(self, "错误", "无法找到父级GUI对象")