import cv2
import numpy as np

from config import *
from helpers import clamp

# Camera motion compensation
# =========================
# Поддерживает три режима оценки ego-motion:
#   - affine  (2x3, 4 DoF через estimateAffinePartial2D)
#   - homography (3x3, 8 DoF через findHomography)
#   - adaptive: пытаемся homography, при плохом качестве падаем на affine
#
# Основные функции:
#   estimate_global_affine_ex     — старая добрая affine + outliers
#   estimate_global_homography_ex — homography + outliers + meta качества
#   estimate_ego_motion_adaptive  — автоматический выбор
# =========================

# Пороги качества homography
HOMO_MIN_INLIERS       = 50
HOMO_MIN_FILLED_CELLS  = 4     # из 3x3 = 9 ячеек
HOMO_GRID_SIZE         = 3
HOMO_MAX_CORNER_SHIFT  = 120.0 # px — защита от вырожденных H


def _rect_mask_exclude(gray, rects_norm):
    h, w = gray.shape[:2]
    mask = np.ones((h, w), dtype=np.uint8) * 255
    for (x1n, y1n, x2n, y2n) in rects_norm:
        x1 = int(clamp(x1n * w, 0, w - 1))
        y1 = int(clamp(y1n * h, 0, h - 1))
        x2 = int(clamp(x2n * w, 0, w))
        y2 = int(clamp(y2n * h, 0, h))
        cv2.rectangle(mask, (x1, y1), (x2, y2), 0, -1)
    return mask


def _lk_correspondences(prev_gray, gray):
    """Общий блок: feature detection + LK tracking. (p0, p1) или (None, None)."""
    if prev_gray is None or gray is None:
        return None, None

    mask = _rect_mask_exclude(prev_gray, CAM_MOTION_EXCLUDE_NORM) if CAM_MOTION_EXCLUDE_NORM else None
    pts = cv2.goodFeaturesToTrack(
        prev_gray,
        maxCorners=int(CAM_MOTION_MAX_CORNERS),
        qualityLevel=float(CAM_MOTION_QL),
        minDistance=int(CAM_MOTION_MIN_DIST),
        mask=mask,
    )
    if pts is None or len(pts) < 25:
        return None, None

    nxt, st, err = cv2.calcOpticalFlowPyrLK(
        prev_gray, gray, pts, None, winSize=(21, 21), maxLevel=3
    )
    if nxt is None or st is None:
        return None, None

    st = st.reshape(-1).astype(bool)
    p0 = pts[st].reshape(-1, 2).astype(np.float32)
    p1 = nxt[st].reshape(-1, 2).astype(np.float32)
    if len(p0) < 25:
        return None, None

    return p0, p1


def _classify_outliers(p0, p1, inlier_mask):
    """
    Разбивает точки текущего кадра (p1) на inliers/outliers.
    Отфильтровывает outliers с слишком малым/большим потоком (LK шум).
    """
    inlier_mask = inlier_mask.reshape(-1).astype(bool)
    inliers_curr = p1[inlier_mask]
    outliers_curr = p1[~inlier_mask]

    if len(outliers_curr) == 0:
        return inliers_curr, outliers_curr

    flow = p1[~inlier_mask] - p0[~inlier_mask]
    mag = np.linalg.norm(flow, axis=1)

    if len(inliers_curr) > 0:
        ego_flow = p1[inlier_mask] - p0[inlier_mask]
        ego_med = float(np.linalg.norm(np.median(ego_flow, axis=0)))
    else:
        ego_med = 0.0

    valid = (mag > max(1.5, ego_med * 1.8)) & (mag < 80.0)
    outliers_curr = outliers_curr[valid]
    return inliers_curr, outliers_curr


# ─────────────────────────────────────────────────────────────
# AFFINE
# ─────────────────────────────────────────────────────────────

def estimate_global_affine_ex(prev_gray, gray):
    empty = (None, np.zeros((0, 2), dtype=np.float32), np.zeros((0, 2), dtype=np.float32))

    p0, p1 = _lk_correspondences(prev_gray, gray)
    if p0 is None:
        return empty

    A, inlier_mask = cv2.estimateAffinePartial2D(
        p0, p1,
        method=cv2.RANSAC,
        ransacReprojThreshold=float(CAM_MOTION_RANSAC_THR),
    )
    if A is None or inlier_mask is None:
        return (A, np.zeros((0, 2), dtype=np.float32), p1)

    inliers_curr, outliers_curr = _classify_outliers(p0, p1, inlier_mask)
    return A, inliers_curr, outliers_curr


def estimate_global_affine(prev_gray, gray):
    """Backward-compatible wrapper."""
    A, _, _ = estimate_global_affine_ex(prev_gray, gray)
    return A


# ─────────────────────────────────────────────────────────────
# HOMOGRAPHY
# ─────────────────────────────────────────────────────────────

