#!/usr/bin/env python3
"""
相机标定工具 - 为车间人脸识别系统生成ROI和三分区参数

Usage:
    python calibrate_camera.py --height 5.0 --pitch 45 --report
    python calibrate_camera.py --height 4.5 --pitch 40 --focal-estimate 2000 -o cam.json

Output:
    生成可直接复制到配置文件的params参数
"""

import json
import argparse
import numpy as np
from dataclasses import dataclass
from typing import List, Tuple, Dict
import sys


@dataclass
class CameraParams:
    """相机参数"""
    height: float      # 安装高度(m)
    pitch_deg: float   # 俯仰角(°)
    focal_px: float    # 像素焦距(px)
    img_w: int         # 图像宽度
    img_h: int         # 图像高度


class CameraCalibrator:
    """相机标定器 - 计算ROI和三分区参数"""
    
    def __init__(self, p: CameraParams):
        self.p = p
        self.cy = p.img_h // 2
        self.cx = p.img_w // 2
        self.theta = np.radians(p.pitch_deg)
        self._build_lut()
    
    def _build_lut(self):
        """预计算距离查找表 LUT[y] = distance(m)"""
        self.lut = np.zeros(self.p.img_h, dtype=np.float32)
        for y in range(self.p.img_h):
            dy = y - self.cy
            angle = self.theta + np.arctan2(dy, self.p.focal_px)
            if abs(angle) > 1e-6 and np.tan(angle) > 0:
                self.lut[y] = self.p.height / np.tan(angle)
            else:
                self.lut[y] = np.inf
    
    def pixel_to_distance(self, y: int) -> float:
        """像素y坐标 -> 距离(m)"""
        if 0 <= y < self.p.img_h:
            return float(self.lut[y])
        return np.inf
    
    def distance_to_pixel(self, d: float) -> int:
        """距离(m) -> 像素y坐标"""
        if d <= 0:
            return self.cy
        angle = np.arctan2(self.p.height, d)
        offset = self.p.focal_px * np.tan(angle - self.theta)
        return int(np.clip(self.cy + offset, 0, self.p.img_h - 1))
    
    def calculate_roi(self, min_d: float = 3.0, max_d: float = 9.0, 
                      margin: int = 20) -> Dict:
        """计算ROI裁剪区域
        
        Args:
            min_d: 最近检测距离（米）
            max_d: 最远检测距离（米）
            margin: 边界余量（像素）
        """
        # 注意：距离越远，y坐标越大（画面下方）
        y_min = self.distance_to_pixel(max_d)  # 远距在下
        y_max = self.distance_to_pixel(min_d)  # 近距在上
        
        # 添加边界余量
        y_min = max(0, y_min - margin)
        y_max = min(self.p.img_h, y_max + margin)
        
        saving = 1 - (y_max - y_min) / self.p.img_h
        
        return {
            "crop": {
                "x": 0,
                "y": int(y_min),
                "w": self.p.img_w,
                "h": int(y_max - y_min)
            },
            "saving_percent": round(saving * 100, 1),
            "y_min": int(y_min),
            "y_max": int(y_max)
        }
    
    def get_zone_boundaries(self, boundaries_m: List[float]) -> List[int]:
        """获取分区边界像素坐标
        
        Args:
            boundaries_m: 分界线距离列表，如[5.0, 7.0]
        """
        return [self.distance_to_pixel(d) for d in boundaries_m]
    
    def estimate_face_size(self, distance: float, real_width: float = 0.16) -> float:
        """估算给定距离的人脸像素大小"""
        if distance <= 0:
            return 0
        return self.p.focal_px * real_width / distance
    
    def generate_config(self, min_d: float = 3.0, max_d: float = 9.0,
                       boundaries_m: List[float] = None) -> Dict:
        """生成配置文件
        
        Returns:
            包含params和meta的配置字典
        """
        if boundaries_m is None:
            boundaries_m = [5.0, 7.0]  # 默认5米和7米分界线
            
        roi = self.calculate_roi(min_d, max_d)
        boundaries_y = self.get_zone_boundaries(boundaries_m)
        
        # 生成可直接复制到instances params的配置
        config = {
            "params": {
                "camera_pitch": self.p.pitch_deg,
                "roi_enabled": True,
                "roi_y": roi["crop"]["y"],
                "roi_h": roi["crop"]["h"],
                "zones_enabled": True,
                "zone_boundary_5m": boundaries_y[0],
                "zone_boundary_7m": boundaries_y[1],
                # 320模型缩放因子
                "zone_scale_near": 1.0,
                "zone_scale_mid": 1.3,
                "zone_scale_far": 1.8
            },
            "meta": {
                "height": self.p.height,
                "pitch": self.p.pitch_deg,
                "focal_px": self.p.focal_px,
                "image_size": [self.p.img_w, self.p.img_h],
                "roi_saving": roi["saving_percent"],
                "zone_boundaries_m": boundaries_m,
                "zone_boundaries_y": boundaries_y,
                "detection_range": [min_d, max_d]
            }
        }
        
        return config
    
    def print_report(self, min_d: float = 3.0, max_d: float = 9.0,
                    boundaries_m: List[float] = None):
        """打印标定报告"""
        if boundaries_m is None:
            boundaries_m = [5.0, 7.0]
            
        roi = self.calculate_roi(min_d, max_d)
        boundaries_y = self.get_zone_boundaries(boundaries_m)
        
        print("=" * 70)
        print("相机标定报告")
        print("=" * 70)
        print(f"\n【相机参数】")
        print(f"  安装高度 H: {self.p.height}m")
        print(f"  俯仰角 θ: {self.p.pitch_deg}°")
        print(f"  像素焦距 f: {self.p.focal_px}px")
        print(f"  图像尺寸: {self.p.img_w}x{self.p.img_h}")
        
        print(f"\n【ROI配置】(检测范围 {min_d}-{max_d}m)")
        print(f"  裁剪区域: y={roi['crop']['y']}, h={roi['crop']['h']}")
        print(f"  算力节省: {roi['saving_percent']}%")
        
        print(f"\n【三分区配置】(RetinaFace_320)")
        print(f"  分界线y坐标: 5m={boundaries_y[0]}px, 7m={boundaries_y[1]}px")
        # y坐标：画面上方（y小）= 远距离，画面下方（y大）= 近距离
        by5, by7 = boundaries_y[0], boundaries_y[1]
        if by5 > by7:
            # 5米分界线在下方（y大），7米在上方（y小）
            print(f"  画面分布: 上方(y小)=远距离, 下方(y大)=近距离")
            print(f"  近区(3-5m): y>{by5}, scale=1.0x  (画面下方)")
            print(f"  中区(5-7m): {by7}<y≤{by5}, scale=1.3x  (画面中部)")
            print(f"  远区(7-9m): y≤{by7}, scale=1.8x  (画面上方)")
        else:
            print(f"  近区(3-5m): y<{by5}, scale=1.0x")
            print(f"  中区(5-7m): {by5}≤y<{by7}, scale=1.3x")
            print(f"  远区(7-9m): y≥{by7}, scale=1.8x")
        
        print(f"\n【距离-像素-人脸大小映射】")
        print(f"  {'距离':>6} | {'像素y':>6} | {'原始人脸':>10} | {'320输入':>9} | {'区域':>6}")
        print(f"  {'-'*60}")
        
        for d in [3, 4, 5, 6, 7, 8, 9]:
            y = self.distance_to_pixel(d)
            face_orig = self.estimate_face_size(d)
            
            # 确定区域和320输入大小
            zone_idx = 0
            if d >= boundaries_m[1]:
                zone_idx = 2
            elif d >= boundaries_m[0]:
                zone_idx = 1
            
            scales = [1.0, 1.3, 1.8]
            face_320 = face_orig * scales[zone_idx]
            zone_name = ["近", "中", "远"][zone_idx]
            
            print(f"  {d:>6.0f}m | {y:>6} | {face_orig:>9.0f}px | {face_320:>9.0f}px | {zone_name:>6}")
        
        print(f"\n【320模型说明】")
        print(f"  - 320最佳检测范围：25-35px（占320输入的8-11%）")
        print(f"  - 8-9米检测率可能降至70%，如不满足可升级到640模型")
        print(f"  - 升级只需改：model_path + scales=[0.7,1.0,1.4]")
        
        print("\n" + "=" * 70)


