import math
import signal
import sys
import time
from dataclasses import dataclass
from typing import Tuple

import cv2
import numpy as np

import motor


FRAME_WIDTH = 320
FRAME_HEIGHT = 240
GUIDE_CENTER_X = 160
GUIDE_CENTER_Y = 120
CONTROL_DEADZONE = 15

MIN_CONTOUR_AREA = 90
MAX_CONTOUR_AREA = 6000
EDGE_IGNORE = 12
TARGET_LOCK_DURATION = 3.0


@dataclass
class TrackState:
    cx: int
    cy: int
    area: float
    bbox: Tuple[int, int, int, int]
    vx: float = 0.0
    vy: float = 0.0
    misses: int = 0
    visible: bool = True


def signal_handler(sig, frame):
    print("\nStop signal received, stopping motors...")
    motor.stop()
    print("Motors stopped, exiting.")
    sys.exit(0)


signal.signal(signal.SIGINT, signal_handler)


def estimate_camera_motion(prev_gray, gray):
    identity = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)

    prev_pts = cv2.goodFeaturesToTrack(
        prev_gray,
        maxCorners=180,
        qualityLevel=0.01,
        minDistance=8,
        blockSize=7,
    )
    if prev_pts is None or len(prev_pts) < 8:
        return identity, 0.0, 0

    curr_pts, status, _ = cv2.calcOpticalFlowPyrLK(
        prev_gray,
        gray,
        prev_pts,
        None,
        winSize=(21, 21),
        maxLevel=3,
        criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01),
    )
    if curr_pts is None or status is None:
        return identity, 0.0, 0

    valid = status.reshape(-1) == 1
    prev_valid = prev_pts[valid].reshape(-1, 2)
    curr_valid = curr_pts[valid].reshape(-1, 2)
    if len(prev_valid) < 8:
        return identity, 0.0, len(prev_valid)

    transform, inliers = cv2.estimateAffinePartial2D(
        prev_valid,
        curr_valid,
        method=cv2.RANSAC,
        ransacReprojThreshold=3.0,
        maxIters=2000,
        confidence=0.99,
    )
    if transform is None:
        return identity, 0.0, len(prev_valid)

    tx = float(transform[0, 2])
    ty = float(transform[1, 2])
    rotation = math.degrees(math.atan2(transform[1, 0], transform[0, 0]))
    motion_level = math.hypot(tx, ty) + abs(rotation) * 1.5

    inlier_count = int(inliers.sum()) if inliers is not None else len(prev_valid)
    return transform.astype(np.float32), motion_level, inlier_count


def build_motion_mask(prev_gray, gray, transform, motion_level):
    height, width = gray.shape

    aligned_prev = cv2.warpAffine(
        prev_gray,
        transform,
        (width, height),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_REPLICATE,
    )

    valid_region = cv2.warpAffine(
        np.full_like(prev_gray, 255),
        transform,
        (width, height),
        flags=cv2.INTER_NEAREST,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )

    diff = cv2.absdiff(aligned_prev, gray)
    diff = cv2.GaussianBlur(diff, (5, 5), 0)

    threshold_value = int(min(40, 18 + motion_level * 1.5))
    _, mask = cv2.threshold(diff, threshold_value, 255, cv2.THRESH_BINARY)

    if EDGE_IGNORE > 0:
        mask[:EDGE_IGNORE, :] = 0
        mask[-EDGE_IGNORE:, :] = 0
        mask[:, :EDGE_IGNORE] = 0
        mask[:, -EDGE_IGNORE:] = 0

    mask = cv2.bitwise_and(mask, valid_region)

    kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    kernel_large = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_small)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_large)
    mask = cv2.dilate(mask, kernel_small, iterations=1)

    return mask


