from dataclasses import dataclass

import cv2
import numpy as np

from config import *
from helpers import box_area, box_center, box_wh, clamp, clip_box, iou


@dataclass(frozen=True)
class MotionGroupEvidence:
    reliable: bool = False
    valid: bool = False
    point_count: int = 0
    coherent_count: int = 0
    coherence: float = 0.0
    residual_px: float = 0.0
    residual_x: float = 0.0
    residual_y: float = 0.0
    raw_motion_px: float = 0.0
    speed_norm_s: float = 0.0
    scale_ratio: float = 1.0
    spread_ratio: float = 0.0
    edge_violation: bool = False
    speed_violation: bool = False
    screen_static: bool = False
    score: float = 0.0


def _transform_points(points, affine):
    if affine is None:
        return points.copy()
    linear = np.asarray(affine[:, :2], dtype=np.float32)
    offset = np.asarray(affine[:, 2], dtype=np.float32)
    return points @ linear.T + offset


def analyze_motion_group(prev_gray, gray, box, affine=None, dt=0.04):
    if prev_gray is None or gray is None or prev_gray.shape != gray.shape:
        return MotionGroupEvidence()

    frame_h, frame_w = gray.shape[:2]
    b = clip_box(box, frame_w, frame_h)
    bw, bh = box_wh(b)
    pad_x = max(float(PHYSICS_BOX_PAD_MIN), float(bw) * float(PHYSICS_BOX_PAD_RATIO))
    pad_y = max(float(PHYSICS_BOX_PAD_MIN), float(bh) * float(PHYSICS_BOX_PAD_RATIO))
    sample_box = clip_box(
        [b[0] - pad_x, b[1] - pad_y, b[2] + pad_x, b[3] + pad_y],
        frame_w,
        frame_h,
    )

    mask = np.zeros_like(prev_gray, dtype=np.uint8)
    x1, y1, x2, y2 = map(int, sample_box)
    mask[y1:y2, x1:x2] = 255
    points = cv2.goodFeaturesToTrack(
        prev_gray,
        maxCorners=int(PHYSICS_MAX_POINTS),
        qualityLevel=float(PHYSICS_QUALITY_LEVEL),
        minDistance=float(PHYSICS_MIN_POINT_DISTANCE),
        mask=mask,
        blockSize=3,
    )
    if points is None or len(points) < int(PHYSICS_MIN_POINTS):
        return MotionGroupEvidence(point_count=0 if points is None else int(len(points)))

    next_points, status, errors = cv2.calcOpticalFlowPyrLK(
        prev_gray,
        gray,
        points,
        None,
        winSize=(21, 21),
        maxLevel=3,
        criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01),
    )
    if next_points is None or status is None:
        return MotionGroupEvidence()

    good = status.reshape(-1).astype(bool)
    if errors is not None:
        good &= errors.reshape(-1) <= float(PHYSICS_MAX_LK_ERROR)
    old = points.reshape(-1, 2)[good].astype(np.float32)
    new = next_points.reshape(-1, 2)[good].astype(np.float32)
    inside = (
        (new[:, 0] >= 0)
        & (new[:, 0] < frame_w)
        & (new[:, 1] >= 0)
        & (new[:, 1] < frame_h)
    )
    old = old[inside]
    new = new[inside]
    point_count = int(len(new))
    if point_count < int(PHYSICS_MIN_POINTS):
        return MotionGroupEvidence(point_count=point_count)

    ego_pred = _transform_points(old, affine)
    residuals = new - ego_pred
    raw_motion_px = float(np.linalg.norm(np.median(new - old, axis=0)))
    median_residual = np.median(residuals, axis=0)
    residual_px = float(np.linalg.norm(median_residual))
    deviations = np.linalg.norm(residuals - median_residual[None, :], axis=1)
    coherent_limit = max(
        float(PHYSICS_COHERENT_MIN_PX),
        float(PHYSICS_COHERENT_RESIDUAL_FACTOR) * residual_px,
    )
    coherent = deviations <= coherent_limit
    coherent_count = int(np.count_nonzero(coherent))
    coherence = float(coherent_count / max(1, point_count))

    coherent_new = new[coherent]
    coherent_pred = ego_pred[coherent]
    spread_ratio = 0.0
    scale_ratio = 1.0
    if coherent_count >= 3:
        span = np.ptp(coherent_new, axis=0)
        spread_ratio = float(
            np.linalg.norm(span) / max(1.0, np.linalg.norm([float(bw), float(bh)]))
        )
        pred_center = np.median(coherent_pred, axis=0)
        new_center = np.median(coherent_new, axis=0)
        old_radius = np.linalg.norm(coherent_pred - pred_center[None, :], axis=1)
        new_radius = np.linalg.norm(coherent_new - new_center[None, :], axis=1)
        usable = old_radius >= 1.0
        if np.count_nonzero(usable) >= 3:
            scale_ratio = float(np.median(new_radius[usable] / old_radius[usable]))

    frame_diag = max(1.0, float(np.hypot(frame_w, frame_h)))
    speed_norm_s = residual_px / max(1e-3, float(dt)) / frame_diag
    area_ratio = float(box_area(b) / max(1.0, float(frame_w * frame_h)))
    near_factor = float(clamp(np.sqrt(area_ratio / max(1e-6, PHYSICS_NEAR_AREA_RATIO)), 0.0, 1.0))
    max_speed = (
        (1.0 - near_factor) * float(PHYSICS_FAR_MAX_SPEED_NORM_S)
        + near_factor * float(PHYSICS_NEAR_MAX_SPEED_NORM_S)
    )
    speed_valid = bool(
        speed_norm_s <= max_speed
        or scale_ratio >= float(PHYSICS_GROWTH_FULL)
    )
    screen_static = bool(
        raw_motion_px < float(PHYSICS_MIN_RAW_MOTION_PX)
        and abs(scale_ratio - 1.0) < float(PHYSICS_SCREEN_STATIC_SCALE_EPS)
    )

    edge_margin = max(
        float(PHYSICS_EDGE_MARGIN_MIN),
        float(PHYSICS_EDGE_MARGIN_RATIO) * min(frame_w, frame_h),
    )
    moving_outward = bool(
        (b[0] <= edge_margin and median_residual[0] < -float(PHYSICS_MIN_RESIDUAL_PX))
        or (b[1] <= edge_margin and median_residual[1] < -float(PHYSICS_MIN_RESIDUAL_PX))
        or (b[2] >= (frame_w - edge_margin) and median_residual[0] > float(PHYSICS_MIN_RESIDUAL_PX))
        or (b[3] >= (frame_h - edge_margin) and median_residual[1] > float(PHYSICS_MIN_RESIDUAL_PX))
    )
    edge_violation = bool(
        area_ratio <= float(PHYSICS_DISTANT_AREA_RATIO)
        and moving_outward
    )

    support = float(clamp(coherent_count / max(1.0, PHYSICS_FULL_SUPPORT_POINTS), 0.0, 1.0))
    independence = float(
        clamp(
            (residual_px - float(PHYSICS_MIN_RESIDUAL_PX))
            / max(1e-6, float(PHYSICS_FULL_RESIDUAL_PX) - float(PHYSICS_MIN_RESIDUAL_PX)),
            0.0,
            1.0,
        )
    )
    growth = float(
        clamp(
            (scale_ratio - float(PHYSICS_GROWTH_START))
            / max(1e-6, float(PHYSICS_GROWTH_FULL) - float(PHYSICS_GROWTH_START)),
            0.0,
            1.0,
        )
    )
    spread = float(clamp(spread_ratio / max(1e-6, PHYSICS_FULL_SPREAD_RATIO), 0.0, 1.0))
    reliable = bool(
        point_count >= int(PHYSICS_MIN_POINTS)
        and coherent_count >= int(PHYSICS_MIN_COHERENT_POINTS)
        and coherence >= float(PHYSICS_MIN_COHERENCE)
    )
    valid = bool(
        reliable
        and speed_valid
        and not edge_violation
        and not screen_static
        and (
            residual_px >= float(PHYSICS_MIN_RESIDUAL_PX)
            or scale_ratio >= float(PHYSICS_MIN_GROWTH_RATIO)
        )
    )
    score = (
        0.30 * coherence
        + 0.20 * support
        + 0.25 * independence
        + 0.15 * spread
        + 0.10 * growth
    )
    if not speed_valid:
        score -= 0.35
    if edge_violation:
        score -= 0.50
    if screen_static:
        score -= 0.35

    return MotionGroupEvidence(
        reliable=reliable,
        valid=valid,
        point_count=point_count,
        coherent_count=coherent_count,
        coherence=coherence,
        residual_px=residual_px,
        residual_x=float(median_residual[0]),
        residual_y=float(median_residual[1]),
        raw_motion_px=raw_motion_px,
        speed_norm_s=speed_norm_s,
        scale_ratio=scale_ratio,
        spread_ratio=spread_ratio,
        edge_violation=edge_violation,
        speed_violation=not speed_valid,
        screen_static=screen_static,
        score=float(clamp(score, 0.0, 1.0)),
    )


def match_motion_evidence(box, entries):
    if not entries:
        return None
    b = np.asarray(box, dtype=np.float32)
    center = box_center(b)
    diag = max(1.0, float(np.linalg.norm(box_wh(b))))
    best = None
    best_score = -1.0
    for candidate_box, evidence in entries:
        overlap = float(iou(b, candidate_box))
        distance = float(np.linalg.norm(center - box_center(candidate_box)))
        if overlap <= 0.0 and distance > max(12.0, 0.75 * diag):
            continue
        match_score = overlap + 1.0 / (1.0 + distance)
        if match_score > best_score:
            best_score = match_score
            best = evidence
    return best