#!/usr/bin/env python3
"""
起重机挂钩数据集生成全流程
============================
Step 1: 视频抽帧 → Step 2: Engagement筛选 → Step 3: Ready注入
→ Step 4: 质量过滤 → Step 5: 高阈值重推理 → Step 6: 质量排序 → Step 7: CVAT导出

用法:
  python pipeline.py                          # 运行全部步骤
  python pipeline.py --skip-extract           # 跳过抽帧
  python pipeline.py --skip-extract --skip-select  # 跳过多个步骤
"""
import os
import sys
import re
import time
import shutil
import zipfile
import subprocess
import numpy as np
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

from config import *


# ============================================================
# Step 1: 视频抽帧 (FFmpeg NVDEC硬件加速)
# ============================================================
def step1_extract():
    print("=" * 60)
    print("Step 1: 视频抽帧 (FFmpeg NVDEC)")
    print("=" * 60)

    if EXTRACTED_DIR.exists():
        print(f"  已存在 {EXTRACTED_DIR}，跳过抽帧")
        return

    videos = []
    for mp4 in VIDEO_DIR.rglob('*.mp4'):
        parts = mp4.parts
        try:
            cam_idx = parts.index(VIDEO_DIR_NAME)
            cam_id = parts[cam_idx + 1]
            cam_name = f'cam{cam_id.zfill(2)}'
            date = parts[cam_idx + 2]
        except (IndexError, ValueError):
            continue
        videos.append((mp4, cam_name, date))

    print(f"  找到 {len(videos)} 个视频")
    EXTRACTED_DIR.mkdir(parents=True, exist_ok=True)
    total_frames = 0

    for i, (mp4, cam_name, date) in enumerate(videos, 1):
        out_dir = EXTRACTED_DIR / cam_name
        out_dir.mkdir(exist_ok=True)
        stem = mp4.stem
        pattern = str(out_dir / f'{cam_name}_{date}_{stem}_%06d.jpg')

        cmd = ['ffmpeg', '-hide_banner', '-loglevel', 'error',
               '-hwaccel', 'nvdec', '-c:v', 'hevc_cuvid',
               '-i', str(mp4),
               '-qmin', '1', '-qmax', '1', '-vsync', '0', pattern]

        print(f"  [{i}/{len(videos)}] {cam_name}/{date}/{stem}...", end=' ', flush=True)
        t0 = time.time()
        subprocess.run(cmd, check=True)

        n_frames = len(list(out_dir.glob(f'{cam_name}_{date}_{stem}_*.jpg')))
        for j in range(1, n_frames + 1):
            lbl = out_dir / f'{cam_name}_{date}_{stem}_{j:06d}.txt'
            if not lbl.exists():
                lbl.touch()

        elapsed = time.time() - t0
        total_frames += n_frames
        print(f"{n_frames} frames in {elapsed:.0f}s")

    print(f"\n  抽帧完成: {total_frames} 帧")


# ============================================================
# Step 2: Engagement评分 + 去重 + 排序选取
# ============================================================
def _score_and_pos(dets):
    """评分：hoist+ear共现=0.7, 仅ear=0.4, 中心距离奖励, 多hoist惩罚"""
    if len(dets) == 0:
        return 0
    has_hoist = np.any(dets[:, 0] == 0)
    has_ear = np.any(dets[:, 0] == 1)
    n_hoist = int(np.sum(dets[:, 0] == 0))

    score = 0.7 if (has_hoist and has_ear) else (0.4 if has_ear else 0.0)

    avg_xc = np.mean(dets[:, 1])
    avg_yc = np.mean(dets[:, 2])
    center_dist = ((avg_xc - 0.5)**2 + (avg_yc - 0.5)**2) ** 0.5
    score += max(0, 0.3 - center_dist) * 0.5
    if n_hoist > 1:
        score -= (n_hoist - 1) * 0.1
    return max(0, score)


def _fix_duplicates(labels):
    """去重：ear保留最小框，hoist保留最大框"""
    for cls_id, keep_func in [(1, 'min'), (0, 'max')]:
        cls_mask = labels[:, 0] == cls_id
        if np.sum(cls_mask) <= 1:
            continue
        cls_dets = labels[cls_mask]
        areas = cls_dets[:, 3] * cls_dets[:, 4]
        keep_idx = np.argmin(areas) if keep_func == 'min' else np.argmax(areas)
        keep_mask = np.ones(len(labels), dtype=bool)
        cls_indices = np.where(cls_mask)[0]
        for ri in range(len(cls_dets)):
            if ri != keep_idx:
                keep_mask[cls_indices[ri]] = False
        labels = labels[keep_mask]
    return labels


