"""
热成像检测参数自动优化 — 基于标注数据调优 C++ 检测器参数

通过调用 thermal_detector.exe (--json 模式) 运行检测，
将结果与手动标注的 ground truth 比对，搜索最优参数。

由于 C++ 参数在头文件中是编译时常量，本工具的优化策略是：
  1. 运行默认参数检测，计算 baseline F1
  2. 生成候选参数值，逐个写入临时配置文件
  3. 输出最优参数供用户手动更新到 thermal_types.h

用法:
  python auto_tune.py                              # 使用 labels.json + test/test_images/
  python auto_tune.py --label-file my_labels.json  # 指定标注文件
  python auto_tune.py --image-dir path/to/images   # 指定图片目录
  python auto_tune.py --metric precision           # 优化目标改为精确率
  python auto_tune.py --export best_params.json    # 导出最优参数
"""

import cv2
import numpy as np
import json
import os
import sys
import glob
import subprocess
import re
import copy

# 标注类型 → 检测器输出类型映射
LABEL_TO_DETECT = {
    'Person': 'Person',
    'Barrel_Overheat': 'Barrel Overheat',
    'Cold_Weapon': 'Cold Weapon',
    'Flash_Residue': 'Flash Residue',
    'Hotspot': 'Hotspot',
}


def load_labels(label_path):
    with open(label_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def iou(box_a, box_b):
    """计算两个 bbox 的 IoU。box = [x, y, w, h]"""
    ax1, ay1 = box_a[0], box_a[1]
    ax2, ay2 = ax1 + box_a[2], ay1 + box_a[3]
    bx1, by1 = box_b[0], box_b[1]
    bx2, by2 = bx1 + box_b[2], by1 + box_b[3]

    ix1 = max(ax1, bx1)
    iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2)
    iy2 = min(ay2, by2)

    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    area_a = box_a[2] * box_a[3]
    area_b = box_b[2] * box_b[3]
    union = area_a + area_b - inter
    return inter / max(union, 1)


def match_detections(gt_labels, detections, iou_threshold=0.1):
    """匹配 ground truth 和检测结果，返回 TP, FP, FN

    使用混合匹配策略：先尝试 IoU，若 IoU < threshold 则用中心点距离匹配
    """
    gt_by_type = {}
    for lb in gt_labels:
        t = LABEL_TO_DETECT.get(lb['type'], lb['type'])
        gt_by_type.setdefault(t, []).append(lb)

    det_by_type = {}
    for d in detections:
        t = d.get('type', '')
        det_by_type.setdefault(t, []).append(d)

    tp, fp, fn = 0, 0, 0

    all_types = set(list(gt_by_type.keys()) + list(det_by_type.keys()))
    for t in all_types:
        gts = gt_by_type.get(t, [])
        dets = det_by_type.get(t, [])

        matched_gt = set()
        matched_det = set()

        pairs = []
        for gi, g in enumerate(gts):
            gx = g['bbox'][0] + g['bbox'][2] / 2
            gy = g['bbox'][1] + g['bbox'][3] / 2
            for di, d in enumerate(dets):
                db = d.get('bbox', {})
                dbox = [db.get('x', 0), db.get('y', 0), db.get('w', 0), db.get('h', 0)]
                dx = dbox[0] + dbox[2] / 2
                dy = dbox[1] + dbox[3] / 2

                # IoU 匹配
                iou_score = iou(g['bbox'], dbox)

                # 中心点距离匹配（当 IoU 低时）
                gt_size = max(g['bbox'][2], g['bbox'][3])
                det_size = max(dbox[2], dbox[3])
                max_size = max(gt_size, det_size, 1)
                dist = ((gx - dx) ** 2 + (gy - dy) ** 2) ** 0.5
                dist_ratio = dist / max_size

                # 综合分数：IoU 优先，中心点距离兜底
                if iou_score >= iou_threshold:
                    score = iou_score
                elif dist_ratio < 1.5:
                    score = iou_threshold * 0.5  # 给一个较低的匹配分数
                else:
                    continue

                pairs.append((score, gi, di))

        pairs.sort(reverse=True)
        for score, gi, di in pairs:
            if gi not in matched_gt and di not in matched_det:
                matched_gt.add(gi)
                matched_det.add(di)

        tp += len(matched_gt)
        fp += len(dets) - len(matched_det)
        fn += len(gts) - len(matched_gt)

    return tp, fp, fn


