# -*- coding: utf-8 -*-
"""Build density features from already extracted frame/crop images.

Use this after frame extraction or after RF-DETR soup masking. Labels are still
read from the CSV only as training targets, never as model inputs.
"""

from __future__ import annotations

import argparse
import csv
import json
import re
from collections import defaultdict
from pathlib import Path
from typing import Any

import cv2
import numpy as np

from build_density_features import base_label_row, feature_keys, image_features, write_csv


ROOT = Path(__file__).resolve().parents[1]
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def main() -> int:
    parser = argparse.ArgumentParser(description="Build density features from saved images.")
    parser.add_argument("--csv", required=True, help="Source CSV containing ID and label columns.")
    parser.add_argument("--image-dir", required=True, help="Directory containing frame or soup crop images.")
    parser.add_argument("--output", default="local_batch_api/density_features/density_features_from_images.csv")
    parser.add_argument("--limit", type=int, default=0)
    args = parser.parse_args()

    csv_path = Path(args.csv).expanduser().resolve()
    image_dir = resolve_path(args.image_dir)
    output_path = resolve_path(args.output)

    labels = read_label_rows(csv_path)
    grouped = group_images_by_id(image_dir)
    ids = list(labels.keys())
    if args.limit > 0:
        ids = ids[: args.limit]

    rows: list[dict[str, Any]] = []
    for video_id in ids:
        images = grouped.get(video_id, [])
        if not images:
            rows.append(error_row(labels[video_id], f"no images for id={video_id}"))
            continue
        print(f"IMAGE_FEATURES {video_id} frames={len(images)}")
        frame_records = []
        for image_path in images:
            image = read_image(image_path)
            if image is None:
                continue
            frame_records.append(
                {
                    "image_path": str(image_path),
                    "ratio": parse_ratio(image_path.name),
                    **image_features(image),
                }
            )
        if not frame_records:
            rows.append(error_row(labels[video_id], f"no readable images for id={video_id}"))
            continue
        rows.append(aggregate_image_features(labels[video_id], frame_records))

    write_csv(output_path, rows)
    jsonl_path = output_path.with_suffix(".jsonl")
    with jsonl_path.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(f"done rows={len(rows)} output={output_path}")
    return 0


def resolve_path(value: str) -> Path:
    path = Path(value)
    return path if path.is_absolute() else ROOT / path


def read_label_rows(csv_path: Path) -> dict[str, Any]:
    from build_density_features import VideoItem

    rows: dict[str, Any] = {}
    with csv_path.open("r", encoding="utf-8-sig", newline="") as fh:
        for idx, row in enumerate(csv.DictReader(fh)):
            video_id = (row.get("ID") or row.get("id") or "").strip()
            if not video_id:
                continue
            rows[video_id] = VideoItem(video_path=Path(f"{video_id}.MOV"), row_index=idx, metadata=dict(row))
    return rows


def group_images_by_id(image_dir: Path) -> dict[str, list[Path]]:
    grouped: dict[str, list[Path]] = defaultdict(list)
    for path in sorted(p for p in image_dir.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS):
        match = re.match(r"(?P<id>\d+)(?:_r\d+)?", path.stem)
        if match:
            grouped[match.group("id")].append(path)
    return dict(grouped)


def read_image(path: Path) -> np.ndarray | None:
    data = np.frombuffer(path.read_bytes(), dtype=np.uint8)
    return cv2.imdecode(data, cv2.IMREAD_COLOR)


def parse_ratio(name: str) -> float:
    match = re.search(r"_r(\d+)", name)
    return float(match.group(1)) / 100.0 if match else 0.0


def aggregate_image_features(item: Any, frames: list[dict[str, Any]]) -> dict[str, Any]:
    row: dict[str, Any] = base_label_row(item)
    row["frame_count"] = len(frames)
    row["frame_ratios"] = ",".join(str(frame["ratio"]) for frame in frames)
    for key in feature_keys():
        values = np.array([float(frame[key]) for frame in frames], dtype=np.float64)
        row[f"{key}_mean"] = float(values.mean())
        row[f"{key}_median"] = float(np.median(values))
        row[f"{key}_min"] = float(values.min())
        row[f"{key}_max"] = float(values.max())
    row["frames_json"] = json.dumps(frames, ensure_ascii=False)
    row["error"] = ""
    return row


def error_row(item: Any, error: str) -> dict[str, Any]:
    row = base_label_row(item)
    row["frame_count"] = 0
    row["frame_ratios"] = ""
    row["frames_json"] = "[]"
    row["error"] = error
    return row


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