def step2_select():
    print("\n" + "=" * 60)
    print("Step 2: Engagement评分 + 去重 + 选取engaged帧")
    print("=" * 60)

    if OUTPUT_DIR.exists():
        shutil.rmtree(str(OUTPUT_DIR))

    cameras = sorted([d.name for d in EXTRACTED_IMG_DIR.iterdir() if d.is_dir()])
    all_kept = {}

    for cam_name in cameras:
        t0 = time.time()
        lbl_dir = EXTRACTED_LBL_DIR / cam_name

        # 读取所有标签
        labels = {}
        with os.scandir(str(lbl_dir)) as it:
            for entry in it:
                if not entry.name.endswith('.txt'):
                    continue
                with open(entry.path) as f:
                    data = f.read().strip()
                stem = entry.name[:-4]
                if data:
                    rows = []
                    for line in data.split('\n'):
                        parts = line.split()
                        if len(parts) >= 5:
                            rows.append([float(x) for x in parts[:5]])
                    labels[stem] = np.array(rows) if rows else np.zeros((0, 5))
                else:
                    labels[stem] = np.zeros((0, 5))

        # 评分
        scores = {name: _score_and_pos(dets) for name, dets in labels.items()}

        # 去重
        for name, dets in labels.items():
            if len(dets) > 0:
                labels[name] = _fix_duplicates(dets)

        # 按分数排序
        ranked = sorted(scores.keys(), key=lambda n: scores[n], reverse=True)

        # 筛选
        selected = []
        last_seq = -999
        for name in ranked:
            if scores[name] < SCORE_THRESH:
                continue
            dets = labels[name]
            if not np.any((dets[:, 0] == 0) | (dets[:, 0] == 1)):
                continue
            m = re.search(r'_(\d{6})$', name)
            seq = int(m.group(1)) if m else 0
            if abs(seq - last_seq) < SEQ_GAP or seq == last_seq:
                continue
            selected.append(name)
            last_seq = seq

        # 创建输出目录 + 符号链接
        out_img = OUTPUT_DIR / 'images' / cam_name
        out_lbl = OUTPUT_DIR / 'labels' / cam_name
        out_img.mkdir(parents=True, exist_ok=True)
        out_lbl.mkdir(parents=True, exist_ok=True)

        img_dir = EXTRACTED_IMG_DIR / cam_name
        for name in selected:
            si = img_dir / f'{name}.jpg'
            sl = EXTRACTED_LBL_DIR / cam_name / f'{name}.txt'
            di = out_img / f'czzt_{name}.jpg'
            dl = out_lbl / f'czzt_{name}.txt'
            if si.exists() and not di.exists():
                os.symlink(str(si.resolve()), str(di))
            if sl.exists() and not dl.exists():
                os.symlink(str(sl.resolve()), str(dl))

        all_kept[cam_name] = len(selected)
        print(f"  {cam_name}: {len(selected)} images in {time.time() - t0:.1f}s")

    total = sum(all_kept.values())
    print(f"\n  Step 2 完成: {total} images")
    return all_kept


# ============================================================
# Step 3: 注入 ready 标签 (YOLOv5推理)
# ============================================================
def _calc_iou(boxA, boxB):
    x1 = max(boxA[0], boxB[0]); y1 = max(boxA[1], boxB[1])
    x2 = min(boxA[2], boxB[2]); y2 = min(boxA[3], boxB[3])
    inter = max(0, x2-x1) * max(0, y2-y1)
    areaA = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
    areaB = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
    union = areaA + areaB - inter
    return inter / union if union > 0 else 0


