# -*- coding: utf-8 -*-
"""多フレーム抽出 → 湯気分類で除外 → 残りからスープ部分を抽出（検出テスト）。

以前作成した2モデルを動画に適用する:
  1) 湯気分類器 (EfficientNet-B0): フレームを「湯気多/少」に分類し、多いものを除外
  2) RF-DETR セグメンテーション: 湯気が少ないフレームから soup_valid_area を抽出

各動画から複数フレーム(既定10枚)を ffmpeg で等間隔抽出して検証する。
"""

from __future__ import annotations

import argparse
import csv
import json
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

import cv2
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
SUITE = (
    Path.home()
    / "Documents"
    / "開発しているもの"
    / "制作中"
    / "AIラーメン開発"
    / "ramen-soup-ai-suite"
)
CROP_SRC = SUITE / "03_soup-density-ai" / "src"
CROP_CKPT = (
    SUITE
    / "03_soup-density-ai"
    / "models"
    / "crop"
    / "soup_valid_area_rfdetr_seg_v2"
    / "checkpoint_best_ema.pth"
)
STEAM_CKPT = (
    SUITE
    / "04_steam-classification"
    / "steam_detection_project"
    / "models"
    / "steam_classifier_best.pt"
)

VIDEO_EXTENSIONS = [".mov", ".mp4", ".m4v", ".avi", ".mkv"]


# --------------------------------------------------------------------------- #
# 湯気分類器
# --------------------------------------------------------------------------- #
def build_steam_model(checkpoint_path: Path, device):
    import torch

    try:
        import timm

        model = timm.create_model("efficientnet_b0", pretrained=False, num_classes=2)
    except Exception:
        from torchvision import models
        import torch.nn as nn

        model = models.efficientnet_b0(weights=None)
        in_features = model.classifier[-1].in_features
        model.classifier[-1] = nn.Linear(in_features, 2)

    checkpoint = torch.load(checkpoint_path, map_location=device)
    state = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
    model.load_state_dict(state)
    model.to(device)
    model.eval()
    return model


def steam_transform(input_size: int = 224):
    import albumentations as A
    from albumentations.pytorch import ToTensorV2

    return A.Compose(
        [
            A.Resize(height=input_size, width=input_size),
            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            ToTensorV2(),
        ]
    )


def steam_probability(model, transform, image_bgr, device) -> float:
    import torch

    rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    tensor = transform(image=rgb)["image"].unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(tensor)
        prob = torch.softmax(logits, dim=1)[0, 1].item()
    return float(prob)


# --------------------------------------------------------------------------- #
# フレーム抽出
# --------------------------------------------------------------------------- #
def probe_duration(video_path: Path) -> float:
    result = subprocess.run(
        [
            "ffprobe", "-v", "error", "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1", str(video_path),
        ],
        capture_output=True, text=True, timeout=30,
    )
    try:
        return float(result.stdout.strip())
    except ValueError:
        return 0.0


def extract_frame(video_path: Path, seek_seconds: float):
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        tmp_path = Path(tmp.name)
    try:
        result = subprocess.run(
            [
                "ffmpeg", "-hide_banner", "-loglevel", "error",
                "-ss", f"{seek_seconds:.3f}", "-i", str(video_path),
                "-frames:v", "1", "-q:v", "2", "-y", str(tmp_path),
            ],
            capture_output=True, text=True, timeout=60,
        )
        if result.returncode != 0 or not tmp_path.exists():
            return None
        data = np.frombuffer(tmp_path.read_bytes(), dtype=np.uint8)
        return cv2.imdecode(data, cv2.IMREAD_COLOR)
    finally:
        try:
            tmp_path.unlink()
        except FileNotFoundError:
            pass


def imwrite_unicode(path: Path, image) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    ok, encoded = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), 92])
    if ok:
        path.write_bytes(encoded.tobytes())


def resolve_video(video_dir: Path, video_id: str) -> Path | None:
    for ext in VIDEO_EXTENSIONS:
        for cand in (video_dir / f"{video_id}{ext}", video_dir / f"{video_id}{ext.upper()}"):
            if cand.exists():
                return cand
    return None


