#!/usr/bin/env python3
"""
用多尺度模板匹配重新分析已拍摄的 PTZ 校准图像，生成更可靠的映射表。
"""
import os
import sys
import json
import argparse
from pathlib import Path
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import cv2
import numpy as np


def multi_scale_template_match(ptz_img: np.ndarray, panorama_img: np.ndarray,
                                scales=(0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5)):
    """
    在全景图中搜索 PTZ 图像的最佳匹配位置。
    返回 (x_ratio, y_ratio, best_scale, score)
    """
    if ptz_img is None or panorama_img is None:
        return None

    ptz_gray = cv2.cvtColor(ptz_img, cv2.COLOR_BGR2GRAY)
    pano_gray = cv2.cvtColor(panorama_img, cv2.COLOR_BGR2GRAY)
    ph, pw = ptz_gray.shape
    pano_h, pano_w = pano_gray.shape

    best = None
    best_score = -1

    for scale in scales:
        resized_w = int(pw * scale)
        resized_h = int(ph * scale)
        if resized_w > pano_w or resized_h > pano_h:
            continue

        resized = cv2.resize(ptz_gray, (resized_w, resized_h), interpolation=cv2.INTER_AREA)
        result = cv2.matchTemplate(pano_gray, resized, cv2.TM_CCOEFF_NORMED)
        _, max_val, _, max_loc = cv2.minMaxLoc(result)

        if max_val > best_score:
            best_score = max_val
            center_x = max_loc[0] + resized_w / 2
            center_y = max_loc[1] + resized_h / 2
            best = (center_x / pano_w, center_y / pano_h, scale, float(max_val))

    return best


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input-dir', type=str, required=True, help='扫描结果目录')
    parser.add_argument('--output-dir', type=str, default=None, help='输出目录')
    parser.add_argument('--score-threshold', type=float, default=0.45, help='匹配分数阈值')
    args = parser.parse_args()

    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir) if args.output_dir else input_dir / 'reanalysis'
    output_dir.mkdir(parents=True, exist_ok=True)

    ptz_dir = input_dir / 'ptz_images'
    panorama_path = input_dir / 'panorama.jpg'

    panorama = cv2.imread(str(panorama_path))
    if panorama is None:
        print(f'无法读取全景图: {panorama_path}')
        return

    pano_h, pano_w = panorama.shape[:2]
    print(f'全景图尺寸: {pano_w}x{pano_h}')

    raw_path = input_dir / 'mapping_raw.json'
    with open(raw_path, 'r', encoding='utf-8') as f:
        raw_data = json.load(f)

    records = []
    for r in raw_data['records']:
        filename = r['filename']
        ptz_path = ptz_dir / filename
        ptz_img = cv2.imread(str(ptz_path))
        if ptz_img is None:
            print(f'  无法读取: {ptz_path}')
            continue

        result = multi_scale_template_match(ptz_img, panorama)
        record = dict(r)
        if result and result[3] >= args.score_threshold:
            x_ratio, y_ratio, scale, score = result
            record['tm_x_ratio'] = round(x_ratio, 4)
            record['tm_y_ratio'] = round(y_ratio, 4)
            record['tm_panorama_x'] = int(x_ratio * pano_w)
            record['tm_panorama_y'] = int(y_ratio * pano_h)
            record['tm_scale'] = round(scale, 3)
            record['tm_score'] = round(score, 3)
            print(f"{filename}: pan={r['pan']:3d} tilt={r['tilt']:+3d} -> x={x_ratio:.3f} y={y_ratio:.3f} scale={scale:.2f} score={score:.3f}")
        else:
            if result:
                print(f"{filename}: pan={r['pan']:3d} tilt={r['tilt']:+3d} -> score too low: {result[3]:.3f}")
            else:
                print(f"{filename}: pan={r['pan']:3d} tilt={r['tilt']:+3d} -> no match")

        records.append(record)

    # 保存结果
    output_path = output_dir / 'mapping_tm.json'
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump({
            'created_at': datetime.now().isoformat(),
            'panorama_size': {'width': pano_w, 'height': pano_h},
            'score_threshold': args.score_threshold,
            'records': records,
        }, f, indent=2, ensure_ascii=False)
    print(f'\n结果已保存: {output_path}')

    # 生成查找表
    valid = [r for r in records if 'tm_x_ratio' in r]
    pan_lookup = sorted([[r['tm_x_ratio'], float(r['pan'])] for r in valid], key=lambda x: x[0])
    tilt_lookup = sorted([[r['tm_y_ratio'], float(r['tilt'])] for r in valid], key=lambda x: x[0])
    lookup = {
        'created_at': datetime.now().isoformat(),
        'pan_lookup': pan_lookup,
        'tilt_lookup': tilt_lookup,
        'valid_count': len(valid),
    }
    lookup_path = output_dir / 'lookup_table_tm.json'
    with open(lookup_path, 'w', encoding='utf-8') as f:
        json.dump(lookup, f, indent=2, ensure_ascii=False)
    print(f'查找表已保存: {lookup_path} (有效点 {len(valid)}/{len(records)})')


if __name__ == '__main__':
    main()