def step3_inject_ready(cam_name):
    import torch
    import cv2

    img_dir = OUTPUT_DIR / 'images' / cam_name
    lbl_dir = OUTPUT_DIR / 'labels' / cam_name
    if not img_dir.exists():
        return 0

    images = sorted(img_dir.glob('czzt_*.jpg'))
    if not images:
        return 0

    sys.path.insert(0, YOLOV5_DIR)
    from models.common import DetectMultiBackend
    from utils.general import non_max_suppression, scale_boxes
    from utils.torch_utils import select_device

    device = select_device('')
    model = DetectMultiBackend(MODEL_PATH, device=device)
    model.eval()

    injected = 0
    for img_path in images:
        lbl_file = lbl_dir / f'{img_path.stem}.txt'

        # 读取已有标签
        existing = []
        if lbl_file.exists():
            with open(lbl_file) as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) >= 5:
                        existing.append([int(parts[0])] + [float(x) for x in parts[1:]])

        # 需要同时有hoist和ear才注入ready
        if not (any(r[0] == 0 for r in existing) and any(r[0] == 1 for r in existing)):
            continue

        # 推理
        img_bgr = cv2.imread(str(img_path))
        if img_bgr is None:
            continue
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        h0, w0 = img_rgb.shape[:2]
        img_resized = cv2.resize(img_rgb, (IMG_WIDTH, IMG_HEIGHT))
        img_tensor = img_resized.transpose(2, 0, 1)
        img_tensor = np.ascontiguousarray(img_tensor)
        img_tensor = torch.from_numpy(img_tensor).to(device).float() / 255.0
        if img_tensor.ndimension() == 3:
            img_tensor = img_tensor.unsqueeze(0)

        pred = model(img_tensor)[0]
        pred = non_max_suppression(pred, READY_CONF_THRESH, READY_IOU_THRESH, classes=[2])

        if len(pred) == 0 or len(pred[0]) == 0:
            continue

        dets = pred[0]
        dets[:, :4] = scale_boxes(img_tensor.shape[2:], dets[:, :4], (h0, w0)).round()
        areas = (dets[:, 2] - dets[:, 0]) * (dets[:, 3] - dets[:, 1])

        # 收集ear框用于IoU过滤
        ear_boxes = []
        for r in existing:
            if r[0] == 1:
                ex, ey, ew, eh = r[1], r[2], r[3], r[4]
                ear_boxes.append([(ex-ew/2)*w0, (ey-eh/2)*h0, (ex+ew/2)*w0, (ey+eh/2)*h0])

        # 找最大的不与ear重叠的ready
        sorted_idx = torch.argsort(areas, descending=True)
        best = None
        for idx in sorted_idx:
            det = dets[idx]
            rbox = [det[0].item(), det[1].item(), det[2].item(), det[3].item()]
            if all(_calc_iou(rbox, ebox) <= READY_EAR_IOU_THRESH for ebox in ear_boxes):
                best = det
                break

        if best is None:
            continue

        xc = (best[0] + best[2]) / 2 / w0
        yc = (best[1] + best[3]) / 2 / h0
        w = (best[2] - best[0]) / w0
        h = (best[3] - best[1]) / h0

        with open(lbl_file, 'a') as f:
            f.write(f'2 {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}\n')
        injected += 1

    return injected


def step3_inject():
    print("\n" + "=" * 60)
    print("Step 3: 注入 ready 标签 (YOLOv5推理)")
    print("=" * 60)

    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
    total_injected = 0
    for cam_name in cameras:
        n = step3_inject_ready(cam_name)
        total_injected += n
        print(f"  {cam_name}: injected ready to {n} images")
    print(f"\n  Step 3 完成: {total_injected} ready labels injected")


# ============================================================
# Step 4: 质量过滤
# ============================================================
def step4_filter():
    print("\n" + "=" * 60)
    print("Step 4: 质量过滤")
    print("=" * 60)

    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
    total_keep = 0

    for cam_name in cameras:
        img_dir = OUTPUT_DIR / 'images' / cam_name
        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
        if not lbl_dir.exists():
            continue

        discard_ratio = discard_flat = discard_both = keep = 0

        for lbl_file in list(lbl_dir.glob('czzt_*.txt')):
            with open(lbl_file) as f:
                lines = f.readlines()

            hoist_areas, ready_areas, ear_aspects = [], [], []
            for line in lines:
                parts = line.strip().split()
                if len(parts) < 5:
                    continue
                cls = int(parts[0])
                w, h = float(parts[3]), float(parts[4])
                if cls == 0:
                    hoist_areas.append(w * h)
                elif cls == 2:
                    ready_areas.append(w * h)
                elif cls == 1 and h > 0:
                    ear_aspects.append(w / h)

            fail_ratio = (hoist_areas and ready_areas and
                          max(ready_areas) / max(hoist_areas) >= READY_HOIST_RATIO_THRESH)
            fail_flat = (ear_aspects and all(a >= EAR_FLAT_ASPECT_THRESH for a in ear_aspects))

            if fail_ratio and fail_flat:
                discard_both += 1
            elif fail_ratio:
                discard_ratio += 1
            elif fail_flat:
                discard_flat += 1
            else:
                keep += 1
                continue

            img_file = img_dir / (lbl_file.stem + '.jpg')
            if img_file.exists():
                img_file.unlink()
            lbl_file.unlink()

        total_keep += keep
        print(f"  {cam_name}: keep={keep}, discard_ratio={discard_ratio}, discard_flat={discard_flat}, discard_both={discard_both}")

    print(f"\n  Step 4 完成: {total_keep} images after filtering")


