"""
Court calibration module for mapping pixel coordinates to real-world coordinates
Uses homography transformation based on court keypoints
"""
import numpy as np
import cv2
from typing import List, Tuple, Optional, Dict
import json


class CourtCalibrator:
    """
    Calibrates camera perspective to map pixel coordinates to real-world court coordinates
    Uses homography transformation
    """

    def __init__(self, court_width_m: float = 6.1, court_length_m: float = 13.4):
        """
        Initialize court calibrator

        Args:
            court_width_m: Width of pickleball court in meters (default: 6.1m)
            court_length_m: Length of pickleball court in meters (default: 13.4m)
        """
        self.court_width = court_width_m
        self.court_length = court_length_m
        self.homography_matrix = None
        self.court_corners_pixel = None
        self.court_corners_real = np.array([
            [0, 0],                                    # Top-left
            [court_length_m, 0],                       # Top-right
            [court_length_m, court_width_m],           # Bottom-right
            [0, court_width_m]                         # Bottom-left
        ], dtype=np.float32)

    def calibrate_manual(self, corner_points: List[Tuple[float, float]]) -> bool:
        """
        Manually calibrate using 4 corner points of the court

        Args:
            corner_points: List of 4 (x, y) tuples representing court corners in pixels
                          Order: [top-left, top-right, bottom-right, bottom-left]

        Returns:
            True if calibration successful, False otherwise
        """
        if len(corner_points) != 4:
            print("Error: Need exactly 4 corner points")
            return False

        self.court_corners_pixel = np.array(corner_points, dtype=np.float32)

        # Calculate homography matrix
        self.homography_matrix, status = cv2.findHomography(
            self.court_corners_pixel,
            self.court_corners_real,
            method=cv2.RANSAC
        )

        if self.homography_matrix is None:
            print("Error: Failed to calculate homography matrix")
            return False

        print("✓ Court calibration successful")
        return True

    def calibrate_auto(self, frame: np.ndarray) -> bool:
        """
        Automatically detect court corners using computer vision
        This is a simplified version - can be enhanced with YOLO keypoint detection

        Args:
            frame: Video frame to detect court corners from

        Returns:
            True if calibration successful, False otherwise
        """
        # Convert to grayscale
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Apply edge detection
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)

        # Detect lines using Hough Transform
        lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100,
                               minLineLength=100, maxLineGap=10)

        if lines is None or len(lines) < 4:
            print("Error: Could not detect enough lines for court corners")
            return False

        # This is a placeholder - in production, you'd use:
        # 1. YOLO keypoint detection model trained on court corners
        # 2. More sophisticated line intersection detection
        # 3. Court line template matching

        print("Warning: Auto-calibration not fully implemented")
        print("Please use calibrate_manual() with corner points")
        return False

    def pixel_to_real(self, pixel_coords: Tuple[float, float]) -> Optional[Tuple[float, float]]:
        """
        Transform pixel coordinates to real-world court coordinates

        Args:
            pixel_coords: (x, y) tuple in pixel space

        Returns:
            (x, y) tuple in real-world meters, or None if not calibrated
        """
        if self.homography_matrix is None:
            print("Error: Court not calibrated. Call calibrate_manual() first")
            return None

        # Convert to homogeneous coordinates
        pixel_point = np.array([[pixel_coords[0], pixel_coords[1]]], dtype=np.float32)
        pixel_point = pixel_point.reshape(-1, 1, 2)

        # Apply homography transformation
        real_point = cv2.perspectiveTransform(pixel_point, self.homography_matrix)

        x, y = real_point[0][0]
        return (float(x), float(y))

    def real_to_pixel(self, real_coords: Tuple[float, float]) -> Optional[Tuple[float, float]]:
        """
        Transform real-world court coordinates to pixel coordinates

        Args:
            real_coords: (x, y) tuple in real-world meters

        Returns:
            (x, y) tuple in pixel space, or None if not calibrated
        """
        if self.homography_matrix is None:
            print("Error: Court not calibrated")
            return None

        # Use inverse homography
        inv_homography = np.linalg.inv(self.homography_matrix)

        real_point = np.array([[real_coords[0], real_coords[1]]], dtype=np.float32)
        real_point = real_point.reshape(-1, 1, 2)

        pixel_point = cv2.perspectiveTransform(real_point, inv_homography)

        x, y = pixel_point[0][0]
        return (float(x), float(y))

    def is_calibrated(self) -> bool:
        """Check if court is calibrated"""
        return self.homography_matrix is not None

    def draw_court_overlay(self, frame: np.ndarray) -> np.ndarray:
        """
        Draw court boundaries on frame for visualization

        Args:
            frame: Input frame

        Returns:
            Frame with court overlay
        """
        if not self.is_calibrated():
            return frame

        overlay = frame.copy()

        # Draw court corners
        if self.court_corners_pixel is not None:
            for point in self.court_corners_pixel:
                cv2.circle(overlay, (int(point[0]), int(point[1])), 10, (0, 255, 0), -1)

            # Draw court boundary lines
            pts = self.court_corners_pixel.astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(overlay, [pts], True, (0, 255, 0), 2)

        return overlay

    def save_calibration(self, filepath: str):
        """Save calibration data to file"""
        if not self.is_calibrated():
            print("Error: No calibration to save")
            return

        data = {
            'court_width_m': self.court_width,
            'court_length_m': self.court_length,
            'homography_matrix': self.homography_matrix.tolist(),
            'court_corners_pixel': self.court_corners_pixel.tolist()
        }

        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)

        print(f"✓ Calibration saved to {filepath}")

    def load_calibration(self, filepath: str) -> bool:
        """Load calibration data from file"""
        try:
            with open(filepath, 'r') as f:
                data = json.load(f)

            self.court_width = data['court_width_m']
            self.court_length = data['court_length_m']
            self.homography_matrix = np.array(data['homography_matrix'])
            self.court_corners_pixel = np.array(data['court_corners_pixel'])

            print(f"✓ Calibration loaded from {filepath}")
            return True

        except Exception as e:
            print(f"Error loading calibration: {e}")
            return False


