AcidMap/app/services/cd_flux_service.py

import os
import csv
import logging
import sys
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
from pathlib import Path
from typing import Dict, Any

from sqlalchemy import func
from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models import FarmlandData, FluxCdOutputData, FluxCdInputData
from app.utils.mapping_utils import MappingUtils, csv_to_raster_workflow

# 配置日志
from app.log.logger import get_logger

logger = get_logger(__name__)


def get_base_dir():
    """获取基础目录路径（与土地数据处理函数一致）"""
    if getattr(sys, 'frozen', False):
        # 打包后的可执行文件
        return os.path.dirname(sys.executable)
    else:
        # 脚本运行模式
        return os.path.dirname(os.path.abspath(__file__))


def get_static_dir():
    """获取静态资源目录"""
    base_dir = get_base_dir()
    return os.path.join(base_dir, "..", "static", "cd_flux")


def get_default_boundary_shp():
    """获取默认边界SHP文件路径"""
    static_dir = get_static_dir()
    # 尝试几个可能的边界文件路径
    possible_paths = [
        os.path.join(static_dir, "lechang.shp"),
    ]

    for path in possible_paths:
        if os.path.exists(path):
            return path

    return None


class FluxCdVisualizationService:
    """
    农田Cd通量可视化服务类

    @description: 提供基于农田位置和Cd通量数据的空间分布图和直方图生成功能
    @version: 1.1.2 (修复绘图问题)
    """

    def __init__(self, db: Session = None):
        """
        初始化可视化服务

        @param db: 可选的数据库会话对象
        """
        self.logger = logger
        self.db = db
        self.mapper = MappingUtils()

    def generate_cd_input_flux_map(self, output_dir: str = None, boundary_shp: str = None) -> Dict[str, Any]:
        """
        生成输入镉通量(In_Cd)的空间分布图和直方图

        @param output_dir: 输出文件目录
        @param boundary_shp: 边界Shapefile文件路径
        @return: 包含输出文件路径的字典
        """
        return self._generate_cd_flux_map(field='in_cd', title_prefix="input",
                                          output_dir=output_dir, boundary_shp=boundary_shp)

    def generate_cd_output_flux_map(self, output_dir: str = None, boundary_shp: str = None) -> Dict[str, Any]:
        """
        生成输出镉通量(Out_Cd)的空间分布图和直方图

        @param output_dir: 输出文件目录
        @param boundary_shp: 边界Shapefile文件路径
        @return: 包含输出文件路径的字典
        """
        return self._generate_cd_flux_map(field='out_cd', title_prefix="output",
                                          output_dir=output_dir, boundary_shp=boundary_shp)

    def generate_cd_net_flux_map(self, output_dir: str = None, boundary_shp: str = None) -> Dict[str, Any]:
        """
        生成净镉通量(Net_Cd)的空间分布图和直方图

        @param output_dir: 输出文件目录
        @param boundary_shp: 边界Shapefile文件路径
        @return: 包含输出文件路径的字典
        """
        return self._generate_cd_flux_map(field='net_cd', title_prefix="net",
                                          output_dir=output_dir, boundary_shp=boundary_shp)

    def generate_end_cd_map(self, output_dir: str = None, boundary_shp: str = None) -> Dict[str, Any]:
        """
        生成当年Cd浓度(End_Cd)的空间分布图和直方图
        """
        return self._generate_cd_flux_map(
            field='end_cd',
            title_prefix="End",
            output_dir=output_dir,
            boundary_shp=boundary_shp
        )

    def _generate_cd_flux_map(self, field: str, title_prefix: str,
                              output_dir: str = None, boundary_shp: str = None) -> Dict[str, Any]:
        """
        通用的Cd通量地图生成方法（直接从数据库获取数据）
        """
        try:
            # 字段配置
            field_config = {
                'in_cd': {"title": "Input", "unit": "g/ha/a", "label": "input Cd flux"},
                'out_cd': {"title": "Output", "unit": "g/ha/a", "label": "output Cd flux"},
                'net_cd': {"title": "Net", "unit": "g/ha/a", "label": "net Cd flux"},
                'end_cd': {"title": "End", "unit": "mg/kg", "label": "end Cd concentration"}
            }

            config = field_config.get(field, field_config['in_cd'])

            # 获取数据库会话
            db = self.db if self.db else SessionLocal()
            should_close = self.db is None

            # 从数据库查询数据
            query = db.query(
                FarmlandData.lon,
                FarmlandData.lan,
                getattr(FluxCdOutputData, field)
            ).join(
                FluxCdOutputData,
                (FarmlandData.farmland_id == FluxCdOutputData.farmland_id) &
                (FarmlandData.sample_id == FluxCdOutputData.sample_id)
            ).filter(
                getattr(FluxCdOutputData, field) != None  # 排除空值
            )

            data = query.all()

            if not data:
                if should_close:
                    db.close()
                return {
                    "success": False,
                    "message": f"数据库中未找到任何{config['title']}Cd数据",
                    "data": None
                }

            # 设置输出目录
            static_dir = get_static_dir()
            if output_dir is None:
                output_dir = static_dir
            os.makedirs(output_dir, exist_ok=True)

            # 设置边界SHP
            if boundary_shp is None:
                boundary_shp = get_default_boundary_shp()

            # 生成CSV文件
            base_name = f"fluxcd_{field}"
            csv_path = os.path.join(output_dir, f"{base_name}.csv")
            self._generate_csv(data, csv_path, field_names=['lon', 'lan', field])

            # 使用模板TIFF
            template_tif = os.path.join(static_dir, "meanTemp.tif")
            if not os.path.exists(template_tif):
                raise FileNotFoundError("未找到模板TIFF文件")

            # 转换为栅格
            raster_path = os.path.join(output_dir, f"{base_name}_raster.tif")
            workflow_result = csv_to_raster_workflow(
                csv_file=csv_path,
                template_tif=template_tif,
                output_dir=output_dir,
                resolution_factor=4.0,
                interpolation_method='linear',
                field_name=field,
                lon_col=0,
                lat_col=1,
                value_col=2,
                enable_interpolation=True
            )

            raster_path = workflow_result['raster']
            stats = workflow_result['statistics']

            # 创建地图
            map_path = os.path.join(output_dir, f"{base_name}_map")
            map_file = self.mapper.create_raster_map(
                shp_path=boundary_shp,
                tif_path=raster_path,
                output_path=map_path,
                colormap='green_yellow_red_purple',
                title=f"{config['title']} Cd {field.split('_')[0]} map",
                output_size=12,
                dpi=300,
                enable_interpolation=False
            )

            # 创建直方图
            histogram_path = self.mapper.create_histogram(
                raster_path,
                save_path=os.path.join(output_dir, f"{base_name}_histogram.jpg"),
                xlabel=f"{config['label']} ({config['unit']})",
                ylabel='frequency',
                title=f"{config['title']} Cd {field.split('_')[0]} histogram",
                bins=100
            )

            result = {
                "success": True,
                "message": f"成功生成{config['title']}Cd数据可视化结果",
                "data": {
                    "csv": csv_path,
                    "raster": raster_path,
                    "map": map_file,
                    "histogram": histogram_path,
                    "statistics": stats,
                    "boundary_used": boundary_shp if boundary_shp else "无"
                }
            }

            if should_close:
                db.close()

            return result

        except Exception as e:
            self.logger.error(f"生成Cd数据可视化结果时发生错误: {str(e)}", exc_info=True)
            return {
                "success": False,
                "message": f"生成失败: {str(e)}",
                "data": None
            }

    def update_from_csv(self, csv_file_path: str):
        """包装更新服务方法"""
        update_service = FluxCdUpdateService(db=self.db)
        return update_service.update_from_csv(csv_file_path)

    def _fetch_fluxcd_data(self, db: Session, field: str = 'in_cd') -> list:
        """
        从数据库查询需要的数据

        @param db: 数据库会话
        @param field: 要查询的字段（'in_cd', 'out_cd'或'net_cd'）
        @returns: 查询结果列表
        """
        try:
            # 查询Farmland_data和FluxCd_output_data的联合数据
            query = db.query(
                FarmlandData.lon,
                FarmlandData.lan,
                getattr(FluxCdOutputData, field)
            ).join(
                FluxCdOutputData,
                (FarmlandData.farmland_id == FluxCdOutputData.farmland_id) &
                (FarmlandData.sample_id == FluxCdOutputData.sample_id)
            )

            return query.all()

        except Exception as e:
            self.logger.error(f"查询农田Cd通量数据时发生错误: {str(e)}")
            return []

    def _generate_csv(self, data: list, output_path: str, field_names: list = None):
        """
        将查询结果生成CSV文件

        @param data: 查询结果列表
        @param output_path: 输出CSV文件路径
        @param field_names: CSV列名列表
        """
        if field_names is None:
            field_names = ['lon', 'lan', 'In_Cd']

        try:
            with open(output_path, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                # 写入表头
                writer.writerow(field_names)

                # 写入数据
                for row in data:
                    writer.writerow(row)

            self.logger.info(f"CSV文件已生成: {output_path}")

        except Exception as e:
            self.logger.error(f"生成CSV文件时发生错误: {str(e)}")
            raise


class FluxCdUpdateService:
    """
    农田Cd通量数据更新服务类
    处理用户上传的CSV文件并更新数据库中的通量值
    """

    def __init__(self, db: Session = None):
        self.db = db if db else SessionLocal()
        self.logger = logging.getLogger(__name__)

    def update_from_csv(self, csv_file_path: str) -> Dict[str, Any]:
        """
        从CSV文件更新Cd通量数据

        @param csv_file_path: CSV文件路径
        @return: 更新结果字典
        """
        try:
            # 读取CSV文件
            df = self._read_csv(csv_file_path)

            # 验证CSV格式
            self._validate_csv(df)

            # 更新数据库
            update_count = 0
            create_count = 0
            error_count = 0

            for _, row in df.iterrows():
                try:
                    # 查找匹配的农田样点
                    farmland = self._find_farmland(row['lon'], row['lan'])

                    if farmland:
                        # 更新现有记录
                        input_updated = self._update_input_data(farmland, row)
                        if input_updated:
                            self._update_output_data(farmland)
                            update_count += 1
                    else:
                        # 创建新记录
                        if self._create_new_records(row):
                            create_count += 1
                        else:
                            error_count += 1
                except Exception as e:
                    self.logger.error(f"处理记录失败: {str(e)}")
                    error_count += 1

            self.db.commit()

            return {
                "success": True,
                "message": f"成功更新 {update_count} 条记录，创建 {create_count} 条新记录，{error_count} 条失败",
                "updated_count": update_count,
                "created_count": create_count,
                "error_count": error_count
            }

        except Exception as e:
            self.db.rollback()
            self.logger.error(f"更新失败: {str(e)}", exc_info=True)
            return {
                "success": False,
                "message": f"更新失败: {str(e)}"
            }
        finally:
            self.db.close()

    def _create_new_records(self, row: pd.Series):
        """为新的经纬度创建完整的记录"""
        try:
            self.logger.info(f"创建新农田记录: 经度={row['lon']}, 纬度={row['lan']}")

            # 获取下一个可用的 Farmland_ID
            max_id = self.db.query(func.max(FarmlandData.farmland_id)).scalar()
            new_farmland_id = (max_id or 0) + 1

            # 创建新的农田样点记录
            farmland = FarmlandData(
                farmland_id=new_farmland_id,
                lon=row['lon'],
                lan=row['lan'],
                type=0.0,  # 默认为旱地
                geom=f"POINT({row['lon']} {row['lan']})"
            )
            self.db.add(farmland)
            self.db.flush()  # 确保ID被分配
            self.logger.info(f"农田记录创建成功: Farmland_ID={farmland.farmland_id}, Sample_ID={farmland.sample_id}")

            # 创建新的输入通量记录
            input_data = FluxCdInputData(
                farmland_id=farmland.farmland_id,
                sample_id=farmland.sample_id,
                atmospheric_deposition=row['DQCJ_Cd'],
                irrigation_input=row['GGS_Cd'],
                agro_chemicals_input=row['NCP_Cd'],
                # 设置合理的默认值
                initial_cd=0.0,
                groundwater_leaching=row.get('DX_Cd', 0.023),  # 添加默认值处理
                surface_runoff=row.get('DB_Cd', 0.368),
                grain_removal=row.get('ZL_Cd', 0.0),
                straw_removal=row.get('JG_Cd', 0.0)
            )
            self.db.add(input_data)
            self.logger.info("输入通量记录创建成功")

            # 计算输入总通量
            in_cd = row['DQCJ_Cd'] + row['GGS_Cd'] + row['NCP_Cd']

            # 计算输出总通量
            out_cd = (input_data.groundwater_leaching +
                      input_data.surface_runoff +
                      input_data.grain_removal +
                      input_data.straw_removal)

            # 创建新的输出通量记录
            output_data = FluxCdOutputData(
                farmland_id=farmland.farmland_id,
                sample_id=farmland.sample_id,
                in_cd=in_cd,
                out_cd=out_cd,
                net_cd=in_cd - out_cd,
                end_cd=0.0  # 默认当年Cd浓度
            )
            self.db.add(output_data)
            self.logger.info("输出通量记录创建成功")

            self.logger.info(f"所有记录创建完成: Farmland_ID={farmland.farmland_id}, Sample_ID={farmland.sample_id}")
            return True

        except Exception as e:
            self.logger.error(f"创建新记录失败: {str(e)}")
            return False

    def _read_csv(self, file_path: str) -> pd.DataFrame:
        """读取CSV文件并验证基本格式"""
        try:
            df = pd.read_csv(file_path)
            self.logger.info(f"成功读取CSV文件: {file_path}")
            return df
        except Exception as e:
            raise ValueError(f"CSV文件读取失败: {str(e)}")

    def _validate_csv(self, df: pd.DataFrame):
        """验证CSV文件包含必要的列"""
        required_columns = {'lon', 'lan', 'DQCJ_Cd', 'GGS_Cd', 'NCP_Cd'}
        # 添加输出通量可选字段
        output_fields = {'DX_Cd', 'DB_Cd', 'ZL_Cd', 'JG_Cd'}

        if not required_columns.issubset(df.columns):
            missing = required_columns - set(df.columns)
            raise ValueError(f"CSV缺少必要列: {', '.join(missing)}")

        # 检查是否包含输出通量字段（可选）
        if not output_fields.issubset(df.columns):
            missing_output = output_fields - set(df.columns)
            self.logger.warning(f"CSV缺少输出通量字段: {', '.join(missing_output)}，将使用默认值")

    def _find_farmland(self, lon: float, lan: float) -> FarmlandData:
        """根据经纬度查找农田样点"""
        # 使用容差匹配（0.001度≈100米）
        tol = 0.001
        return self.db.query(FarmlandData).filter(
            FarmlandData.lon.between(lon - tol, lon + tol),
            FarmlandData.lan.between(lan - tol, lan + tol)
        ).first()

    def _update_input_data(self, farmland: FarmlandData, row: pd.Series) -> bool:
        """更新输入通量数据，返回是否有更新"""
        input_data = self.db.query(FluxCdInputData).filter(
            FluxCdInputData.farmland_id == farmland.farmland_id,
            FluxCdInputData.sample_id == farmland.sample_id
        ).first()

        if not input_data:
            self.logger.warning(f"未找到输入数据: Farmland_ID={farmland.farmland_id}, Sample_ID={farmland.sample_id}")
            return False

        # 检查是否需要更新
        updated = False
        if input_data.atmospheric_deposition != row['DQCJ_Cd']:
            input_data.atmospheric_deposition = row['DQCJ_Cd']
            updated = True
        if input_data.irrigation_input != row['GGS_Cd']:
            input_data.irrigation_input = row['GGS_Cd']
            updated = True
        if input_data.agro_chemicals_input != row['NCP_Cd']:
            input_data.agro_chemicals_input = row['NCP_Cd']
            updated = True

        if updated:
            self.logger.info(f"更新输入通量: Farmland_ID={farmland.farmland_id}, Sample_ID={farmland.sample_id}")

        return updated

    def _update_output_data(self, farmland: FarmlandData):
        """更新输出通量数据"""
        output_data = self.db.query(FluxCdOutputData).filter(
            FluxCdOutputData.farmland_id == farmland.farmland_id,
            FluxCdOutputData.sample_id == farmland.sample_id
        ).first()

        if not output_data:
            self.logger.warning(f"未找到输出数据: Farmland_ID={farmland.farmland_id}, Sample_ID={farmland.sample_id}")
            return

        # 重新获取更新后的输入数据
        input_data = self.db.query(FluxCdInputData).filter(
            FluxCdInputData.farmland_id == farmland.farmland_id,
            FluxCdInputData.sample_id == farmland.sample_id
        ).first()

        # 计算输出总通量
        out_cd = (input_data.groundwater_leaching +  # DX_Cd
                  input_data.surface_runoff +  # DB_Cd
                  input_data.grain_removal +  # ZL_Cd
                  input_data.straw_removal)  # JG_Cd

        # 更新输出通量记录
        output_data.in_cd = input_data.input_flux()
        output_data.out_cd = out_cd
        output_data.net_cd = output_data.in_cd - out_cd

        # 新增：计算并更新当年Cd浓度
        try:
            output_data.calculate_end_cd(self.db)
            self.logger.info(f"成功更新当年Cd浓度: {output_data.end_cd} mg/kg")
        except Exception as e:
            self.logger.error(f"更新当年Cd浓度失败: {str(e)}")

        self.logger.info(f"更新输出通量: Farmland_ID={farmland.farmland_id}, Sample_ID={farmland.sample_id}")


# 测试主函数
# 测试主函数
if __name__ == "__main__":
    import tempfile
    import shutil
    import json

    # 创建数据库会话
    db = SessionLocal()

    try:
        # 初始化服务
        service = FluxCdVisualizationService(db=db)

        print("=" * 50)
        print("测试Cd通量可视化服务")
        print("=" * 50)

        # 测试生成可视化结果
        print("\n>>> 测试生成Cd输入通量可视化地图")
        input_result = service.generate_cd_input_flux_map()
        print(json.dumps(input_result, indent=2, ensure_ascii=False))

        print("\n>>> 测试生成Cd输出通量可视化地图")
        output_result = service.generate_cd_output_flux_map()
        print(json.dumps(output_result, indent=2, ensure_ascii=False))

        print("\n>>> 测试生成Cd净通量可视化地图")
        net_result = service.generate_cd_net_flux_map()
        print(json.dumps(net_result, indent=2, ensure_ascii=False))

        # 测试更新服务
        print("\n>>> 测试从CSV更新Cd通量数据")

        # 创建临时目录和CSV文件
        temp_dir = tempfile.mkdtemp()
        test_csv_path = os.path.join(temp_dir, "test_update.csv")

        # 创建测试CSV文件
        test_data = [
            "lon,lan,DQCJ_Cd,GGS_Cd,NCP_Cd,DX_Cd,DB_Cd,ZL_Cd,JG_Cd",
            "113.123,25.456,1.24,4.56,7.89,0.1,0.2,0.05,0.03",
            "113.125,25.457,2.35,5.67,8.90,0.15,0.25,0.06,0.04"
        ]

        with open(test_csv_path, 'w', encoding='utf-8') as f:
            f.write("\n".join(test_data))

        # 调用更新服务
        update_result = service.update_from_csv(test_csv_path)
        print(json.dumps(update_result, indent=2, ensure_ascii=False))

        # 清理临时文件
        shutil.rmtree(temp_dir)
        print(f"\n清理临时目录: {temp_dir}")

        print("\n测试完成!")

    except Exception as e:
        logger.error(f"测试过程中发生错误: {str(e)}")
        import traceback

        traceback.print_exc()
    finally:
        # 关闭数据库会话
        db.close()