"""Asset: Transform ball positions to 3D coordinates (X, Y, Z) in meters"""

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


@asset(
    io_manager_key="json_io_manager",
    compute_kind="opencv",
    description="Compute 3D ball coordinates (X, Y, Z) using camera calibration"
)
def compute_ball_3d_coordinates(
    context: AssetExecutionContext,
    detect_court_keypoints: Dict,
    detect_ball_positions: List[Dict],
    calibrate_camera_3d: Dict
) -> List[Dict]:
    """
    Transform ball pixel coordinates to 3D world coordinates (meters)

    Strategy:
    - If ball on ground: use homography (Z=0)
    - If ball in air: use ray casting + bbox size estimation (Z>0)

    Inputs:
        - detect_court_keypoints: court corners
        - detect_ball_positions: ball positions in pixels
        - calibrate_camera_3d: camera calibration

    Returns:
        List of dicts with 3D coordinates for each frame:
        [
            {
                "frame": 0,
                "x_m": float or null,  # X position on court (0-13.4m)
                "y_m": float or null,  # Y position on court (0-6.1m)
                "z_m": float or null,  # Z height above court (meters)
                "on_ground": bool,     # True if ball touching court
                "confidence": float    # Detection confidence
            },
            ...
        ]
    """
    if not calibrate_camera_3d.get('calibrated'):
        context.log.error("Camera not calibrated, cannot compute 3D coordinates")
        return [{"frame": i, "x_m": None, "y_m": None, "z_m": None, "on_ground": False}
                for i in range(len(detect_ball_positions))]

    # Extract calibration parameters
    camera_matrix = np.array(calibrate_camera_3d['camera_matrix'], dtype=np.float32)
    rotation_vec = np.array(calibrate_camera_3d['rotation_vector'], dtype=np.float32)
    translation_vec = np.array(calibrate_camera_3d['translation_vector'], dtype=np.float32)
    rotation_matrix = np.array(calibrate_camera_3d['rotation_matrix'], dtype=np.float32)

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

    # Build homography for ground plane (Z=0)
    court_corners_pixel = np.array(detect_court_keypoints['corners_pixel'], dtype=np.float32)
    court_corners_meters = np.array([
        [0, 0],
        [13.4, 0],
        [13.4, 6.1],
        [0, 6.1]
    ], dtype=np.float32)

    homography_matrix = cv2.getPerspectiveTransform(court_corners_pixel, court_corners_meters)

    # Camera position in world coordinates
    camera_position = -rotation_matrix.T @ translation_vec.reshape(3, 1)

    context.log.info(f"Processing {len(detect_ball_positions)} frames for 3D coordinate transformation")

    results = []

    for i, ball_det in enumerate(detect_ball_positions):
        if ball_det['x'] is None:
            # No ball detected
            results.append({
                "frame": i,
                "x_m": None,
                "y_m": None,
                "z_m": None,
                "on_ground": False,
                "confidence": 0.0
            })
            continue

        ball_x = ball_det['x']
        ball_y = ball_det['y']
        ball_diameter = ball_det['diameter_px']
        bbox = ball_det['bbox']
        confidence = ball_det['confidence']

        # Strategy 1: Try ground plane projection (assume Z=0)
        ball_point_2d = np.array([[ball_x, ball_y]], dtype=np.float32)
        ball_ground = cv2.perspectiveTransform(ball_point_2d.reshape(-1, 1, 2), homography_matrix)
        x_ground, y_ground = ball_ground[0][0]

        # Check if ball is likely on ground by comparing bbox size
        # If ball bbox is large → likely on ground
        # Simple heuristic: if diameter > threshold, assume on ground
        on_ground_threshold = 30  # pixels - tune this based on typical ball size on ground

        if ball_diameter and ball_diameter > on_ground_threshold:
            # Ball likely on ground
            results.append({
                "frame": i,
                "x_m": float(x_ground),
                "y_m": float(y_ground),
                "z_m": 0.0,
                "on_ground": True,
                "confidence": float(confidence)
            })
        else:
            # Ball likely in air - use ray casting
            try:
                # Unproject 2D point to 3D ray
                point_2d_normalized = np.array([
                    (ball_x - cx) / fx,
                    (ball_y - cy) / fy,
                    1.0
                ])

                # Ray direction in camera coordinates
                ray_camera = point_2d_normalized / np.linalg.norm(point_2d_normalized)

                # Transform ray to world coordinates
                ray_world = rotation_matrix.T @ ray_camera

                # Estimate distance using ball size
                # Real pickleball diameter: 74mm = 0.074m
                ball_diameter_real = 0.074  # meters
                ball_diameter_pixels = ball_diameter if ball_diameter else 20  # fallback to 20px

                # Distance formula: D = (real_size * focal_length) / pixel_size
                distance = (ball_diameter_real * focal_length) / ball_diameter_pixels

                # Compute 3D point along ray
                ball_3d = camera_position.flatten() + ray_world * distance

                x_3d, y_3d, z_3d = ball_3d

                # Sanity checks
                if z_3d < 0 or z_3d > 10:  # Height should be 0-10m
                    # Invalid, fallback to ground projection
                    results.append({
                        "frame": i,
                        "x_m": float(x_ground),
                        "y_m": float(y_ground),
                        "z_m": 0.0,
                        "on_ground": True,
                        "confidence": float(confidence)
                    })
                else:
                    results.append({
                        "frame": i,
                        "x_m": float(x_3d),
                        "y_m": float(y_3d),
                        "z_m": float(z_3d),
                        "on_ground": False,
                        "confidence": float(confidence)
                    })

            except Exception as e:
                context.log.warning(f"Frame {i}: Error computing 3D coords, using ground projection: {e}")
                results.append({
                    "frame": i,
                    "x_m": float(x_ground),
                    "y_m": float(y_ground),
                    "z_m": 0.0,
                    "on_ground": True,
                    "confidence": float(confidence)
                })

        if (i + 1) % 20 == 0:
            context.log.info(f"Processed {i + 1}/{len(detect_ball_positions)} frames")

    # Statistics
    detected_count = sum(1 for r in results if r['x_m'] is not None)
    on_ground_count = sum(1 for r in results if r.get('on_ground', False))
    in_air_count = detected_count - on_ground_count

    context.log.info(f"✓ Computed 3D coordinates: {detected_count} detections")
    context.log.info(f"  On ground: {on_ground_count}, In air: {in_air_count}")

    return results
