AcidMap/Cd_Prediction_Integrated_System/analysis/visualization.py

"""
可视化模块
Visualization Module

基于通用绘图模块 app.utils.mapping_utils 的封装
提供与原有接口兼容的可视化功能
"""

import os
import sys
import logging

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

# 导入通用绘图模块
from app.utils.mapping_utils import MappingUtils, COLORMAPS

import config

class Visualizer:
    """
    可视化器
    负责创建栅格地图和直方图
    基于通用绘图模块的封装，保持原有接口兼容性
    """
    
    def __init__(self):
        """
        初始化可视化器
        """
        self.logger = logging.getLogger(__name__)
        self.mapping_utils = MappingUtils()
        

    def create_raster_map(self, 
                         shp_path=None, 
                         tif_path=None, 
                         color_map_name=None,
                         title_name="Prediction Cd",
                         output_path=None,
                         output_size=None,
                         high_res=False):
        """
        创建栅格地图
        
        @param shp_path: 输入的矢量数据的路径
        @param tif_path: 输入的栅格数据的路径
        @param color_map_name: 使用的色彩方案
        @param title_name: 输出数据的图的名称
        @param output_path: 输出保存的图片的路径
        @param output_size: 图片尺寸
        @param high_res: 是否使用高分辨率输出（默认False，DPI=300）
        @return: 输出图片路径
        """
        try:
            # 使用默认值
            if shp_path is None:
                shp_path = config.ANALYSIS_CONFIG["boundary_shp"]
            if tif_path is None:
                tif_path = config.ANALYSIS_CONFIG["output_raster"]
            if color_map_name is None:
                color_map_name = config.VISUALIZATION_CONFIG["color_maps"][config.VISUALIZATION_CONFIG["default_colormap"]]
            if output_path is None:
                output_path = os.path.join(config.OUTPUT_PATHS["figures_dir"], "Prediction_results")
            if output_size is None:
                output_size = config.VISUALIZATION_CONFIG["figure_size"]
            
            # 检查边界文件是否存在
            if shp_path and not os.path.exists(shp_path):
                self.logger.warning(f"边界文件不存在: {shp_path}，将跳过边界绘制")
                shp_path = None
                
            # 设置DPI
            dpi = 600 if high_res else config.VISUALIZATION_CONFIG["dpi"]
            
            # 调用通用绘图模块
            return self.mapping_utils.create_raster_map(
                shp_path=shp_path,
                tif_path=tif_path,
                output_path=output_path,
                colormap=color_map_name,
                title=title_name,
                output_size=output_size,
                dpi=dpi
            )
            
        except Exception as e:
            self.logger.error(f"栅格地图创建失败: {str(e)}")
            raise
    
    def create_histogram(self, 
                        file_path=None,
                        figsize=None,
                        xlabel='Cd content',
                        ylabel='Frequency',
                        title='County level Cd Frequency',
                        save_path=None,
                        high_res=False):
        """
        绘制GeoTIFF文件的直方图
        
        @param file_path: GeoTIFF 文件路径
        @param figsize: 图像尺寸，如 (10, 6)
        @param xlabel: 横坐标标签
        @param ylabel: 纵坐标标签
        @param title: 图标题
        @param save_path: 可选，保存图片路径（含文件名和扩展名）
        @param high_res: 是否使用高分辨率输出（默认False，DPI=300）
        @return: 输出图片路径
        """
        try:
            # 使用默认值
            if file_path is None:
                file_path = config.ANALYSIS_CONFIG["output_raster"]
            if figsize is None:
                figsize = (6, 6)
            if save_path is None:
                save_path = os.path.join(config.OUTPUT_PATHS["figures_dir"], "Prediction_frequency.jpg")
            
            # 设置DPI
            dpi = 600 if high_res else config.VISUALIZATION_CONFIG["dpi"]
            
            # 调用通用绘图模块
            return self.mapping_utils.create_histogram(
                file_path=file_path,
                save_path=save_path,
                figsize=figsize,
                xlabel=xlabel,
                ylabel=ylabel,
                title=title,
                dpi=dpi
            )
            
        except Exception as e:
            self.logger.error(f"直方图创建失败: {str(e)}")
            raise

if __name__ == "__main__":
    # 测试代码
    visualizer = Visualizer()
    
    # 测试栅格地图创建
    try:
        map_output = visualizer.create_raster_map()
        print(f"栅格地图创建完成: {map_output}")
    except Exception as e:
        print(f"栅格地图创建失败: {e}")
    
    # 测试直方图创建
    try:
        histogram_output = visualizer.create_histogram()
        print(f"直方图创建完成: {histogram_output}")
    except Exception as e:
        print(f"直方图创建失败: {e}")