"""
3D camera calibration using court geometry.
Extracted from dagster_project/assets/camera_calibration.py and coordinate_transform_3d.py.

For our setup: two cameras mounted at the net center, looking at opposite halves.
Each camera gets its own calibration.
"""

import cv2
import numpy as np
import json
from typing import Dict, Optional, Tuple

# Pickleball court dimensions (meters)
COURT_LENGTH = 13.4
COURT_WIDTH = 6.1
HALF_COURT_LENGTH = COURT_LENGTH / 2  # 6.7m
NET_HEIGHT = 0.914  # meters at center
BALL_DIAMETER = 0.074  # meters (74mm)


class CameraCalibrator:
    """3D camera calibration using solvePnP with known court geometry."""

    def __init__(self):
        self.camera_matrix = None
        self.rotation_vec = None
        self.translation_vec = None
        self.rotation_matrix = None
        self.homography = None  # ground plane homography
        self.calibrated = False

    def calibrate(self, court_corners_pixel: np.ndarray,
                  court_corners_3d: np.ndarray,
                  image_width: int, image_height: int) -> bool:
        """Calibrate camera using court corner correspondences.

        Args:
            court_corners_pixel: Nx2 array of corner positions in pixels
            court_corners_3d: Nx3 array of corresponding 3D world coordinates
            image_width, image_height: frame dimensions
        """
        # Estimate camera matrix
        focal_length = float(max(image_width, image_height))
        cx, cy = image_width / 2.0, image_height / 2.0

        self.camera_matrix = np.array([
            [focal_length, 0, cx],
            [0, focal_length, cy],
            [0, 0, 1]
        ], dtype=np.float32)

        try:
            success, self.rotation_vec, self.translation_vec = cv2.solvePnP(
                court_corners_3d.astype(np.float32),
                court_corners_pixel.astype(np.float32),
                self.camera_matrix,
                None,
                flags=cv2.SOLVEPNP_ITERATIVE
            )
            if not success:
                return False

            self.rotation_matrix, _ = cv2.Rodrigues(self.rotation_vec)

            # Build ground plane homography (Z=0)
            ground_2d = court_corners_3d[:, :2].astype(np.float32)
            self.homography = cv2.getPerspectiveTransform(
                court_corners_pixel[:4].astype(np.float32),
                ground_2d[:4]
            )

            self.calibrated = True
            return True
        except Exception as e:
            print(f"Calibration failed: {e}")
            return False

    def pixel_to_ground(self, px: float, py: float) -> Optional[Tuple[float, float]]:
        """Project pixel to ground plane (Z=0) using homography."""
        if self.homography is None:
            return None
        pt = np.array([[[px, py]]], dtype=np.float32)
        result = cv2.perspectiveTransform(pt, self.homography)
        return float(result[0][0][0]), float(result[0][0][1])

    def pixel_to_3d(self, px: float, py: float,
                    ball_diameter_px: float = 20.0) -> Optional[Tuple[float, float, float]]:
        """Estimate 3D position using ray casting and ball size.

        If ball appears large (>30px diameter), assume on ground (Z=0).
        Otherwise, estimate height from apparent size.
        """
        if not self.calibrated:
            return None

        # Large ball → on ground
        if ball_diameter_px > 30:
            ground = self.pixel_to_ground(px, py)
            if ground:
                return (ground[0], ground[1], 0.0)
            return None

        fx = self.camera_matrix[0, 0]
        fy = self.camera_matrix[1, 1]
        cx = self.camera_matrix[0, 2]
        cy = self.camera_matrix[1, 2]
        focal = (fx + fy) / 2

        # Ray direction in camera coords
        ray_cam = np.array([
            (px - cx) / fx,
            (py - cy) / fy,
            1.0
        ])
        ray_cam /= np.linalg.norm(ray_cam)

        # Transform to world coords
        ray_world = self.rotation_matrix.T @ ray_cam

        # Distance from apparent ball size
        distance = (BALL_DIAMETER * focal) / max(ball_diameter_px, 1.0)

        # Camera position in world
        cam_pos = (-self.rotation_matrix.T @ self.translation_vec).flatten()

        point_3d = cam_pos + ray_world * distance
        x, y, z = point_3d

        # Sanity check
        if z < 0 or z > 10:
            ground = self.pixel_to_ground(px, py)
            if ground:
                return (ground[0], ground[1], 0.0)
            return None

        return (float(x), float(y), float(z))

    def save(self, filepath: str):
        if not self.calibrated:
            return
        data = {
            'camera_matrix': self.camera_matrix.tolist(),
            'rotation_vector': self.rotation_vec.flatten().tolist(),
            'translation_vector': self.translation_vec.flatten().tolist(),
            'rotation_matrix': self.rotation_matrix.tolist(),
            'homography': self.homography.tolist() if self.homography is not None else None,
        }
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=2)

    def load(self, filepath: str) -> bool:
        try:
            with open(filepath, 'r') as f:
                data = json.load(f)
            self.camera_matrix = np.array(data['camera_matrix'], dtype=np.float32)
            self.rotation_vec = np.array(data['rotation_vector'], dtype=np.float32).reshape(3, 1)
            self.translation_vec = np.array(data['translation_vector'], dtype=np.float32).reshape(3, 1)
            self.rotation_matrix = np.array(data['rotation_matrix'], dtype=np.float32)
            if data.get('homography'):
                self.homography = np.array(data['homography'], dtype=np.float32)
            self.calibrated = True
            return True
        except Exception as e:
            print(f"Failed to load calibration: {e}")
            return False


def get_half_court_3d_points(side: str) -> np.ndarray:
    """Get 3D coordinates for one half of the court.

    Args:
        side: 'left' (CAM 0, X: 0 to 6.7m) or 'right' (CAM 1, X: 6.7 to 13.4m)

    Returns:
        4x3 array of court corners in world coords (X=along court, Y=across, Z=up)
    """
    if side == 'left':
        return np.array([
            [0, 0, 0],                        # near-left (baseline)
            [0, COURT_WIDTH, 0],               # near-right
            [HALF_COURT_LENGTH, COURT_WIDTH, 0],  # far-right (net)
            [HALF_COURT_LENGTH, 0, 0],            # far-left (net)
        ], dtype=np.float32)
    else:
        return np.array([
            [COURT_LENGTH, 0, 0],              # near-left (baseline)
            [COURT_LENGTH, COURT_WIDTH, 0],    # near-right
            [HALF_COURT_LENGTH, COURT_WIDTH, 0],  # far-right (net)
            [HALF_COURT_LENGTH, 0, 0],            # far-left (net)
        ], dtype=np.float32)