# ============================================================
# Step 5: 重新YOLO推理 (高置信度阈值, 清除重叠框)
# ============================================================
def step5_reinfer():
    import torch
    import cv2

    print("\n" + "=" * 60)
    print(f"Step 5: 重新YOLO推理 (conf>={REINFER_CONF}, iou={REINFER_IOU})")
    print("=" * 60)

    sys.path.insert(0, YOLOV5_DIR)
    from models.common import DetectMultiBackend
    from utils.general import non_max_suppression, scale_boxes
    from utils.torch_utils import select_device

    device = select_device('')
    model = DetectMultiBackend(MODEL_PATH, device=device)
    model.eval()

    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
    total_reinferred = 0

    for cam_name in cameras:
        img_dir = OUTPUT_DIR / 'images' / cam_name
        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
        if not img_dir.exists():
            continue

        images = sorted(img_dir.glob('czzt_*.jpg'))
        if not images:
            continue

        reinferred = 0
        for img_path in images:
            img_bgr = cv2.imread(str(img_path))
            if img_bgr is None:
                continue
            img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
            h0, w0 = img_rgb.shape[:2]
            img_resized = cv2.resize(img_rgb, (IMG_WIDTH, IMG_HEIGHT))
            img_tensor = img_resized.transpose(2, 0, 1)
            img_tensor = np.ascontiguousarray(img_tensor)
            img_tensor = torch.from_numpy(img_tensor).to(device).float() / 255.0
            if img_tensor.ndimension() == 3:
                img_tensor = img_tensor.unsqueeze(0)

            pred = model(img_tensor)[0]
            pred = non_max_suppression(pred, REINFER_CONF, REINFER_IOU)

            lbl_file = lbl_dir / f'{img_path.stem}.txt'

            if len(pred) == 0 or len(pred[0]) == 0:
                lbl_file.write_text('')
                reinferred += 1
                continue

            dets = pred[0]
            dets[:, :4] = scale_boxes(img_tensor.shape[2:], dets[:, :4], (h0, w0)).round()

            with open(lbl_file, 'w') as f:
                for det in dets:
                    cls_id = int(det[5])
                    if cls_id not in CLASSES:
                        continue
                    xc = (det[0] + det[2]) / 2 / w0
                    yc = (det[1] + det[3]) / 2 / h0
                    w = (det[2] - det[0]) / w0
                    h = (det[3] - det[1]) / h0
                    f.write(f'{cls_id} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}\n')
            reinferred += 1

        total_reinferred += reinferred
        print(f"  {cam_name}: reinferred {reinferred} images")

    print(f"\n  Step 5 完成: {total_reinferred} images reinferred")


# ============================================================
# Step 6: 质量排序取top MAX_PER_CAM
# ============================================================
def step6_select_top():
    print("\n" + "=" * 60)
    print(f"Step 6: 质量排序取top {MAX_PER_CAM}")
    print("=" * 60)

    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])

    for cam_name in cameras:
        img_dir = OUTPUT_DIR / 'images' / cam_name
        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
        if not img_dir.exists():
            continue

        images = sorted(img_dir.glob('czzt_*.jpg'))
        if len(images) <= MAX_PER_CAM:
            print(f"  {cam_name}: {len(images)} images (<= {MAX_PER_CAM}, no trimming)")
            continue

        def compute_quality(img_path):
            try:
                import cv2
                img = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
                if img is None:
                    return 0
                brightness = np.mean(img) / 255.0
                laplacian = cv2.Laplacian(img, cv2.CV_64F)
                sharpness = laplacian.var() / 1000.0
                return brightness * QUALITY_BRIGHTNESS_WEIGHT + min(sharpness, 1.0) * QUALITY_SHARPNESS_WEIGHT
            except Exception:
                return 0

        quality = {}
        with ThreadPoolExecutor(max_workers=16) as pool:
            for img_path in images:
                quality[img_path.name] = compute_quality(img_path)

        ranked = sorted(quality.keys(), key=lambda n: quality[n], reverse=True)
        keep_names = set(ranked[:MAX_PER_CAM])

        removed = 0
        for img_path in images:
            if img_path.name not in keep_names:
                img_path.unlink()
                lbl = lbl_dir / (img_path.stem + '.txt')
                if lbl.exists():
                    lbl.unlink()
                removed += 1

        print(f"  {cam_name}: kept {len(keep_names)}, removed {removed}")


