#!/usr/bin/env python3
"""
光湖视频AI系统 · OpenCV 平面追踪引擎
D140 · 铸渊 ICE-GL-ZY001 · 2026-06-23

用途:
  1. 特征点追踪 — 追踪画面中物体运动轨迹
  2. 画面稳定化 — 去除抖动
  3. 文字/图形贴合 — 把文字/水印/PIN到追踪的平面上
  4. 运动数据导出 — 供 video-editor.js 的 keyframes 使用

调用方式:
  python3 planar-tracker.py <mode> <video> [options]

  mode:
    track     — 追踪特征点，输出运动轨迹JSON
    stabilize — 画面稳定化，输出稳定后视频
    pin_text  — 追踪+贴合文字到指定区域
    motion_data — 导出运动数据供FFmpeg使用

输出:
  track     → JSON文件 (轨迹数据)
  stabilize → MP4文件 (稳定后视频)
  pin_text  → MP4文件 (带贴合文字的视频)
  motion_data → JSON文件 (zoom/pan关键帧数据)

铸渊 ICE-GL-ZY001 · D140
"""

import cv2
import numpy as np
import json
import sys
import os
import argparse
from pathlib import Path

# ==================== 特征点追踪 ====================

def track_features(video_path, output_json, max_points=200, min_distance=15, quality=0.01):
    """
    Shi-Tomasi角点检测 + Lucas-Kanade光流追踪
    
    输出JSON格式:
    {
      "fps": 24,
      "frame_count": 120,
      "resolution": [1920, 1080],
      "tracks": [
        {
          "point_id": 0,
          "points": [{"frame": 0, "x": 100, "y": 200}, ...],
          "start": [100, 200],
          "end": [105, 198],
          "displacement": [5, -2],
          "avg_velocity": [0.04, -0.017]
        }
      ]
    }
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"无法打开视频: {video_path}")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    print(f"[Tracker] 视频: {Path(video_path).name}  {width}x{height} @ {fps:.1f}fps  {frame_count}帧")
    
    # 读取第一帧
    ret, prev_frame = cap.read()
    if not ret:
        raise RuntimeError("无法读取视频帧")
    
    prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
    
    # Shi-Tomasi角点检测
    prev_points = cv2.goodFeaturesToTrack(
        prev_gray, maxCorners=max_points, qualityLevel=quality,
        minDistance=min_distance, blockSize=7
    )
    
    if prev_points is None:
        raise RuntimeError("第一帧未检测到特征点")
    
    print(f"[Tracker] 检测到 {len(prev_points)} 个特征点")
    
    # 初始化轨迹
    tracks = []
    for i, pt in enumerate(prev_points):
        tracks.append({
            "point_id": i,
            "points": [{"frame": 0, "x": float(pt[0][0]), "y": float(pt[0][1])}],
            "start": [float(pt[0][0]), float(pt[0][1])],
            "active": True
        })
    
    # Lucas-Kanade参数
    lk_params = dict(
        winSize=(15, 15),
        maxLevel=2,
        criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
    )
    
    frame_idx = 1
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # 计算光流
        next_points, status, error = cv2.calcOpticalFlowPyrLK(
            prev_gray, frame_gray, prev_points, None, **lk_params
        )
        
        # 更新轨迹
        active_count = 0
        for i, (track, st) in enumerate(zip(tracks, status)):
            if not track["active"]:
                continue
            if st == 1 and next_points is not None:
                x, y = float(next_points[i][0]), float(next_points[i][1])
                track["points"].append({"frame": frame_idx, "x": x, "y": y})
                track["active"] = True
                active_count += 1
            else:
                track["active"] = False
        
        # 只保留活跃点供下一帧追踪
        if active_count > 0:
            mask = status.flatten() == 1
            prev_points = next_points[mask].reshape(-1, 1, 2)
            # 对应更新tracks索引
            new_tracks = [t for t, m in zip(tracks, mask) if m]
            # 但要保留丢失的轨迹（标记为非活跃）
            lost_tracks = [t for t, m in zip(tracks, mask) if not m]
            tracks = new_tracks + lost_tracks
            # 重新编号并调整
            # 注意：prev_points和新tracks顺序需要对齐
            # 上面的操作保证了new_tracks和prev_points对齐
        
        prev_gray = frame_gray
        frame_idx += 1
        
        if frame_idx % 30 == 0:
            print(f"  [Tracker] 帧 {frame_idx}/{frame_count}  活跃点: {active_count}")
    
    cap.release()
    
    # 计算位移和速度
    for track in tracks:
        if len(track["points"]) >= 2:
            start = track["start"]
            end = track["points"][-1]
            track["end"] = [end["x"], end["y"]]
            track["displacement"] = [end["x"] - start[0], end["y"] - start[1]]
            frames = len(track["points"])
            track["avg_velocity"] = [
                track["displacement"][0] / frames,
                track["displacement"][1] / frames
            ]
    
    # 只保留有足够数据的轨迹
    valid_tracks = [t for t in tracks if len(t["points"]) >= frame_count * 0.3]
    print(f"[Tracker] 有效轨迹: {len(valid_tracks)}/{len(tracks)}")
    
    result = {
        "fps": fps,
        "frame_count": frame_idx - 1,
        "resolution": [width, height],
        "tracks": valid_tracks
    }
    
    with open(output_json, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    print(f"[Tracker] ✅ 轨迹数据已保存: {output_json}")
    return result


# ==================== 画面稳定化 ====================

def stabilize_video(video_path, output_path, smoothing_radius=25, crop_percent=0.08):
    """
    基于特征点轨迹的画面稳定化
    
    原理:
      1. 追踪相邻帧的特征点
      2. 计算相邻帧之间的仿射变换矩阵
      3. 累积变换轨迹
      4. 平滑轨迹（滑动平均）
      5. 应用平滑后的变换
    
    smoothing_radius: 平滑窗口大小（帧数），越大越平滑但延迟越大
    crop_percent: 边缘裁剪比例，补偿稳定后的黑边
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"无法打开视频: {video_path}")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    print(f"[Stabilize] {Path(video_path).name}  {width}x{height} @ {fps:.1f}fps")
    
    # 读取所有帧
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)
    cap.release()
    
    if len(frames) < 2:
        raise RuntimeError("视频帧数不足")
    
    # 追踪特征点
    prev_gray = cv2.cvtColor(frames[0], cv2.COLOR_BGR2GRAY)
    transforms = []  # 每帧的变换 [dx, dy, da(角度)]
    
    for i in range(1, len(frames)):
        curr_gray = cv2.cvtColor(frames[i], cv2.COLOR_BGR2GRAY)
        
        prev_pts = cv2.goodFeaturesToTrack(prev_gray, maxCorners=200, qualityLevel=0.01,
                                           minDistance=15, blockSize=7)
        if prev_pts is None or len(prev_pts) < 10:
            transforms.append([0, 0, 0])
            prev_gray = curr_gray
            continue
        
        curr_pts, status, _ = cv2.calcOpticalFlowPyrLK(prev_gray, curr_gray, prev_pts, None)
        
        # 筛选有效点
        valid_prev = prev_pts[status == 1]
        valid_curr = curr_pts[status == 1]
        
        if len(valid_prev) < 10:
            transforms.append([0, 0, 0])
            prev_gray = curr_gray
            continue
        
        # 计算仿射变换 (只取平移和旋转)
        m = cv2.estimateAffinePartial2D(valid_prev, valid_curr)[0]
        if m is None:
            transforms.append([0, 0, 0])
            prev_gray = curr_gray
            continue
        
        dx = m[0, 2]
        dy = m[1, 2]
        da = np.arctan2(m[1, 0], m[0, 0])
        transforms.append([dx, dy, da])
        prev_gray = curr_gray
        
        if i % 30 == 0:
            print(f"  [Stabilize] 帧 {i}/{len(frames)}")
    
    # 累积轨迹
    trajectory = np.zeros((len(transforms) + 1, 3))
    for i, t in enumerate(transforms):
        trajectory[i + 1] = trajectory[i] + np.array(t)
    
    # 平滑轨迹（滑动平均）
    smoothed = np.zeros_like(trajectory)
    for i in range(len(trajectory)):
        start = max(0, i - smoothing_radius)
        end = min(len(trajectory), i + smoothing_radius + 1)
        smoothed[i] = np.mean(trajectory[start:end], axis=0)
    
    # 计算平滑后的变换
    smooth_diff = smoothed - trajectory
    
    # 写出稳定后的视频
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    crop_x = int(width * crop_percent)
    crop_y = int(height * crop_percent)
    out_w = width - 2 * crop_x
    out_h = height - 2 * crop_y
    
    writer = cv2.VideoWriter(output_path, fourcc, fps, (out_w, out_h))
    
    for i, frame in enumerate(frames):
        dx = smooth_diff[i, 0]
        dy = smooth_diff[i, 1]
        da = smooth_diff[i, 2]
        
        # 构造仿射变换矩阵
        m = np.array([
            [np.cos(da), -np.sin(da), dx],
            [np.sin(da), np.cos(da), dy]
        ], dtype=np.float32)
        
        stabilized = cv2.warpAffine(frame, m, (width, height))
        # 裁剪黑边
        cropped = stabilized[crop_y:crop_y + out_h, crop_x:crop_x + out_w]
        writer.write(cropped)
    
    writer.release()
    print(f"[Stabilize] ✅ 稳定后视频: {output_path}")
    return output_path