def estimate_global_homography_ex(prev_gray, gray):
    """
    Homography 3x3. Возвращает (H, inliers_curr, outliers_curr, meta).
    meta: n_inliers, filled_cells, spatial_ok, corner_shift_ok
    """
    empty_meta = {
        "n_inliers": 0,
        "filled_cells": 0,
        "spatial_ok": False,
        "corner_shift_ok": False,
    }
    empty = (
        None,
        np.zeros((0, 2), dtype=np.float32),
        np.zeros((0, 2), dtype=np.float32),
        empty_meta,
    )

    p0, p1 = _lk_correspondences(prev_gray, gray)
    if p0 is None:
        return empty

    H, inlier_mask = cv2.findHomography(
        p0, p1,
        method=cv2.RANSAC,
        ransacReprojThreshold=float(CAM_MOTION_RANSAC_THR),
        maxIters=2000,
        confidence=0.995,
    )
    if H is None or inlier_mask is None:
        return empty

    inlier_mask_flat = inlier_mask.reshape(-1).astype(bool)
    n_inliers = int(inlier_mask_flat.sum())

    # Распределение inliers по 3x3 сетке
    h, w = gray.shape[:2]
    gs = int(HOMO_GRID_SIZE)
    cells = np.zeros((gs, gs), dtype=np.uint8)
    if n_inliers > 0:
        inl = p1[inlier_mask_flat]
        gx = np.clip((inl[:, 0] / w * gs).astype(int), 0, gs - 1)
        gy = np.clip((inl[:, 1] / h * gs).astype(int), 0, gs - 1)
        for x, y in zip(gx, gy):
            cells[y, x] = 1
    filled_cells = int(cells.sum())
    spatial_ok = filled_cells >= int(HOMO_MIN_FILLED_CELLS)

    # Проверка вырождения: сдвиг углов кадра
    corners = np.array(
        [[0, 0], [w - 1, 0], [w - 1, h - 1], [0, h - 1]], dtype=np.float32
    ).reshape(-1, 1, 2)
    try:
        warped_corners = cv2.perspectiveTransform(corners, H).reshape(-1, 2)
        shifts = np.linalg.norm(warped_corners - corners.reshape(-1, 2), axis=1)
        corner_shift_ok = bool(np.max(shifts) < float(HOMO_MAX_CORNER_SHIFT))
    except cv2.error:
        corner_shift_ok = False

    meta = {
        "n_inliers": n_inliers,
        "filled_cells": filled_cells,
        "spatial_ok": spatial_ok,
        "corner_shift_ok": corner_shift_ok,
    }

    inliers_curr, outliers_curr = _classify_outliers(p0, p1, inlier_mask_flat)
    return H, inliers_curr, outliers_curr, meta


def homography_is_plausible(H, meta):
    """Проходит ли homography пороги качества."""
    if H is None:
        return False
    if meta.get("n_inliers", 0) < int(HOMO_MIN_INLIERS):
        return False
    if not meta.get("spatial_ok", False):
        return False
    if not meta.get("corner_shift_ok", False):
        return False
    try:
        det = float(np.linalg.det(H[:2, :2]))
    except Exception:
        return False
    if not (0.3 < abs(det) < 3.0):
        return False
    return True


# ─────────────────────────────────────────────────────────────
# ADAPTIVE SELECTION
# ─────────────────────────────────────────────────────────────

def estimate_ego_motion_adaptive(prev_gray, gray):
    """
    Сначала homography, при плохом качестве — affine.

    Returns:
        M        — 3x3 (homography) / 2x3 (affine) / None
        kind     — "homography" | "affine" | "none"
        inliers  — Nx2 точки статичного мира
        outliers — Nx2 точки независимо движущихся объектов
    """
    H, h_in, h_out, meta = estimate_global_homography_ex(prev_gray, gray)
    if homography_is_plausible(H, meta):
        return H, "homography", h_in, h_out

    A, a_in, a_out = estimate_global_affine_ex(prev_gray, gray)
    if A is not None and affine_is_plausible(A):
        return A, "affine", a_in, a_out

    empty = np.zeros((0, 2), dtype=np.float32)
    return None, "none", empty, empty


# ─────────────────────────────────────────────────────────────
# Utilities / backward compat
# ─────────────────────────────────────────────────────────────

def apply_affine_to_point(A, x, y):
    return (
        A[0, 0] * x + A[0, 1] * y + A[0, 2],
        A[1, 0] * x + A[1, 1] * y + A[1, 2],
    )


def affine_is_plausible(A):
    if A is None:
        return False
    tx = float(A[0, 2])
    ty = float(A[1, 2])
    shift = float(np.hypot(tx, ty))
    if shift > float(CAM_MOTION_MAX_SHIFT_PX):
        return False

    a00 = float(A[0, 0])
    a10 = float(A[1, 0])
    col0_norm = float(np.hypot(a00, a10))
    if col0_norm < 1e-6:
        return False
    if not (float(CAM_MOTION_SCALE_MIN) <= col0_norm <= float(CAM_MOTION_SCALE_MAX)):
        return False

    ang = float(np.degrees(np.arctan2(a10, a00)))
    if abs(ang) > float(CAM_MOTION_MAX_ROT_DEG):
        return False

    return True


def warp_is_plausible(M, kind):
    if kind == "homography":
        return M is not None
    if kind == "affine":
        return affine_is_plausible(M)
    return False