# --------------------------------------------------------------------------- #
# メイン
# --------------------------------------------------------------------------- #
def main() -> int:
    parser = argparse.ArgumentParser(description="湯気除外 + スープ抽出 の多フレーム検出テスト")
    parser.add_argument("--video-dir", default=str(ROOT / "学習データ"))
    parser.add_argument("--ids", default="191-200", help="例: 191-200 または 191,195,200")
    parser.add_argument("--frames", type=int, default=10, help="1動画あたりの抽出フレーム数")
    parser.add_argument("--steam-threshold", type=float, default=0.5,
                        help="この値以上を『湯気多』として除外")
    parser.add_argument("--crop-threshold", type=float, default=0.35)
    parser.add_argument("--output-dir", default=str(ROOT / "local_batch_api" / "detect_test"))
    args = parser.parse_args()

    import torch

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device={device}")

    video_dir = Path(args.video_dir)
    output_dir = Path(args.output_dir)
    frames_dir = output_dir / "frames"
    crops_dir = output_dir / "soup_crops"
    output_dir.mkdir(parents=True, exist_ok=True)

    ids = parse_ids(args.ids)

    # --- モデル読み込み ---
    print("loading steam classifier ...")
    steam_model = build_steam_model(STEAM_CKPT, device)
    steam_tf = steam_transform()

    print("loading RF-DETR soup extractor ...")
    sys.path.insert(0, str(CROP_SRC))
    from crop.extract_soup_area import build_predictor, extract_one_image

    crop_predictor = build_predictor("small", str(CROP_CKPT), args.crop_threshold)

    frame_rows: list[dict[str, Any]] = []
    summary_rows: list[dict[str, Any]] = []

    for vid in ids:
        video_path = resolve_video(video_dir, vid)
        if video_path is None:
            print(f"SKIP {vid}: video not found")
            continue
        duration = probe_duration(video_path)
        ratios = list(np.linspace(0.1, 0.9, args.frames))
        print(f"\n=== {vid} ({video_path.name}) dur={duration:.1f}s frames={args.frames} ===")

        kept = 0
        excluded = 0
        extracted = 0
        for i, ratio in enumerate(ratios):
            seek = max(0.0, duration * ratio)
            image = extract_frame(video_path, seek)
            if image is None:
                continue
            frame_path = frames_dir / f"{vid}_f{i:02d}.jpg"
            imwrite_unicode(frame_path, image)

            steam_prob = steam_probability(steam_model, steam_tf, image, device)
            is_high_steam = steam_prob >= args.steam_threshold

            row: dict[str, Any] = {
                "id": vid,
                "frame_index": i,
                "ratio": round(float(ratio), 3),
                "seek_seconds": round(seek, 2),
                "frame_path": str(frame_path),
                "steam_high_probability": round(steam_prob, 4),
                "steam_high_flag": int(is_high_steam),
                "decision": "excluded_high_steam" if is_high_steam else "kept_low_steam",
                "soup_extract_status": "",
                "soup_extract_confidence": "",
                "soup_crop_path": "",
            }

            if is_high_steam:
                excluded += 1
                print(f"  f{i:02d} steam={steam_prob:.3f} -> 除外(湯気多)")
            else:
                kept += 1
                result = extract_one_image(
                    image_path=frame_path,
                    output_dir=crops_dir,
                    checkpoint_path=str(CROP_CKPT),
                    model_size="small",
                    threshold=args.crop_threshold,
                    low_confidence_threshold=args.crop_threshold,
                    predictor=crop_predictor,
                    save_alpha=False,
                    preserve_filename=False,
                    output_kind="crop",
                    skip_fallback_output=True,
                )
                status = int(result["extract_status"])
                conf = float(result["extract_confidence"])
                row["soup_extract_status"] = status
                row["soup_extract_confidence"] = round(conf, 4)
                row["soup_crop_path"] = result["soup_crop_path"] if status > 0 else ""
                if status > 0:
                    extracted += 1
                tag = {0: "抽出失敗", 1: "抽出OK", 2: "抽出OK(低信頼)"}.get(status, str(status))
                print(f"  f{i:02d} steam={steam_prob:.3f} -> 採用 / スープ{tag} conf={conf:.3f}")

            frame_rows.append(row)

        summary_rows.append({
            "id": vid,
            "video": video_path.name,
            "duration_sec": round(duration, 1),
            "frames_total": args.frames,
            "kept_low_steam": kept,
            "excluded_high_steam": excluded,
            "soup_extracted": extracted,
            "soup_extract_rate": round(extracted / kept, 3) if kept else 0.0,
        })

    write_csv(output_dir / "frame_detections.csv", frame_rows)
    write_csv(output_dir / "summary.csv", summary_rows)
    (output_dir / "summary.json").write_text(
        json.dumps({"frames": frame_rows, "summary": summary_rows}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    print("\n================ SUMMARY ================")
    for s in summary_rows:
        print(
            f"{s['id']}: 抽出{s['frames_total']}枚 / 採用{s['kept_low_steam']} "
            f"除外{s['excluded_high_steam']} / スープ抽出成功{s['soup_extracted']}"
        )
    print(f"\noutput -> {output_dir}")
    return 0


def parse_ids(spec: str) -> list[str]:
    spec = spec.strip()
    if "-" in spec and "," not in spec:
        lo, hi = spec.split("-", 1)
        return [str(n) for n in range(int(lo), int(hi) + 1)]
    return [s.strip() for s in spec.split(",") if s.strip()]


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        path.write_text("", encoding="utf-8-sig")
        return
    fields: list[str] = []
    for row in rows:
        for key in row:
            if key not in fields:
                fields.append(key)
    with path.open("w", encoding="utf-8-sig", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


if __name__ == "__main__":
    raise SystemExit(main())