# ==================== 文字贴合 ====================

def pin_text_to_video(video_path, output_path, text, 
                       start_frame=0, end_frame=None,
                       init_x=None, init_y=None,
                       font_scale=1.0, color=(255, 255, 255),
                       thickness=2, font=cv2.FONT_HERSHEY_SIMPLEX,
                       track_mode='feature'):
    """
    追踪画面特征点 + 将文字PIN到追踪位置
    
    track_mode:
      'feature' — 特征点追踪（自动选择画面中心区域的强特征点）
      'fixed'   — 固定位置（不追踪）
    
    文字会跟随画面运动，保持相对位置不变
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"无法打开视频: {video_path}")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    if end_frame is None:
        end_frame = frame_count
    
    if init_x is None:
        init_x = width // 2
    if init_y is None:
        init_y = height // 2
    
    print(f"[PinText] 文字: '{text}'  位置: ({init_x},{init_y})  帧: {start_frame}-{end_frame}")
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    # 跳到起始帧
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    
    prev_gray = None
    prev_point = np.array([[init_x, init_y]], dtype=np.float32).reshape(-1, 1, 2)
    current_offset = np.array([0.0, 0.0])
    
    for frame_idx in range(start_frame, min(end_frame, frame_count)):
        ret, frame = cap.read()
        if not ret:
            break
        
        curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        if track_mode == 'feature' and prev_gray is not None:
            # 追踪文字所在位置附近的特征点
            next_point, status, _ = cv2.calcOpticalFlowPyrLK(
                prev_gray, curr_gray, prev_point, None,
                winSize=(31, 31), maxLevel=2
            )
            
            if status is not None and status[0] == 1 and next_point is not None:
                # 计算位移
                dx = next_point[0][0] - prev_point[0][0]
                dy = next_point[0][1] - prev_point[0][1]
                current_offset += np.array([dx, dy])
                prev_point = next_point
        
        # 计算文字绘制位置
        text_x = int(init_x + current_offset[0])
        text_y = int(init_y + current_offset[1])
        
        # 绘制文字（带描边）
        # 描边
        cv2.putText(frame, text, (text_x + 2, text_y + 2), font, font_scale,
                    (0, 0, 0), thickness + 2, cv2.LINE_AA)
        # 主文字
        cv2.putText(frame, text, (text_x, text_y), font, font_scale,
                    color, thickness, cv2.LINE_AA)
        
        writer.write(frame)
        prev_gray = curr_gray
        
        if (frame_idx - start_frame) % 30 == 0:
            print(f"  [PinText] 帧 {frame_idx}/{end_frame}  偏移: ({current_offset[0]:.1f}, {current_offset[1]:.1f})")
    
    cap.release()
    writer.release()
    print(f"[PinText] ✅ 文字贴合视频: {output_path}")
    return output_path


# ==================== 运动数据导出 ====================

def export_motion_data(track_json, output_json, mode='zoom_pan'):
    """
    从追踪数据中提取运动趋势，导出为 video-editor.js 可用的 keyframes 格式
    
    mode='zoom_pan': 
      分析整体运动趋势 → 输出 zoom(缩放) + pan(平移) 关键帧
      
    输出格式:
    {
      "keyframes": {
        "zoom": { "from": 1.0, "to": 1.05 },
        "pan": { "from": [0, 0], "to": [12, -8] }
      },
      "stability": 0.85,   // 稳定性评分 0-1
      "motion_intensity": "low" | "medium" | "high"
    }
    """
    with open(track_json, 'r') as f:
        track_data = json.load(f)
    
    tracks = track_data.get('tracks', [])
    if not tracks:
        return {"keyframes": {}, "stability": 0, "motion_intensity": "unknown"}
    
    # 分析所有轨迹的位移
    displacements = []
    for t in tracks:
        if 'displacement' in t:
            displacements.append(t['displacement'])
    
    if not displacements:
        return {"keyframes": {}, "stability": 0, "motion_intensity": "unknown"}
    
    disp_arr = np.array(displacements)
    
    # 计算平均位移
    mean_dx = float(np.mean(disp_arr[:, 0]))
    mean_dy = float(np.mean(disp_arr[:, 1]))
    
    # 计算位移标准差（稳定性指标）
    std_dx = float(np.std(disp_arr[:, 0]))
    std_dy = float(np.std(disp_arr[:, 1]))
    stability = 1.0 / (1.0 + (std_dx + std_dy) / 100.0)
    stability = max(0.0, min(1.0, stability))
    
    # 运动强度
    total_motion = abs(mean_dx) + abs(mean_dy) + (std_dx + std_dy) * 0.5
    if total_motion < 20:
        motion_intensity = 'low'
    elif total_motion < 80:
        motion_intensity = 'medium'
    else:
        motion_intensity = 'high'
    
    # 缩放趋势（通过分析点之间距离的变化）
    # 简化版：用位移方差近似
    zoom_trend = 1.0 + min(0.2, max(-0.1, (std_dx + std_dy) / 5000.0))
    
    result = {
        "keyframes": {
            "zoom": {
                "from": 1.0,
                "to": round(zoom_trend, 4)
            },
            "pan": {
                "from": [0, 0],
                "to": [round(mean_dx, 2), round(mean_dy, 2)]
            }
        },
        "stability": round(stability, 4),
        "motion_intensity": motion_intensity,
        "avg_displacement": [round(mean_dx, 2), round(mean_dy, 2)],
        "track_count": len(tracks)
    }
    
    with open(output_json, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    print(f"[MotionData] ✅ 运动数据: {output_json}")
    print(f"  稳定性: {result['stability']}  强度: {motion_intensity}")
    print(f"  平均位移: ({mean_dx:.1f}, {mean_dy:.1f})")
    return result


# ==================== CLI ====================

def main():
    parser = argparse.ArgumentParser(description='光湖OpenCV平面追踪引擎')
    parser.add_argument('mode', choices=['track', 'stabilize', 'pin_text', 'motion_data'],
                        help='运行模式')
    parser.add_argument('video', help='输入视频路径')
    parser.add_argument('-o', '--output', help='输出路径')
    parser.add_argument('--text', help='pin_text模式的文字内容')
    parser.add_argument('--x', type=int, help='文字初始X坐标')
    parser.add_argument('--y', type=int, help='文字初始Y坐标')
    parser.add_argument('--track-json', help='motion_data模式的追踪数据JSON')
    parser.add_argument('--smoothing', type=int, default=25, help='稳定化平滑窗口')
    
    args = parser.parse_args()
    
    if args.mode == 'track':
        out = args.output or args.video.rsplit('.', 1)[0] + '_tracks.json'
        track_features(args.video, out)
    
    elif args.mode == 'stabilize':
        out = args.output or args.video.rsplit('.', 1)[0] + '_stabilized.mp4'
        stabilize_video(args.video, out, smoothing_radius=args.smoothing)
    
    elif args.mode == 'pin_text':
        if not args.text:
            print("错误: pin_text模式需要 --text 参数")
            sys.exit(1)
        out = args.output or args.video.rsplit('.', 1)[0] + '_pinned.mp4'
        pin_text_to_video(args.video, out, args.text, init_x=args.x, init_y=args.y)
    
    elif args.mode == 'motion_data':
        if not args.track_json:
            print("错误: motion_data模式需要 --track-json 参数")
            sys.exit(1)
        out = args.output or args.video.rsplit('.', 1)[0] + '_motion.json'
        export_motion_data(args.track_json, out)


if __name__ == '__main__':
    main()