WQ_GUI/src/core/workspace_manager.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
工作空间管理器

负责工作目录文件扫描、步骤输出路径发现、配置裁剪等业务逻辑，
与 GUI 组件解耦，不直接引用任何 UI 类。
"""

import copy
from pathlib import Path

from src.gui.core.event_bus import global_event_bus


class WorkspaceManager:
    """管理步骤默认输出路径、文件扫描与配置裁剪"""

    # 白名单：科学数据格式后缀
    SCIENTIFIC_EXTENSIONS = {'.dat', '.tif', '.tiff', '.shp'}
    # 临时文件关键词黑名单
    TMP_KEYWORDS = ('__tmp', '_tmp')
    # 掩膜类型集合
    MASK_TYPES = {'water_mask', 'glint_mask', 'boundary_mask'}

    def __init__(self):
        self.step_default_outputs = {
            'step1': {'water_mask': [
                "1_water_mask/water_mask_out.dat",
                "1_water_mask/water_mask_from_ndwi.dat",
                "1_water_mask/water_mask_from_shp.dat",
            ]},
            'step2': {'glint_mask': "2_Glint_Detection/severe_glint_area.dat"},
            'step3': {'deglint_image': [
                "3_deglint/deglint_image.bsq",
                "3_deglint/deglint_goodman.bsq",
            ]},
            'step4_sampling': {'sampling_points': "4_sampling/sampling_spectra.csv"},
            'step5_clean': {'processed_data': "5_Data_Cleaning/processed_data.csv"},
            'step6_feature': {'training_spectra': "6_Spectral_Feature_Extraction/training_spectra.csv"},
            'step7_index': {'training_spectra_indices': "7_Water_Quality_Indices/training_spectra_indices.csv"},
            'step8_ml_train': {'Supervised_Model_Training': "8_Supervised_Model_Training/"},
            'step9_ml_predict': {'9_ML_Prediction': "9_ML_Prediction/"},
            'step10_watercolor': {'WaterIndex_CSV': "10_WaterIndex_CSV/"},
            'step11_map': {'11_Thematic_Map': "11_Thematic_Map/"},
        }
        self.step_outputs = {}

    # ★ 2026-07-01：切换工作目录时清空缓存，防止旧目录的 OutputUpdated 重放
    def clear_all_outputs(self):
        """清空所有缓存的步骤输出路径。工作目录切换时调用。"""
        self.step_outputs.clear()

    def _publish_outputs(self, step_id: str, outputs: dict):
        """将发现的产出发布到 EventBus。

        Args:
            step_id: 面板 step_id（如 'step1', 'step5_clean'）
            outputs: {output_type: path_str}
        """
        for output_type, path in outputs.items():
            if path:
                global_event_bus.publish('OutputUpdated', {
                    'step_id': step_id,
                    'output_type': output_type,
                    'path': path,
                })

    @staticmethod
    def _is_scientific_mask(path_str):
        """白名单判断：只有 .dat .tif .tiff .shp 才算科学数据格式"""
        p = Path(path_str)
        name_lower = str(path_str).lower()
        if any(kw in name_lower for kw in WorkspaceManager.TMP_KEYWORDS):
            return False
        return p.suffix.lower() in WorkspaceManager.SCIENTIFIC_EXTENSIONS

    def find_step_output(self, work_path, step_id, output_type, ref_img_path=None):
        """查找指定步骤的输出文件

        Args:
            work_path: 工作目录 Path 对象
            step_id: 步骤 ID
            output_type: 输出类型（如 'water_mask', 'deglint_image' 等）
            ref_img_path: 参考影像路径（仅 output_type='reference_img' 时需要）

        Returns:
            找到的文件路径字符串，或 None
        """
        if step_id not in self.step_default_outputs:
            return None

        raw = self.step_default_outputs[step_id]

        rel_path = None
        if isinstance(raw, str):
            rel_path = raw
        elif isinstance(raw, dict):
            rel_path = raw.get(output_type) or list(raw.values())[0]

        if not rel_path:
            return None

        # 特殊处理：从 step_outputs 记录中查找实际输出路径
        if step_id in self.step_outputs:
            actual_outputs = self.step_outputs[step_id]
            if output_type in actual_outputs:
                candidate = actual_outputs[output_type]
                if output_type in self.MASK_TYPES and not self._is_scientific_mask(candidate):
                    pass
                else:
                    return candidate

        if output_type == 'water_mask':
            if isinstance(rel_path, list):
                for candidate in rel_path:
                    mask_path = work_path / candidate
                    if mask_path.exists():
                        return str(mask_path)
            elif rel_path:
                mask_path = work_path / rel_path
                if mask_path.exists():
                    return str(mask_path)
        elif output_type == 'reference_img':
            if ref_img_path and Path(ref_img_path).exists():
                return ref_img_path
        elif output_type == 'deglint_image':
            if isinstance(rel_path, list):
                for candidate in rel_path:
                    deglint_path = work_path / candidate
                    if deglint_path.exists():
                        return str(deglint_path)
            elif rel_path:
                deglint_path = work_path / rel_path
                if deglint_path.exists():
                    return str(deglint_path)
            deglint_dir = work_path / "3_deglint"
            if deglint_dir.exists():
                for file_path in deglint_dir.glob("deglint_*.bsq"):
                    return str(file_path)
                for file_path in deglint_dir.glob("interpolated_*.bsq"):
                    return str(file_path)
        elif isinstance(rel_path, str):
            if rel_path.endswith('/'):
                output_path = work_path / rel_path.rstrip('/')
                if output_path.exists() and output_path.is_dir():
                    return str(output_path)
            else:
                output_path = work_path / rel_path
                if output_path.exists():
                    return str(output_path)

        return None

    def scan_work_directory_for_files(self, work_path):
        """扫描工作目录，自动发现各步骤的输出文件

        Returns:
            discovered_outputs: dict, {step_id: {output_type: path_str}}
        """
        discovered_outputs = {}

        subdirs = {
            '1_water_mask': 'step1',
            '2_Glint_Detection': 'step2',
            '3_deglint': 'step3',
            '4_sampling': 'step4_sampling',
            '5_Data_Cleaning': 'step5_clean',
            '6_Spectral_Feature_Extraction': 'step6_feature',
            '7_Water_Quality_Indices': 'step7_index',
            '8_Supervised_Model_Training': 'step8_ml_train',
            '9_ML_Prediction': 'step9_ml_predict',
            '10_WaterIndex_CSV': 'step10_watercolor',
            '11_Thematic_Map': 'step11_map',
            '12_visualization': ['step12_viz', 'step13_report'],
        }

        for subdir, step_ids in subdirs.items():
            subdir_path = work_path / subdir
            if not subdir_path.exists():
                continue

            if isinstance(step_ids, str):
                step_ids = [step_ids]

            for file_path in subdir_path.rglob('*'):
                if file_path.is_file():
                    file_name = file_path.name.lower()
                    file_stem = file_path.stem.lower()

                    for step_id in step_ids:
                        if step_id not in discovered_outputs:
                            discovered_outputs[step_id] = {}

                        # 2026-06-30 加强匹配：用更精确的边界匹配替代简单的子串包含
                        # water_mask 匹配：必须是独立文件名的一部分（如 water_mask_out.dat, water_mask_from_ndwi.dat）
                        if step_id == 'step1' and self._is_scientific_mask(file_path):
                            if ('water_mask' in file_stem
                                and 'glint' not in file_stem):
                                discovered_outputs[step_id]['water_mask'] = str(file_path)
                        # glint_mask 匹配：必须同时含 glint 且文件名主体以 mask 或 area 结尾
                        elif step_id == 'step2' and self._is_scientific_mask(file_path):
                            if ('glint' in file_stem and
                                (file_stem.endswith('mask') or file_stem.endswith('area')
                                 or 'severe_glint' in file_stem)):
                                discovered_outputs[step_id]['glint_mask'] = str(file_path)
                        elif 'deglint' in file_stem and step_id == 'step3':
                            discovered_outputs[step_id]['deglint_image'] = str(file_path)
                        elif file_name == 'processed_data.csv' and step_id == 'step5_clean':
                            discovered_outputs[step_id]['processed_data'] = str(file_path)
                        elif ('training_spectra' in file_stem and file_path.suffix == '.csv'
                              and step_id == 'step6_feature'):
                            discovered_outputs[step_id]['training_spectra'] = str(file_path)
                        elif ('water_quality_indices' in file_stem and file_path.suffix == '.csv'
                              and step_id == 'step7_index'):
                            discovered_outputs[step_id]['water_indices'] = str(file_path)
                        elif file_name == 'sampling_spectra.csv' and step_id == 'step4_sampling':
                            discovered_outputs[step_id]['sampling_points'] = str(file_path)
                        elif file_name.endswith('.csv') and step_id in ['step9_ml_predict', 'step11_map', 'step12_viz']:
                            discovered_outputs[step_id]['predictions'] = str(file_path)

        for step_id, outputs in discovered_outputs.items():
            if step_id not in self.step_outputs:
                self.step_outputs[step_id] = {}
            self.step_outputs[step_id].update(outputs)
            # ★ 发布 EventBus 事件，驱动下游面板自动填充
            self._publish_outputs(step_id, outputs)

        return discovered_outputs

    def update_step_outputs(self, step_name, work_path):
        """更新指定步骤的输出路径记录并发布 EventBus 事件。"""
        if step_name not in self.step_default_outputs:
            return

        step_outputs = self.step_default_outputs[step_name]
        published = {}

        for output_type, relative_path in step_outputs.items():
            if isinstance(relative_path, list):
                for candidate in relative_path:
                    output_path = work_path / candidate
                    if output_path.exists():
                        path_str = str(output_path)
                        self.step_outputs.setdefault(step_name, {})[output_type] = path_str
                        published[output_type] = path_str
                        break
            elif '*' in relative_path:
                pattern_path = work_path / relative_path.replace('*', '*')
                matching_files = list(pattern_path.parent.glob(pattern_path.name))
                if matching_files:
                    latest_file = max(matching_files, key=lambda p: p.stat().st_mtime)
                    path_str = str(latest_file)
                    self.step_outputs.setdefault(step_name, {})[output_type] = path_str
                    published[output_type] = path_str
            else:
                output_path = work_path / relative_path
                if output_path.exists():
                    path_str = str(output_path)
                    self.step_outputs.setdefault(step_name, {})[output_type] = path_str
                    published[output_type] = path_str

        if published:
            self._publish_outputs(step_name, published)

    @staticmethod
    def prune_config_for_prediction_mode(config: dict) -> dict:
        """Prediction-only 模式：禁用训练相关步骤，保留预测和成图步骤。

        2026-06-30 修复：步骤 ID 从旧的 PIPELINE_STEPS 体系改为 PANEL_REGISTRY 体系，
        确保与 get_current_config() 返回的 key 一致，避免 enabled: False 写入无效 key。

        被禁用的 step dict 中统一写入 'enabled': False，
        这些配置最终传给 PipelineRunner，Runner 会跳过它们。
        同时，被跳过的步骤的 required_input_files 在 build_missing_items
        中不会被检查，从而自然规避了"CSV 缺失"等训练模式下的误报。

        Args:
            config: 完整配置字典（来自 get_current_config，key 为 PANEL_REGISTRY step_id）

        Returns:
            裁剪后的 config（深拷贝，原 config 不被修改）
        """
        cfg = copy.deepcopy(config)

        training_steps = [
            "step4_sampling",
            "step5_clean",
            "step6_feature",
            "step7_index",
            "step8_ml_train",
        ]
        for step_id in training_steps:
            step_cfg = cfg.setdefault(step_id, {})
            step_cfg["enabled"] = False

        return cfg