"""
批量处理脚本

批量分析图像并生成CSV报告
"""

import cv2
import pandas as pd
from pathlib import Path
from tqdm import tqdm
import logging
import argparse
from datetime import datetime
import json

try:
    from .integrated_analyzer import IntegratedMushroomAnalyzer
except ImportError:  # pragma: no cover - fallback for direct script usage
    from integrated_analyzer import IntegratedMushroomAnalyzer


class BatchProcessor:
    """批量处理器"""

    def __init__(
        self,
        yolo_model_path: str,
        sam_model_path: str,
        apriltag_size_mm: float = 37.58,
        device: str = 'cuda'
    ):
        """
        初始化批量处理器

        参数:
            yolo_model_path: YOLO模型路径
            sam_model_path: SAM模型路径
            apriltag_size_mm: AprilTag尺寸
            device: 运行设备
        """
        self.analyzer = IntegratedMushroomAnalyzer(
            yolo_model_path=yolo_model_path,
            sam_model_path=sam_model_path,
            apriltag_size_mm=apriltag_size_mm,
            device=device
        )

        logging.info("批量处理器初始化完成")

    def process_directory(
        self,
        input_dir: str,
        output_dir: str,
        image_extensions: list = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'],
        save_visualizations: bool = True
    ):
        """
        批量处理目录中的所有图像

        参数:
            input_dir: 输入图像目录
            output_dir: 输出目录
            image_extensions: 支持的图像格式
            save_visualizations: 是否保存可视化图像
        """
        input_path = Path(input_dir)
        output_path = Path(output_dir)

        # 创建输出目录
        output_path.mkdir(parents=True, exist_ok=True)
        vis_dir = output_path / 'visualizations'
        if save_visualizations:
            vis_dir.mkdir(exist_ok=True)

        # 查找所有图像文件
        image_files = []
        for ext in image_extensions:
            image_files.extend(list(input_path.glob(f'*{ext}')))

        if len(image_files) == 0:
            logging.error(f"在 {input_dir} 中未找到图像文件")
            return

        logging.info(f"找到 {len(image_files)} 张图像")

        # 存储结果
        all_results = []
        failed_images = []

        # 逐张处理
        for image_file in tqdm(image_files, desc="处理进度"):
            try:
                result = self.analyzer.analyze_image(
                    str(image_file),
                    visualize=save_visualizations
                )

                if result and result['success']:
                    # 展平结果（每个香菇一行）
                    for mushroom in result['mushrooms']:
                        if 'error' not in mushroom:
                            row = {
                                'filename': result['filename'],
                                'apriltag_detected': result['apriltag_detected'],
                                'px_per_mm': result['px_per_mm'],
                                'mushroom_id': mushroom['mushroom_id'],
                                'detection_confidence': mushroom['detection_confidence'],
                                'diameter_px': mushroom['diameter_px'],
                                'diameter_mm': mushroom['diameter_mm'],
                                'cap_area_px': mushroom['cap_area_px'],
                                'cap_area_mm2': mushroom['cap_area_mm2'],
                                'crack_area_px': mushroom['crack_area_px'],
                                'crack_ratio': mushroom['crack_ratio'],
                                'crack_ratio_percent': mushroom['crack_ratio_percent'],
                                'pattern_type': mushroom['pattern_type'],
                                'pattern_category': mushroom['pattern_category'],
                                'white_pattern_ratio': mushroom['white_pattern_ratio'],
                                'brown_pattern_ratio': mushroom['brown_pattern_ratio'],
                                'flatness_score': mushroom.get('flatness_score', 0.0),
                                'suspected_ban_gu': mushroom.get('suspected_ban_gu', False),
                                'v_threshold': mushroom['v_threshold'],
                                's_threshold': mushroom['s_threshold']
                            }
                            all_results.append(row)

                    # 保存可视化
                    if save_visualizations and 'visualization' in result:
                        vis_path = vis_dir / f"{image_file.stem}_analysis.jpg"
                        cv2.imwrite(str(vis_path), result['visualization'])

                else:
                    error_msg = result.get('error', '未知错误') if result else '分析失败'
                    failed_images.append({
                        'filename': image_file.name,
                        'error': error_msg
                    })
                    logging.warning(f"{image_file.name} 处理失败: {error_msg}")

            except Exception as e:
                logging.error(f"处理 {image_file.name} 时出错: {str(e)}")
                failed_images.append({
                    'filename': image_file.name,
                    'error': str(e)
                })

        # 保存结果到CSV
        if len(all_results) > 0:
            df = pd.DataFrame(all_results)

            # 保存完整结果
            csv_path = output_path / 'analysis_results.csv'
            df.to_csv(csv_path, index=False, encoding='utf-8-sig')
            logging.info(f"结果已保存到：{csv_path}")

            # 生成统计报告
            self._generate_report(df, output_path)

        # 保存失败记录
        if len(failed_images) > 0:
            failed_df = pd.DataFrame(failed_images)
            failed_csv = output_path / 'failed_images.csv'
            failed_df.to_csv(failed_csv, index=False, encoding='utf-8-sig')
            logging.warning(f"{len(failed_images)} 张图像处理失败，详见：{failed_csv}")

        # 生成处理日志
        log_data = {
            'timestamp': datetime.now().isoformat(),
            'total_images': len(image_files),
            'successful_images': len(image_files) - len(failed_images),
            'failed_images': len(failed_images),
            'total_mushrooms': len(all_results),
            'input_dir': str(input_path),
            'output_dir': str(output_path)
        }

        log_path = output_path / 'processing_log.json'
        with open(log_path, 'w', encoding='utf-8') as f:
            json.dump(log_data, f, indent=2, ensure_ascii=False)

        print("\n" + "="*60)
        print("批量处理完成")
        print("="*60)
        print(f"总图像数: {log_data['total_images']}")
        print(f"成功处理: {log_data['successful_images']}")
        print(f"失败数量: {log_data['failed_images']}")
        print(f"检测到香菇: {log_data['total_mushrooms']} 个")
        print(f"结果保存在: {output_path}")
        print("="*60)

    def _generate_report(self, df: pd.DataFrame, output_dir: Path):
        """
        生成统计报告

        参数:
            df: 结果DataFrame
            output_dir: 输出目录
        """
        report_lines = []
        report_lines.append("="*60)
        report_lines.append("香菇分析统计报告")
        report_lines.append("="*60)
        report_lines.append(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        report_lines.append(f"总样本数: {len(df)}")
        report_lines.append("")

        # 直径统计
        report_lines.append("【直径统计】")
        report_lines.append(f"  平均值: {df['diameter_mm'].mean():.2f} mm")
        report_lines.append(f"  中位数: {df['diameter_mm'].median():.2f} mm")
        report_lines.append(f"  标准差: {df['diameter_mm'].std():.2f} mm")
        report_lines.append(f"  最小值: {df['diameter_mm'].min():.2f} mm")
        report_lines.append(f"  最大值: {df['diameter_mm'].max():.2f} mm")
        report_lines.append("")

        # 面积统计
        report_lines.append("【面积统计】")
        report_lines.append(f"  平均值: {df['cap_area_mm2'].mean():.2f} mm²")
        report_lines.append(f"  中位数: {df['cap_area_mm2'].median():.2f} mm²")
        report_lines.append(f"  标准差: {df['cap_area_mm2'].std():.2f} mm²")
        report_lines.append(f"  最小值: {df['cap_area_mm2'].min():.2f} mm²")
        report_lines.append(f"  最大值: {df['cap_area_mm2'].max():.2f} mm²")
        report_lines.append("")

        # 开裂比统计
        report_lines.append("【开裂比统计】")
        report_lines.append(f"  平均值: {df['crack_ratio_percent'].mean():.2f}%")
        report_lines.append(f"  中位数: {df['crack_ratio_percent'].median():.2f}%")
        report_lines.append(f"  标准差: {df['crack_ratio_percent'].std():.2f}%")
        report_lines.append(f"  最小值: {df['crack_ratio_percent'].min():.2f}%")
        report_lines.append(f"  最大值: {df['crack_ratio_percent'].max():.2f}%")
        report_lines.append("")

        # 花纹类型分布
        report_lines.append("【花纹类型分布】")
        pattern_counts = df['pattern_type'].value_counts().sort_index()
        for pattern, count in pattern_counts.items():
            percentage = (count / len(df)) * 100
            report_lines.append(f"  {pattern}: {count} 个 ({percentage:.1f}%)")
        report_lines.append("")

        # 标定参数统计
        report_lines.append("【标定参数】")
        report_lines.append(f"  平均像素/毫米: {df['px_per_mm'].mean():.3f}")
        report_lines.append(f"  标定参数标准差: {df['px_per_mm'].std():.3f}")
        report_lines.append("")

        report_lines.append("="*60)

        # 保存报告
        report_path = output_dir / 'analysis_report.txt'
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(report_lines))

        # 同时打印到控制台
        print('\n'.join(report_lines))

        logging.info(f"统计报告已保存到：{report_path}")


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(description='香菇图像批量分析工具')

    parser.add_argument(
        'input_dir',
        type=str,
        help='输入图像目录'
    )

    parser.add_argument(
        '--output-dir',
        type=str,
        default='./outputs',
        help='输出目录（默认：./outputs）'
    )

    parser.add_argument(
        '--yolo-model',
        type=str,
        default='../models/yolo_mushroom.pt',
        help='YOLO模型路径'
    )

    parser.add_argument(
        '--sam-model',
        type=str,
        default='../models/sam_vit_b.pth',
        help='SAM模型路径'
    )

    parser.add_argument(
        '--apriltag-size',
        type=float,
        default=37.58,
        help='AprilTag尺寸（毫米，默认：37.58）'
    )

    parser.add_argument(
        '--device',
        type=str,
        default='cuda',
        choices=['cuda', 'cpu'],
        help='运行设备（默认：cuda）'
    )

    parser.add_argument(
        '--no-visualizations',
        action='store_true',
        help='不保存可视化图像'
    )

    args = parser.parse_args()

    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )

    # 创建处理器
    processor = BatchProcessor(
        yolo_model_path=args.yolo_model,
        sam_model_path=args.sam_model,
        apriltag_size_mm=args.apriltag_size,
        device=args.device
    )

    # 批量处理
    processor.process_directory(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        save_visualizations=not args.no_visualizations
    )


if __name__ == "__main__":
    main()