def contour_candidates(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    candidates = []
    height, width = mask.shape

    for contour in contours:
        area = cv2.contourArea(contour)
        if area < MIN_CONTOUR_AREA or area > MAX_CONTOUR_AREA:
            continue

        x, y, w, h = cv2.boundingRect(contour)
        if x <= EDGE_IGNORE or y <= EDGE_IGNORE:
            continue
        if x + w >= width - EDGE_IGNORE or y + h >= height - EDGE_IGNORE:
            continue

        aspect_ratio = float(w) / h if h else 0.0
        if aspect_ratio < 0.25 or aspect_ratio > 4.0:
            continue

        rect_area = float(w * h)
        fill_ratio = area / rect_area if rect_area else 0.0
        if fill_ratio < 0.18:
            continue

        cx = x + w // 2
        cy = y + h // 2
        candidates.append(
            {
                "contour": contour,
                "bbox": (x, y, w, h),
                "area": float(area),
                "cx": cx,
                "cy": cy,
            }
        )

    return candidates


def bbox_iou(box_a, box_b):
    ax, ay, aw, ah = box_a
    bx, by, bw, bh = box_b

    left = max(ax, bx)
    top = max(ay, by)
    right = min(ax + aw, bx + bw)
    bottom = min(ay + ah, by + bh)

    if right <= left or bottom <= top:
        return 0.0

    intersection = (right - left) * (bottom - top)
    union = aw * ah + bw * bh - intersection
    return intersection / union if union > 0 else 0.0


def select_target(candidates, track_state):
    if not candidates:
        return None

    if track_state is None:
        return max(
            candidates,
            key=lambda c: c["area"] - 0.6 * abs(c["cx"] - GUIDE_CENTER_X),
        )

    predicted_x = track_state.cx + track_state.vx
    predicted_y = track_state.cy + track_state.vy
    best_candidate = None
    best_score = -1e9

    base_radius = max(55.0, math.sqrt(max(track_state.area, 1.0)) * 2.5)
    search_radius = base_radius + min(track_state.misses * 18.0, 60.0)

    for candidate in candidates:
        distance = math.hypot(candidate["cx"] - predicted_x, candidate["cy"] - predicted_y)
        overlap = bbox_iou(candidate["bbox"], track_state.bbox)
        area_ratio = candidate["area"] / max(track_state.area, 1.0)
        area_penalty = abs(math.log(max(area_ratio, 1e-6)))

        if distance > search_radius and overlap < 0.02:
            continue

        score = (
            candidate["area"] * 0.08
            + overlap * 60.0
            - distance * 1.4
            - area_penalty * 22.0
        )

        if score > best_score:
            best_score = score
            best_candidate = candidate

    if best_candidate is not None:
        return best_candidate

    return None


def update_track_state(track_state, candidate):
    if candidate is None:
        return None

    if track_state is None:
        return TrackState(
            cx=candidate["cx"],
            cy=candidate["cy"],
            area=candidate["area"],
            bbox=candidate["bbox"],
        )

    if track_state.misses > 0:
        smooth_prev = 0.35
    else:
        smooth_prev = 0.60

    new_cx = int(smooth_prev * track_state.cx + (1.0 - smooth_prev) * candidate["cx"])
    new_cy = int(smooth_prev * track_state.cy + (1.0 - smooth_prev) * candidate["cy"])

    raw_vx = new_cx - track_state.cx
    raw_vy = new_cy - track_state.cy
    new_vx = 0.6 * track_state.vx + 0.4 * raw_vx
    new_vy = 0.6 * track_state.vy + 0.4 * raw_vy

    track_state.cx = new_cx
    track_state.cy = new_cy
    track_state.area = candidate["area"]
    track_state.bbox = candidate["bbox"]
    track_state.vx = new_vx
    track_state.vy = new_vy
    track_state.misses = 0
    track_state.visible = True
    return track_state


def draw_debug(frame, mask, track_state, motion_level, inlier_count, lock_remaining):
    frame_width = frame.shape[1]
    cv2.circle(frame, (GUIDE_CENTER_X, GUIDE_CENTER_Y), 5, (255, 0, 0), -1)

    cv2.putText(
        frame,
        f"cam:{motion_level:4.1f} feat:{inlier_count:3d}",
        (8, 18),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.45,
        (0, 255, 255),
        1,
        cv2.LINE_AA,
    )

    if lock_remaining > 0:
        cv2.putText(
            frame,
            f"LOCK {lock_remaining:0.1f}s",
            (8, 38),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.55,
            (0, 0, 255),
            2,
            cv2.LINE_AA,
        )

    if track_state is not None:
        x, y, w, h = track_state.bbox
        color = (0, 255, 0) if track_state.visible else (0, 165, 255)
        cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
        cv2.circle(frame, (track_state.cx, track_state.cy), 5, (0, 0, 255), -1)

        status = "track" if track_state.visible else f"hold:{track_state.misses}"
        cv2.putText(
            frame,
            status,
            (x, max(15, y - 8)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.45,
            color,
            1,
            cv2.LINE_AA,
        )

    mask_preview = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
    preview_h, preview_w = 60, 80
    mask_preview = cv2.resize(mask_preview, (preview_w, preview_h))
    frame[0:preview_h, frame_width - preview_w : frame_width] = mask_preview


def drive_to_target(track_state, lock_active=False):
    if lock_active or track_state is None or not track_state.visible:
        motor.stop()
        return

    dx = track_state.cx - GUIDE_CENTER_X
    dy = track_state.cy - GUIDE_CENTER_Y

    if abs(dx) < CONTROL_DEADZONE and abs(dy) < CONTROL_DEADZONE:
        motor.stop()
    elif abs(dx) < CONTROL_DEADZONE:
        if dy > CONTROL_DEADZONE:
            motor.move_right(speed=1.0)
        else:
            motor.move_left(speed=1.0)
    elif abs(dy) < CONTROL_DEADZONE:
        if dx > CONTROL_DEADZONE:
            motor.forward(speed=1.0)
        else:
            motor.backward(speed=1.0)
    else:
        if dx > CONTROL_DEADZONE and dy > CONTROL_DEADZONE:
            motor.move_right_backward(speed=1.0)
        elif dx > CONTROL_DEADZONE and dy < -CONTROL_DEADZONE:
            motor.move_left_backward(speed=1.0)
        elif dx < -CONTROL_DEADZONE and dy > CONTROL_DEADZONE:
            motor.move_right_forward(speed=1.0)
        else:
            motor.move_left_forward(speed=1.0)


print("Starting camera tracking...")
print(f"OpenCV version: {cv2.__version__}")

cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("Unable to open camera.")
    sys.exit(1)

cap.set(cv2.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

print(f"Camera ready: {cap.get(cv2.CAP_PROP_FRAME_WIDTH)}x{cap.get(cv2.CAP_PROP_FRAME_HEIGHT)}")

success_count = 0
for _ in range(10):
    ret, _ = cap.read()
    if ret:
        success_count += 1
    time.sleep(0.05)
print(f"Warm-up finished, successful reads: {success_count}/10")

ret, prev_frame = cap.read()
if not ret:
    print("Unable to read initial frame.")
    cap.release()
    sys.exit(1)

prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
prev_gray = cv2.GaussianBlur(prev_gray, (5, 5), 0)
track_state = None
lock_until = 0.0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)

    transform, motion_level, inlier_count = estimate_camera_motion(prev_gray, gray)
    motion_mask = build_motion_mask(prev_gray, gray, transform, motion_level)
    candidates = contour_candidates(motion_mask)

    now = time.monotonic()
    lock_remaining = max(0.0, lock_until - now)
    lock_active = lock_remaining > 0.0

    if lock_active:
        track_state = None
        selected = None
    else:
        selected = select_target(candidates, track_state)
        if track_state is not None and selected is None:
            motor.stop()
            track_state = None
            lock_until = now + TARGET_LOCK_DURATION
            lock_remaining = TARGET_LOCK_DURATION
            lock_active = True
        else:
            track_state = update_track_state(track_state, selected)

    draw_debug(frame, motion_mask, track_state, motion_level, inlier_count, lock_remaining)
    drive_to_target(track_state, lock_active=lock_active)

    if lock_active:
        print(f"track locked {lock_remaining:0.1f}s, cam={motion_level:4.1f}")
    elif track_state is not None:
        print(
            f"track=({track_state.cx:3d},{track_state.cy:3d}) "
            f"area={track_state.area:6.1f} cam={motion_level:4.1f}"
        )
    else:
        print(f"waiting target, cam={motion_level:4.1f}")

    cv2.imshow("tracking", frame)
    prev_gray = gray.copy()

    if cv2.waitKey(1) & 0xFF == 27:
        break

cap.release()
cv2.destroyAllWindows()
motor.stop()