def _get_env_with_opencv():
    """构建包含 OpenCV DLL 路径的环境变量"""
    import copy
    env = copy.deepcopy(os.environ)
    # 常见 OpenCV 安装位置
    opencv_paths = [
        os.path.join(os.path.expanduser('~'), '.conda', 'envs', 'opencv-build', 'Library', 'bin'),
        os.path.join(os.path.expanduser('~'), 'miniconda3', 'envs', 'opencv-build', 'Library', 'bin'),
        os.path.join(os.path.expanduser('~'), 'anaconda3', 'envs', 'opencv-build', 'Library', 'bin'),
        os.path.join(os.path.expanduser('~'), 'opencv', 'build', 'x64', 'vc16', 'bin'),
        os.path.join(os.path.expanduser('~'), 'opencv', 'build', 'x64', 'vc17', 'bin'),
    ]
    extra = ';'.join(p for p in opencv_paths if os.path.isdir(p))
    if extra:
        env['PATH'] = extra + ';' + env.get('PATH', '')
    return env


def run_detector(exe_path, image_path):
    """调用 C++ 检测器，返回 JSON 结果"""
    try:
        env = _get_env_with_opencv()
        result = subprocess.run(
            [exe_path, image_path, '--json', '--no-show'],
            capture_output=True, text=True, timeout=30,
            encoding='utf-8', errors='replace',
            env=env,
        )
        output = result.stdout.strip()

        # 找到 JSON 块
        start = output.find('{')
        end = output.rfind('}') + 1
        if start >= 0 and end > start:
            json_str = output[start:end]
            # 移除 file 字段（中文路径含反斜杠会导致 JSON 解析失败）
            import re
            json_str = re.sub(r'"file"\s*:\s*"[^"]*"\s*,?', '', json_str)
            # 修复尾部逗号
            json_str = re.sub(r',\s*}', '}', json_str)
            json_str = re.sub(r',\s*]', ']', json_str)
            return json.loads(json_str)
    except Exception as e:
        print(f"  检测器运行失败: {e}")
    return None


def evaluate_default(exe_path, image_paths, all_labels):
    """用默认参数评估所有图片"""
    total_tp, total_fp, total_fn = 0, 0, 0

    for path in image_paths:
        fname = os.path.basename(path)
        if fname not in all_labels:
            continue

        result = run_detector(exe_path, path)
        if result is None:
            continue

        detections = result.get('detections', [])
        tp, fp, fn = match_detections(all_labels[fname], detections)
        total_tp += tp
        total_fp += fp
        total_fn += fn

    precision = total_tp / max(total_tp + total_fp, 1)
    recall = total_tp / max(total_tp + total_fn, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-8)

    return {
        'f1': f1, 'precision': precision, 'recall': recall,
        'tp': total_tp, 'fp': total_fp, 'fn': total_fn,
    }


def evaluate_with_filter(exe_path, image_paths, all_labels, filters):
    """用默认参数运行检测器，但对结果做后处理过滤来模拟参数变化

    可以模拟的参数：
      - confidence_threshold: 过滤低置信度检测
      - hotspot_max_count: 限制热点数量
      - 按类型过滤面积范围
    """
    total_tp, total_fp, total_fn = 0, 0, 0

    conf_thresh = filters.get('confidence_threshold', 0.0)
    hotspot_max = filters.get('hotspot_max_count', 999)
    person_min_area = filters.get('person_min_area', 0)

    for path in image_paths:
        fname = os.path.basename(path)
        if fname not in all_labels:
            continue

        result = run_detector(exe_path, path)
        if result is None:
            continue

        detections = result.get('detections', [])

        # 后处理过滤
        filtered = []
        hotspot_count = 0
        for d in detections:
            # 置信度过滤
            if d.get('confidence', 1.0) < conf_thresh:
                continue
            # 面积过滤 (Person)
            bbox = d.get('bbox', {})
            area = bbox.get('w', 0) * bbox.get('h', 0)
            if d.get('type') == 'Person' and area < person_min_area:
                continue
            # 热点数量限制
            if d.get('type') == 'Hotspot':
                hotspot_count += 1
                if hotspot_count > hotspot_max:
                    continue
            filtered.append(d)

        tp, fp, fn = match_detections(all_labels[fname], filtered)
        total_tp += tp
        total_fp += fp
        total_fn += fn

    precision = total_tp / max(total_tp + total_fp, 1)
    recall = total_tp / max(total_tp + total_fn, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-8)

    return {
        'f1': f1, 'precision': precision, 'recall': recall,
        'tp': total_tp, 'fp': total_fp, 'fn': total_fn,
    }


