import cv2
import numpy as np

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

# Kalman and KLT trackers
# =========================
# KALMAN 8D
# =========================

class Kalman8D:
    def __init__(self):
        self.x = np.zeros((8, 1), dtype=np.float32)
        self.P = np.eye(8, dtype=np.float32) * 50.0

        self.H = np.zeros((4, 8), dtype=np.float32)
        self.H[0, 0] = 1.0
        self.H[1, 1] = 1.0
        self.H[2, 2] = 1.0
        self.H[3, 3] = 1.0

        self.R = np.eye(4, dtype=np.float32)
        # Снижено с 16→9 и 25→16: новая YOLO даёт точные детекты,
        # Kalman должен им больше доверять и быстрее подстраиваться.
        self.R[0, 0] = 9.0
        self.R[1, 1] = 9.0
        self.R[2, 2] = 16.0
        self.R[3, 3] = 16.0

        self.initialized = False
        # Процесс-шум повышен для реактивности на резкие манёвры.
        # Было: q_pos=2.0, q_size=3.0, q_vel=4.0 — для плавного tracking'а
        # Стало: q_pos=5.0, q_size=3.5, q_vel=12.0 — цель до 160 px/s со швырками
        self.q_pos = 5.0
        self.q_size = 3.5
        self.q_vel = 12.0

    def F(self, dt):
        F = np.eye(8, dtype=np.float32)
        F[0, 4] = dt
        F[1, 5] = dt
        F[2, 6] = dt
        F[3, 7] = dt
        return F

    def Q(self, dt, q_scale=1.0):
        q_scale = float(max(0.25, q_scale))
        q = np.zeros((8, 8), dtype=np.float32)
        s1 = dt * dt
        s2 = dt
        q[0, 0] = self.q_pos * s1 * q_scale
        q[1, 1] = self.q_pos * s1 * q_scale
        q[2, 2] = self.q_size * s1 * q_scale
        q[3, 3] = self.q_size * s1 * q_scale
        q[4, 4] = self.q_vel * s2 * q_scale
        q[5, 5] = self.q_vel * s2 * q_scale
        q[6, 6] = self.q_vel * s2 * q_scale
        q[7, 7] = self.q_vel * s2 * q_scale
        return q

    def predict(self, dt, q_scale=1.0):
        dt = float(max(1e-3, dt))
        F = self.F(dt)
        self.x = F @ self.x
        self.P = F @ self.P @ F.T + self.Q(dt, q_scale=q_scale)
        return self.x[:4].flatten()

    def update(self, z):
        z = np.array(z, dtype=np.float32).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        self.x = self.x + K @ y
        I = np.eye(8, dtype=np.float32)
        self.P = (I - K @ self.H) @ self.P

    def init_from_box(self, box):
        cx, cy = box_center(box)
        w, h = box_wh(box)
        self.x[:] = 0
        self.x[0, 0] = cx
        self.x[1, 0] = cy
        self.x[2, 0] = w
        self.x[3, 0] = h
        self.P = np.eye(8, dtype=np.float32) * 50.0
        self.initialized = True

    def to_box(self):
        cx = float(self.x[0, 0])
        cy = float(self.x[1, 0])
        w = float(max(2.0, self.x[2, 0]))
        h = float(max(2.0, self.x[3, 0]))
        return np.array([cx - w * 0.5, cy - h * 0.5, cx + w * 0.5, cy + h * 0.5], dtype=np.float32)

    def uncertainty(self):
        return float(self.P[0, 0] + self.P[1, 1] + self.P[2, 2] + self.P[3, 3])


# =========================
# KLT TRACKER
# =========================

class KLTTracker:
    def __init__(self):
        self.prev_gray = None
        self.pts = None
        self.box = None
        self.good_count = 0
        self.quality = 0.0

    def reset(self):
        self.prev_gray = None
        self.pts = None
        self.box = None
        self.good_count = 0
        self.quality = 0.0

    def init(self, frame_bgr, box):
        gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
        x1, y1, x2, y2 = map(int, box)
        mask = np.zeros_like(gray)
        cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)

        pts = cv2.goodFeaturesToTrack(
            gray,
            maxCorners=KLT_MAX_CORNERS,
            qualityLevel=KLT_QUALITY,
            minDistance=KLT_MIN_DIST,
            mask=mask
        )
        self.prev_gray = gray
        self.pts = pts
        self.box = box.copy()
        self.good_count = 0
        self.quality = 0.0

    def update(self, frame_bgr):
        if self.prev_gray is None or self.pts is None or len(self.pts) == 0 or self.box is None:
            self.quality = 0.0
            self.good_count = 0
            return None

        gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
        next_pts, st, err = cv2.calcOpticalFlowPyrLK(
            self.prev_gray, gray, self.pts, None,
            winSize=KLT_WIN,
            maxLevel=KLT_MAX_LEVEL,
            criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01)
        )

        if next_pts is None or st is None:
            self.quality = 0.0
            self.good_count = 0
            self.prev_gray = gray
            return None

        st = st.reshape(-1)
        good_old = self.pts[st == 1].reshape(-1, 2)
        good_new = next_pts[st == 1].reshape(-1, 2)

        self.good_count = int(len(good_new))
        if self.good_count < 6:
            self.quality = 0.0
            self.prev_gray = gray
            self.pts = good_new.reshape(-1, 1, 2) if self.good_count > 0 else None
            return None

        disp = good_new - good_old
        med = np.median(disp, axis=0)
        d = np.linalg.norm(disp - med[None, :], axis=1)
        inliers = d < KLT_OUTLIER_THRESH
        in_cnt = int(np.count_nonzero(inliers))

        if in_cnt < 6:
            self.quality = 0.0
            self.prev_gray = gray
            self.pts = good_new.reshape(-1, 1, 2)
            return None

        disp_in = disp[inliers]
        med = np.median(disp_in, axis=0)
        dx, dy = float(med[0]), float(med[1])

        x1, y1, x2, y2 = self.box
        new_box = np.array([x1 + dx, y1 + dy, x2 + dx, y2 + dy], dtype=np.float32)

        frac = float(in_cnt) / float(len(inliers))
        self.quality = clamp(frac * (self.good_count / max(1, KLT_MAX_CORNERS)), 0.0, 1.0)

        self.prev_gray = gray
        self.pts = good_new[inliers].reshape(-1, 1, 2)
        self.box = new_box
        return new_box


# =========================