"""Video scene-cut detection, clip extraction, and subject-centered crop/resize.

Uses plain (non-relative) imports of sibling modules so this file works whether
invoked as ``python dataset_builder/build.py`` or ``python -m dataset_builder.build``
(build.py inserts its own directory onto sys.path before importing).
"""
import math
import statistics

import cv2
import imageio
import numpy as np
from PIL import Image

import crop_resize
from detect import detect_subject_bbox, subdivide_bbox


def _frame_hist(frame_bgr):
    hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
    hist = cv2.calcHist([hsv], [0, 1], None, [50, 60], [0, 180, 0, 256])
    cv2.normalize(hist, hist)
    return hist


def find_scene_cuts(video_path, diff_thresh=0.45, sample_stride=1):
    """Returns a sorted list of frame indices where a scene cut is detected,
    via Bhattacharyya distance between consecutive HSV histograms."""
    cap = cv2.VideoCapture(video_path)
    cuts = []
    prev_hist = None
    idx = 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if idx % sample_stride == 0:
            hist = _frame_hist(frame)
            if prev_hist is not None:
                score = cv2.compareHist(prev_hist, hist, cv2.HISTCMP_BHATTACHARYYA)
                if score > diff_thresh:
                    cuts.append(idx)
            prev_hist = hist
        idx += 1
    cap.release()
    return cuts


def build_clip_ranges(total_frames, fps, cuts, clip_min_sec=2.0, clip_max_sec=6.0):
    """Turns scene-cut frame indices into (start, end) frame ranges honoring
    clip_min/max seconds: too-short segments are merged forward, segments
    longer than clip_max are chopped into roughly-equal sub-segments."""
    min_frames = max(1, int(clip_min_sec * fps))
    max_frames = max(min_frames, int(clip_max_sec * fps))

    boundaries = [0] + sorted(set(c for c in cuts if 0 < c < total_frames)) + [total_frames]
    raw_ranges = []
    start = boundaries[0]
    for b in boundaries[1:]:
        if b - start >= min_frames:
            raw_ranges.append((start, b))
            start = b
    if start < total_frames:
        if raw_ranges and (total_frames - start) < min_frames:
            s, _ = raw_ranges[-1]
            raw_ranges[-1] = (s, total_frames)
        else:
            raw_ranges.append((start, total_frames))

    ranges = []
    for s, e in raw_ranges:
        length = e - s
        if length <= max_frames:
            ranges.append((s, e))
            continue
        n_chunks = math.ceil(length / max_frames)
        chunk_len = length // n_chunks
        for i in range(n_chunks):
            cs = s + i * chunk_len
            ce = e if i == n_chunks - 1 else cs + chunk_len
            ranges.append((cs, ce))

    return [(s, e) for s, e in ranges if e - s >= max(1, min_frames // 2)]


def _sample_indices(start, end, n=3):
    if end - start <= n:
        return list(range(start, end))
    step = (end - start) / (n + 1)
    return [int(start + step * (i + 1)) for i in range(n)]


def _median_bbox(bboxes):
    return (
        statistics.median(b[0] for b in bboxes),
        statistics.median(b[1] for b in bboxes),
        statistics.median(b[2] for b in bboxes),
        statistics.median(b[3] for b in bboxes),
    )


def _motion_tag(frames_gray):
    if len(frames_gray) < 2:
        return "static, minimal motion"
    mags = []
    for a, b in zip(frames_gray, frames_gray[1:]):
        flow = cv2.calcOpticalFlowFarneback(a, b, None, 0.5, 3, 15, 3, 5, 1.2, 0)
        mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        mags.append(float(np.mean(mag)))
    avg = sum(mags) / len(mags)
    if avg < 0.5:
        return "static, minimal motion"
    if avg < 2.0:
        return "subtle motion"
    if avg < 5.0:
        return "moderate motion"
    return "fast, dynamic motion"


def extract_clip(video_path, start, end, target_fps, max_tier="1080p", pad_overrides=None,
                  body_region="auto", body_region_frac=0.45):
    """Reads frames [start, end) from video_path, computes one stable subject
    crop for the whole clip (median bbox over sampled frames), resamples to
    target_fps, trims to a 4n+1 frame count, and returns
    (frames_bgr, (target_w, target_h), representative_pil_image, motion_tag),
    or None if no frames could be read."""
    cap = cv2.VideoCapture(video_path)
    src_fps = cap.get(cv2.CAP_PROP_FPS) or target_fps
    img_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    img_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    bboxes, kinds = [], []
    for idx in _sample_indices(start, end, n=3):
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ok, frame = cap.read()
        if not ok:
            continue
        bbox, kind = detect_subject_bbox(frame)
        bboxes.append(bbox)
        kinds.append(kind)

    if not bboxes:
        cap.release()
        return None

    bbox = _median_bbox(bboxes)
    kind = max(set(kinds), key=kinds.count)
    if kind == "person" and body_region != "auto":
        bbox = subdivide_bbox(bbox, body_region, body_region_frac)
        kind = f"person_{body_region}"
    crop_box, (tw, th), _bucket_name, _tier = crop_resize.compute_crop(
        img_w, img_h, bbox, kind, pad_overrides, max_tier
    )

    step = src_fps / target_fps
    raw_indices = []
    f = float(start)
    while f < end:
        raw_indices.append(int(round(f)))
        f += step
    if len(raw_indices) >= 5:
        n = ((len(raw_indices) - 1) // 4) * 4 + 1
        raw_indices = raw_indices[:n]

    frames_out, frames_gray = [], []
    for idx in raw_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ok, frame = cap.read()
        if not ok:
            continue
        cropped = crop_resize.crop_resize_np(frame, crop_box, tw, th)
        frames_out.append(cropped)
        if len(frames_gray) < 6:
            frames_gray.append(cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY))
    cap.release()

    if not frames_out:
        return None

    motion_tag = _motion_tag(frames_gray)
    mid = frames_out[len(frames_out) // 2]
    representative = Image.fromarray(cv2.cvtColor(mid, cv2.COLOR_BGR2RGB))

    return frames_out, (tw, th), representative, motion_tag


def write_clip_mp4(frames_bgr, out_path, fps):
    with imageio.get_writer(str(out_path), fps=fps, codec="libx264", quality=8) as writer:
        for frame in frames_bgr:
            writer.append_data(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))