# ============================================================
# Step 7: 生成 CVAT ZIP
# ============================================================
def step7_generate_cvat():
    print("\n" + "=" * 60)
    print("Step 7: 生成 CVAT ZIP")
    print("=" * 60)

    CVAT_ZIP_DIR.mkdir(parents=True, exist_ok=True)
    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])

    for cam_name in cameras:
        img_dir = OUTPUT_DIR / 'images' / cam_name
        lbl_dir = OUTPUT_DIR / 'labels' / cam_name

        img_names = sorted([f.name for f in img_dir.glob('czzt_*.jpg')])
        if not img_names:
            print(f"  {cam_name}: 0 images, skipping")
            continue

        zip_path = CVAT_ZIP_DIR / f'{cam_name}.zip'
        obj_names_str = '\n'.join(CVAT_CLASSES) + '\n'
        obj_data = f'classes = {len(CVAT_CLASSES)}\ntrain = data/train.txt\nnames = data/obj.names\nbackup = backup/\n'
        train_lines = '\n'.join(f'data/obj_train_data/data/obj_train_data/{n}' for n in img_names) + '\n'

        with zipfile.ZipFile(str(zip_path), 'w', zipfile.ZIP_STORED) as zf:
            zf.writestr('obj.names', obj_names_str)
            zf.writestr('obj.data', obj_data)
            zf.writestr('train.txt', train_lines)
            for img_name in img_names:
                zf.write(str(img_dir / img_name), f'obj_train_data/data/obj_train_data/{img_name}')
                stem = img_name[:-4]
                lbl_file = lbl_dir / f'{stem}.txt'
                if lbl_file.exists():
                    zf.write(str(lbl_file), f'obj_train_data/data/obj_train_data/{stem}.txt')

        print(f"  {cam_name}: {len(img_names)} images -> {zip_path.name}")

    print(f"\n  Step 7 完成: CVAT zips in {CVAT_ZIP_DIR}")


# ============================================================
# 最终统计
# ============================================================
def print_stats():
    print("\n" + "=" * 60)
    print("最终统计")
    print("=" * 60)

    cameras = sorted([d.name for d in (OUTPUT_DIR / 'images').iterdir() if d.is_dir()])
    total = 0
    label_counts = {0: 0, 1: 0, 2: 0}

    for cam_name in cameras:
        img_dir = OUTPUT_DIR / 'images' / cam_name
        lbl_dir = OUTPUT_DIR / 'labels' / cam_name
        n_img = len(list(img_dir.glob('czzt_*.jpg'))) if img_dir.exists() else 0
        n_lbl = {0: 0, 1: 0, 2: 0}
        if lbl_dir.exists():
            for lf in lbl_dir.glob('czzt_*.txt'):
                with open(lf) as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) >= 5:
                            c = int(parts[0])
                            if c in n_lbl:
                                n_lbl[c] += 1
                                label_counts[c] += 1
        total += n_img
        print(f"  {cam_name}: {n_img} images, hoist={n_lbl[0]}, ear={n_lbl[1]}, ready={n_lbl[2]}")

    print(f"\n  总计: {total} images")
    print(f"  hoist(0): {label_counts[0]}")
    print(f"  ear(1): {label_counts[1]}")
    print(f"  ready(2): {label_counts[2]}")


# ============================================================
# Main
# ============================================================
def main():
    skip = set(sys.argv[1:])

    if '--skip-extract' not in skip:
        step1_extract()
    if '--skip-select' not in skip:
        step2_select()
    if '--skip-inject' not in skip:
        step3_inject()
    if '--skip-filter' not in skip:
        step4_filter()
    if '--skip-reinfer' not in skip:
        step5_reinfer()
    if '--skip-select-top' not in skip:
        step6_select_top()
    if '--skip-zip' not in skip:
        step7_generate_cvat()

    print_stats()


if __name__ == '__main__':
    main()