"""Shared HD-mating-arena ROI geometry, used by both pick_targets.py
(for live overlay) and track_videos.py (for actual tracking).

Pure numpy + cv2; no ethoscope dependency.
"""

from __future__ import annotations

import itertools

import cv2
import numpy as np

# Layout from
# ethoscope/.../roi_builders/roi_templates/builtin/HD_Mating_Arena_6_ROIS.json
HD_MATING_ARENA = {
    "n_rows": 2,
    "n_cols": 3,
    "top_margin": -0.21,
    "bottom_margin": -0.13,
    "left_margin": 0.05,
    "right_margin": 0.05,
    "horizontal_fill": 0.85,
    "vertical_fill": 1.3,
}

HD_FG_DATA = {
    "sample_size": 400,
    "normal_limits": [800, 2000],
    "tolerance": 0.8,
}


def compute_roi_polygons(reference_points, layout=HD_MATING_ARENA):
    """Map 3 L-shape reference points to 6 ROI polygons, in the order ROI 1..6.

    Reference points must be ordered:
        [TOP, CORNER, LEFT]
    matching ethoscope's dst_points = [(0, -1), (0, 0), (-1, 0)].

    Returns:
        list[np.ndarray]  # 6 arrays, each shape (4, 2), int32, in image coords
    """
    ref = np.asarray(reference_points, dtype=np.float32)
    if ref.shape != (3, 2):
        raise ValueError(f"reference_points must be 3x2, got shape {ref.shape}")

    dst_points = np.array([(0, -1), (0, 0), (-1, 0)], dtype=np.float32)
    wrap_mat = cv2.getAffineTransform(dst_points, ref)

    n_col = layout["n_cols"]
    n_row = layout["n_rows"]
    tm, bm = layout["top_margin"], layout["bottom_margin"]
    lm, rm = layout["left_margin"], layout["right_margin"]
    hf, vf = layout["horizontal_fill"], layout["vertical_fill"]

    y_positions = (np.arange(n_row) * 2.0 + 1) * (1 - tm - bm) / (2 * n_row) + tm
    x_positions = (np.arange(n_col) * 2.0 + 1) * (1 - lm - rm) / (2 * n_col) + lm
    centres = [np.array([x, y]) for x, y in itertools.product(x_positions, y_positions)]
    sign_mat = np.array([[-1, -1], [+1, -1], [+1, +1], [-1, +1]])
    xy_size = np.array([hf / float(n_col), vf / float(n_row)]) / 2.0
    rectangles = [sign_mat * xy_size + c for c in centres]

    shift = np.dot(wrap_mat, [1, 1, 0]) - ref[1]

    polys = []
    for r in rectangles:
        r3 = np.append(r, np.zeros((4, 1)), axis=1)
        mapped = np.dot(wrap_mat, r3.T).T - shift
        polys.append(mapped.astype(np.int32))
    return polys