class InteractiveCalibrator:
    """
    Interactive tool for manual court calibration
    Click on 4 corners of the court in order: top-left, top-right, bottom-right, bottom-left
    """

    def __init__(self):
        self.points = []
        self.window_name = "Court Calibration - Click 4 corners (TL, TR, BR, BL)"

    def _mouse_callback(self, event, x, y, flags, param):
        """Handle mouse clicks"""
        if event == cv2.EVENT_LBUTTONDOWN and len(self.points) < 4:
            self.points.append((x, y))
            print(f"Point {len(self.points)}: ({x}, {y})")

    def calibrate_interactive(self, frame: np.ndarray) -> Optional[List[Tuple[float, float]]]:
        """
        Interactive calibration - user clicks 4 corners

        Args:
            frame: First frame of video to calibrate on

        Returns:
            List of 4 corner points, or None if cancelled
        """
        display = frame.copy()
        cv2.namedWindow(self.window_name)
        cv2.setMouseCallback(self.window_name, self._mouse_callback)

        print("\nClick on 4 court corners in this order:")
        print("1. Top-left")
        print("2. Top-right")
        print("3. Bottom-right")
        print("4. Bottom-left")
        print("Press 'q' to cancel, 'r' to reset")

        while True:
            # Draw points
            temp = display.copy()
            for i, point in enumerate(self.points):
                cv2.circle(temp, point, 5, (0, 255, 0), -1)
                cv2.putText(temp, str(i+1), (point[0]+10, point[1]),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            # Draw lines between points
            if len(self.points) > 1:
                for i in range(len(self.points) - 1):
                    cv2.line(temp, self.points[i], self.points[i+1], (0, 255, 0), 2)

            cv2.imshow(self.window_name, temp)
            key = cv2.waitKey(1) & 0xFF

            if key == ord('q'):
                print("Calibration cancelled")
                cv2.destroyAllWindows()
                return None

            if key == ord('r'):
                print("Reset points")
                self.points = []

            if len(self.points) == 4:
                print("\n✓ All 4 points selected")
                cv2.destroyAllWindows()
                return self.points

        return None