def search_best_params(exe_path, image_paths, all_labels, metric='f1'):
    """搜索最优参数组合"""
    print(f"\n{'='*60}")
    print(f"参数优化 — 调用 C++ 检测器")
    print(f"目标: {metric}")
    print(f"标注图片: {len(all_labels)} 张, 标注数: {sum(len(v) for v in all_labels.values())}")
    print(f"检测器: {exe_path}")
    print(f"{'='*60}\n")

    # 基准评估
    print("基准参数评估...")
    baseline = evaluate_default(exe_path, image_paths, all_labels)
    print(f"  基准: F1={baseline['f1']:.3f}  P={baseline['precision']:.3f}  R={baseline['recall']:.3f}")
    print(f"         TP={baseline['tp']}  FP={baseline['fp']}  FN={baseline['fn']}\n")

    # 可搜索参数 (通过后处理模拟)
    search_grid = {
        'confidence_threshold': [0.10, 0.20, 0.25, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80],
        'hotspot_max_count':    [3, 5, 8, 10, 15, 20],
        'person_min_area':      [100, 200, 300, 400, 500, 600, 800],
    }

    best_params = {}
    best_score = baseline[metric]
    improved = True
    iteration = 0

    while improved:
        improved = False
        iteration += 1
        print(f"--- 第 {iteration} 轮搜索 ---")

        for param_name, values in search_grid.items():
            current_val = best_params.get(param_name, values[len(values)//2])
            for val in values:
                test_params = dict(best_params)
                test_params[param_name] = val

                result = evaluate_with_filter(
                    exe_path, image_paths, all_labels, test_params
                )
                score = result[metric]

                if score > best_score + 0.001:
                    old_val = best_params.get(param_name, '?')
                    best_params[param_name] = val
                    best_score = score
                    improved = True
                    print(f"  {param_name}: {old_val} -> {val}  "
                          f"{metric}={score:.3f} "
                          f"(P={result['precision']:.3f} R={result['recall']:.3f} "
                          f"TP={result['tp']} FP={result['fp']} FN={result['fn']})")

        if not improved:
            print("  未找到更优参数，搜索结束")

    # 最终结果
    print(f"\n{'='*60}")
    print(f"优化结果 — {metric}={best_score:.3f}")
    print(f"{'='*60}")

    if best_params:
        print("\n需要修改的参数 (在 cpp/include/thermal_types.h 中):")
        param_map = {
            'confidence_threshold': ('confidence_threshold', '通用置信度下限'),
            'hotspot_max_count':    ('hotspot_max_count', '最多保留热点数'),
            'person_min_area':      ('person_min_area', '人员最小面积'),
        }
        for k, v in best_params.items():
            info = param_map.get(k, (k, ''))
            cxx_name = info[0]
            desc = info[1]
            print(f"  {cxx_name} = {v}  // {desc}")

            if cxx_name == 'confidence_threshold':
                print(f"    -> 在 thermal_detector.cpp 的 detect() 中修改阈值判断")
            elif cxx_name == 'person_min_area':
                print(f"    -> 在 DetectorConfig::person_min_area 中修改")
            elif cxx_name == 'hotspot_max_count':
                print(f"    -> 在 DetectorConfig::hotspot_max_count 中修改")
    else:
        print("\n默认参数已经是最优的，无需修改。")

    final = evaluate_with_filter(exe_path, image_paths, all_labels, best_params) if best_params else baseline
    print(f"\n最终指标: F1={final['f1']:.3f}  P={final['precision']:.3f}  R={final['recall']:.3f}")
    print(f"          TP={final['tp']}  FP={final['fp']}  FN={final['fn']}")

    # 对比
    if best_params:
        print(f"\n对比:")
        print(f"  优化前: F1={baseline['f1']:.3f}  TP={baseline['tp']}  FP={baseline['fp']}  FN={baseline['fn']}")
        print(f"  优化后: F1={final['f1']:.3f}  TP={final['tp']}  FP={final['fp']}  FN={final['fn']}")

    return best_params, final


def export_params(params, metrics, output_path):
    data = {
        'params': params,
        'metrics': metrics,
        'note': '由 auto_tune.py 自动优化生成，需手动更新到 cpp/include/thermal_types.h',
    }
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"\n参数已导出: {output_path}")


def find_exe():
    """查找 thermal_detector.exe"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    # 同级 cpp 目录
    exe = os.path.join(script_dir, 'cpp', 'thermal_detector.exe')
    if os.path.isfile(exe):
        return exe
    # build 目录
    exe = os.path.join(script_dir, 'cpp', 'build', 'Release', 'thermal_detector.exe')
    if os.path.isfile(exe):
        return exe
    return None


def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))

    args = sys.argv[1:]
    label_file = None
    image_dir = os.path.join(script_dir, 'test', 'test_images')
    metric = 'f1'
    export_path = None
    exe_path = None

    i = 0
    while i < len(args):
        if args[i] == '--label-file' and i + 1 < len(args):
            label_file = args[i + 1]
            i += 2
        elif args[i] == '--image-dir' and i + 1 < len(args):
            image_dir = args[i + 1]
            i += 2
        elif args[i] == '--metric' and i + 1 < len(args):
            metric = args[i + 1]
            i += 2
        elif args[i] == '--export' and i + 1 < len(args):
            export_path = args[i + 1]
            i += 2
        elif args[i] == '--exe' and i + 1 < len(args):
            exe_path = args[i + 1]
            i += 2
        else:
            i += 1

    # 检测器
    if exe_path is None:
        exe_path = find_exe()
    if exe_path is None:
        print("找不到 thermal_detector.exe")
        print("用法: python auto_tune.py --exe path/to/thermal_detector.exe")
        return
    print(f"检测器: {exe_path}")

    # 标注
    if label_file is None:
        label_file = os.path.join(script_dir, 'test', 'labels.json')
    if not os.path.exists(label_file):
        label_file_alt = os.path.join(script_dir, 'labels.json')
        if os.path.exists(label_file_alt):
            label_file = label_file_alt
        else:
            print(f"标注文件不存在: {label_file}")
            print("请先运行 label_tool.py 创建标注:")
            print("  python label_tool.py")
            return

    all_labels = load_labels(label_file)
    # 过滤掉 result_ 开头的图片（不对已标注的输出图做评估）
    all_labels = {k: v for k, v in all_labels.items() if not os.path.basename(k).startswith('result_')}
    print(f"标注文件: {label_file}")
    print(f"标注图片数: {len(all_labels)} (已过滤 result_ 图像)")

    if not all_labels:
        print("标注为空，请先标注图片")
        return

    # 图片
    image_paths = []
    for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif']:
        image_paths.extend(glob.glob(os.path.join(image_dir, ext)))
    image_paths.sort()
    labeled_images = [p for p in image_paths if os.path.basename(p) in all_labels]
    print(f"图片目录: {image_dir}")
    print(f"有标注的图片: {len(labeled_images)} / {len(image_paths)}")

    if not labeled_images:
        print("没有匹配到有标注的图片")
        return

    # 搜索
    best_params, metrics = search_best_params(exe_path, labeled_images, all_labels, metric)

    # 导出
    if export_path:
        export_params(best_params, metrics, export_path)
    else:
        default_export = os.path.join(script_dir, 'best_params.json')
        export_params(best_params, metrics, default_export)


if __name__ == '__main__':
    main()