def main():
    parser = argparse.ArgumentParser(
        description='相机标定工具 - 生成ROI和三分区参数',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 基础标定（推荐）
  python calibrate_camera.py --height 5.0 --pitch 45 --report
  
  # 自定义参数并保存
  python calibrate_camera.py --height 4.5 --pitch 40 --focal-estimate 2000 -o cam.json --report
  
  # 调整分区边界
  python calibrate_camera.py --height 5.0 --zones 4.5 6.5 --report
        """
    )
    parser.add_argument('--height', type=float, required=True,
                       help='相机安装高度（米），如 5.0')
    parser.add_argument('--pitch', type=float, default=45,
                       help='俯仰角（度），默认45')
    parser.add_argument('--focal-estimate', type=float, default=2200,
                       help='像素焦距估算值，默认2200（2.5K相机4-6mm镜头）')
    parser.add_argument('--image-size', type=int, nargs=2,
                       default=[2560, 1440],
                       help='图像尺寸 宽 高，默认 2560 1440')
    parser.add_argument('--range', type=float, nargs=2,
                       default=[3.0, 9.0],
                       help='检测范围 最小距离 最大距离，默认 3.0 9.0')
    parser.add_argument('--zones', type=float, nargs=2,
                       default=[5.0, 7.0],
                       help='分区边界距离（米），默认 5.0 7.0（形成3-5,5-7,7-9三区）')
    parser.add_argument('-o', '--output', default='camera_calib.json',
                       help='输出文件路径，默认 camera_calib.json')
    parser.add_argument('--report', action='store_true',
                       help='打印详细报告')
    
    args = parser.parse_args()
    
    # 创建标定器
    params = CameraParams(
        height=args.height,
        pitch_deg=args.pitch,
        focal_px=args.focal_estimate,
        img_w=args.image_size[0],
        img_h=args.image_size[1]
    )
    
    calib = CameraCalibrator(params)
    config = calib.generate_config(
        args.range[0],
        args.range[1],
        list(args.zones)
    )
    
    # 保存配置
    with open(args.output, 'w', encoding='utf-8') as f:
        json.dump(config, f, indent=2, ensure_ascii=False)
    
    print(f"[OK] 配置已保存: {args.output}")
    
    # 打印报告
    if args.report:
        calib.print_report(args.range[0], args.range[1], list(args.zones))
        print(f"\n【可直接复制到配置文件的 params】")
        print(json.dumps(config["params"], indent=2, ensure_ascii=False))
    else:
        # 简洁输出关键参数
        print(f"\n关键参数:")
        print(f"  roi_y: {config['params']['roi_y']}")
        print(f"  roi_h: {config['params']['roi_h']}")
        print(f"  zone_boundary_5m: {config['params']['zone_boundary_5m']}")
        print(f"  zone_boundary_7m: {config['params']['zone_boundary_7m']}")


if __name__ == '__main__':
    main()