import cv2
import numpy as np
import mediapipe as mp
from scipy.stats import entropy
import warnings

warnings.filterwarnings("ignore")

# Initialize MediaPipe Face Detection (Faster than full Face Mesh for simple bounding boxes)
mp_face_detection = mp.solutions.face_detection


class FaceKalman:
    """Smoothes the bounding box to prevent jitter from corrupting optical flow/motion calculations."""

    def __init__(self):
        self.kf = cv2.KalmanFilter(4, 4)
        self.kf.measurementMatrix = np.array(
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32
        )
        self.kf.transitionMatrix = np.array(
            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32
        )
        self.kf.processNoiseCov = np.eye(4, dtype=np.float32) * 0.01

    def update(self, x, y, w, h):
        measured = np.array(
            [[np.float32(x)], [np.float32(y)], [np.float32(w)], [np.float32(h)]]
        )
        self.kf.correct(measured)
        return self.kf.predict().flatten().astype(int)


class PrismFeatureExtractor:
    def __init__(self):
        # The exact 10 features expected by your trained MLPRegressor
        self.detector = mp_face_detection.FaceDetection(
            model_selection=1, min_detection_confidence=0.5
        )

    def _process_single_frame(self, frame, prev_gray_full, kf):
        """Extracts the 11 raw optical primitives from a single frame."""
        h_f, w_f, _ = frame.shape
        gray_full = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Fast Face Detection
        results = self.detector.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        if not results.detections:
            raw_x, raw_y, raw_w, raw_h = w_f // 4, h_f // 4, w_f // 2, h_f // 2
        else:
            bbox = results.detections[0].location_data.relative_bounding_box
            raw_x, raw_y = int(bbox.xmin * w_f), int(bbox.ymin * h_f)
            raw_w, raw_h = int(bbox.width * w_f), int(bbox.height * h_f)

        x, y, w, h = kf.update(raw_x, raw_y, raw_w, raw_h)
        x, y = max(0, x), max(0, y)
        w, h = max(1, min(w_f - x, w)), max(1, min(h_f - y, h))

        face_roi = frame[y : y + h, x : x + w]
        if face_roi.size == 0:
            return np.zeros(10), gray_full

        face_g = face_roi[:, :, 1]
        face_gray = gray_full[y : y + h, x : x + w]

        # Physics Metrics
        phi = np.mean(face_g) / 255.0
        ycrcb = cv2.cvtColor(face_roi, cv2.COLOR_BGR2YCrCb)
        mask = cv2.inRange(ycrcb, np.array([0, 133, 77]), np.array([255, 173, 127]))

        mu = np.mean(ycrcb[:, :, 0][mask > 0]) / 255.0 if np.sum(mask) > 0 else 0.5
        sigma = (
            np.std(face_g[mask > 0]) / 50.0
            if np.sum(mask) > 0
            else np.std(face_g) / 50.0
        )
        chi = min(cv2.Laplacian(face_gray, cv2.CV_64F).var(), 1000.0) / 1000.0

        hist, _ = np.histogram(face_g.flatten(), bins=20, density=True)
        H = entropy(hist + 1e-10) / 5.0
        clip = (np.sum(face_g > 250) + np.sum(face_g < 5)) / face_g.size

        # Motion Calculation (Optical Flow)
        m_val = 0.0
        if prev_gray_full is not None:
            try:
                prev_face = prev_gray_full[y : y + h, x : x + w]
                if prev_face.shape == face_gray.shape:
                    flow = cv2.calcOpticalFlowFarneback(
                        prev_face, face_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
                    )
                    m_val = np.mean(np.linalg.norm(flow, axis=2))
            except:
                pass

        res = (w * h) / (w_f * h_f)
        f_fft = np.fft.fft2(gray_full)
        mag = 20 * np.log(np.abs(np.fft.fftshift(f_fft)) + 1e-10)
        cy, cx = h_f // 2, w_f // 2
        ghost = np.mean(mag[cy - 10 : cy + 10, cx - 10 : cx + 10]) / (
            np.mean(mag) + 1e-5
        )

        return [
            phi,
            sigma,
            mu,
            chi,
            H,
            clip,
            m_val,
            res,
            np.mean(face_g),
            ghost,
        ], gray_full

    def extract_from_video(self, video_path: str, max_samples=300) -> np.ndarray:
        """Reads video, uniformly samples frames, and returns the final 10 engineered features."""
        kf = FaceKalman()
        accumulators = {i: [] for i in range(10)}
        prev_gray = None

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Could not open video file: {video_path}")

        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        step = max(1, total // max_samples)

        count = processed = 0
        while processed < max_samples:
            ret, frame = cap.read()
            if not ret:
                break

            if count % step == 0:
                stats, prev_gray = self._process_single_frame(frame, prev_gray, kf)
                if np.sum(stats) > 0:
                    for idx, val in enumerate(stats):
                        accumulators[idx].append(val)
                    processed += 1
            count += 1
        cap.release()

        if processed < 10:
            raise ValueError("Video too short or face not reliably detected.")

        # 1. Aggregate into Base Primitives
        phi = np.mean(accumulators[0])
        sigma = np.mean(accumulators[1])
        mu = np.mean(accumulators[2])
        chi = np.mean(accumulators[3])
        motion = min(np.percentile(accumulators[6], 95), 5.0) / 5.0
        ghost = min(np.mean(accumulators[9]), 5.0) / 5.0

        # 2. Compute the 10 Engineered Features (Matching Training Config)
        sigma_safe = sigma + 1e-6
        snr = phi / sigma_safe
        stability = mu / (sigma_safe + 1e-3)
        motion_delta = 0.0  # For a single inference video, cross-video delta is 0
        purity_score = (phi * chi) / (sigma_safe * ghost + 1e-6)

        # Output shape: (1, 10)
        engineered_features = np.array(
            [
                [
                    snr,
                    stability,
                    motion,
                    motion_delta,
                    purity_score,
                    phi,
                    sigma,
                    chi,
                    ghost,
                    np.log1p(np.abs(phi)),
                ]
            ]
        )

        return engineered_features