"""Asset: Calibrate camera using court corners and net position"""

import cv2
import numpy as np
from typing import Dict
from dagster import asset, AssetExecutionContext


@asset(
    io_manager_key="json_io_manager",
    compute_kind="opencv",
    description="Calibrate camera using cv2.solvePnP with court corners and net"
)
def calibrate_camera_3d(
    context: AssetExecutionContext,
    detect_court_keypoints: Dict,
    detect_net: Dict
) -> Dict:
    """
    Calibrate camera to get 3D pose using known 3D↔2D point correspondences

    Inputs:
        - detect_court_keypoints: 4 court corners in pixels
        - detect_net: 4 net corners in pixels

    Known 3D points:
        - Court: 13.4m × 6.1m rectangle at Z=0
        - Net: height 0.914m at middle of court (Y=3.05m)

    Outputs:
        Camera calibration parameters

    Returns:
        Dict with:
            - camera_matrix: [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
            - rotation_vector: [rx, ry, rz]
            - translation_vector: [tx, ty, tz]
            - rotation_matrix: 3x3 matrix
            - reprojection_error: RMS error in pixels
            - calibrated: success flag
    """
    # Get 2D points (pixels)
    court_corners = np.array(detect_court_keypoints['corners_pixel'], dtype=np.float32)
    net_corners = np.array(detect_net['net_corners_pixel'], dtype=np.float32)

    # Define 3D points (meters) in world coordinates
    # Court corners (Z=0, on ground)
    court_3d = np.array([
        [0, 0, 0],       # TL
        [13.4, 0, 0],    # TR
        [13.4, 6.1, 0],  # BR
        [0, 6.1, 0]      # BL
    ], dtype=np.float32)

    # Net endpoints (2 points)
    # Net is at Y=3.05m (middle of 6.1m court width)
    # We have 2 endpoints: left and right side at top of net
    net_3d = np.array([
        [0, 3.05, 0.914],       # Left endpoint (top of net)
        [13.4, 3.05, 0.914],    # Right endpoint (top of net)
    ], dtype=np.float32)

    # Combine all 3D and 2D points
    object_points = np.vstack([court_3d, net_3d])  # 6 points total (4 court + 2 net)
    image_points = np.vstack([court_corners, net_corners])

    context.log.info(f"Calibrating with {len(object_points)} point correspondences")
    context.log.info(f"3D points shape: {object_points.shape}")
    context.log.info(f"2D points shape: {image_points.shape}")

    # Initial camera matrix estimate
    # Assume principal point at image center
    w = detect_court_keypoints['frame_width']
    h = detect_court_keypoints['frame_height']
    cx = w / 2
    cy = h / 2
    # Estimate focal length (typical for drone/action cameras)
    focal_length = max(w, h)  # Initial guess

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

    # No lens distortion (assume corrected or minimal)
    dist_coeffs = None

    # Solve PnP to get camera pose
    try:
        success, rotation_vec, translation_vec = cv2.solvePnP(
            object_points,
            image_points,
            camera_matrix,
            dist_coeffs,
            flags=cv2.SOLVEPNP_ITERATIVE
        )

        if not success:
            context.log.error("cv2.solvePnP failed")
            return {
                "calibrated": False,
                "error": "solvePnP failed"
            }

        # Convert rotation vector to rotation matrix
        rotation_matrix, _ = cv2.Rodrigues(rotation_vec)

        # Calculate reprojection error
        projected_points, _ = cv2.projectPoints(
            object_points,
            rotation_vec,
            translation_vec,
            camera_matrix,
            dist_coeffs
        )
        projected_points = projected_points.reshape(-1, 2)
        reprojection_error = np.sqrt(np.mean((image_points - projected_points) ** 2))

        context.log.info(f"✓ Camera calibration successful")
        context.log.info(f"  Reprojection error: {reprojection_error:.2f} pixels")
        context.log.info(f"  Focal length: {focal_length:.1f}")
        context.log.info(f"  Rotation vector: {rotation_vec.flatten()}")
        context.log.info(f"  Translation vector: {translation_vec.flatten()}")

        return {
            "camera_matrix": camera_matrix.tolist(),
            "rotation_vector": rotation_vec.flatten().tolist(),
            "translation_vector": translation_vec.flatten().tolist(),
            "rotation_matrix": rotation_matrix.tolist(),
            "reprojection_error": float(reprojection_error),
            "focal_length": float(focal_length),
            "principal_point": [float(cx), float(cy)],
            "image_size": [w, h],
            "calibrated": True
        }

    except Exception as e:
        context.log.error(f"Error during camera calibration: {e}")
        return {
            "calibrated": False,
            "error": str(e)
        }