import numpy as np
from segment_anything import sam_model_registry, SamPredictor
import cv2
from scipy.optimize import curve_fit

class GolfTrajectoryPredictor:
    def __init__(self, sam_checkpoint):
        # Initialize SAM model
        self.sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
        self.predictor = SamPredictor(self.sam)

    def get_club_metrics(self, frame, point):
        """Extract golf club angle and position using SAM"""
        self.predictor.set_image(frame)
        masks, scores, _ = self.predictor.predict(
            point_coords=np.array([point]),
            point_labels=np.array([1])
        )

        # Get the best mask
        club_mask = masks[np.argmax(scores)]

        # Calculate club angle from mask
        coords = np.column_stack(np.where(club_mask))
        if len(coords) < 2:
            return None, None

        # Fit line to get club angle
        vx, vy, x0, y0 = cv2.fitLine(coords, cv2.DIST_L2, 0, 0.01, 0.01)
        angle = np.arctan2(vy, vx)

        return angle, (x0[0], y0[0])

    def physics_trajectory(self, t, v0, theta, h0, g=9.81):
        """Model the physics of projectile motion"""
        # Convert angle to radians
        theta_rad = np.radians(theta)

        # Initial velocities
        v0x = v0 * np.cos(theta_rad)
        v0y = v0 * np.sin(theta_rad)

        # Position equations
        x = v0x * t
        y = h0 + v0y * t - 0.5 * g * t**2

        return np.column_stack((x, y))

    def fit_trajectory(self, points, initial_height, club_angle=None):
        """Fit trajectory to user-selected points"""
        times = np.linspace(0, 1, len(points))

        # Initial guess for parameters
        # Use club angle if available to better estimate initial velocity direction
        initial_theta = club_angle if club_angle is not None else 45

        # Calculate approximate initial velocity from first two points
        if len(points) >= 2:
            dx = points[1][0] - points[0][0]
            dy = points[1][1] - points[0][1]
            initial_v0 = np.sqrt(dx**2 + dy**2) / (times[1] - times[0])
        else:
            initial_v0 = 50  # Default initial guess

        # Fit physics model to points
        try:
            params, _ = curve_fit(
                lambda t, v0, theta: self.physics_trajectory(t, v0, theta, initial_height),
                times,
                points,
                p0=[initial_v0, initial_theta]
            )
            return params
        except RuntimeError:
            return None

    def predict_full_trajectory(self, video_path, user_selected_points):
        """Main function to predict full trajectory"""
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)

        # Get initial frame for club analysis
        ret, first_frame = cap.read()
        if not ret:
            return None

        # Get club metrics from first frame
        club_angle, club_pos = self.get_club_metrics(first_frame, user_selected_points[0])

        # Convert pixel coordinates to physical space
        # This requires camera calibration in practice
        scale_factor = 0.01  # meters per pixel
        physical_points = np.array(user_selected_points) * scale_factor

        # Fit trajectory
        params = self.fit_trajectory(
            physical_points,
            initial_height=physical_points[0][1],
            club_angle=club_angle
        )

        if params is None:
            return None

        # Generate full trajectory
        t = np.linspace(0, len(user_selected_points)/fps, 100)
        full_trajectory = self.physics_trajectory(t, params[0], params[1], physical_points[0][1])

        # Convert back to pixel coordinates
        pixel_trajectory = full_trajectory / scale_factor

        return pixel_trajectory, params

    def visualize_trajectory(self, frame, trajectory, user_points):
        """Visualize the predicted trajectory and user-selected points"""
        vis_frame = frame.copy()

        # Draw predicted trajectory
        trajectory = trajectory.astype(np.int32)
        for i in range(len(trajectory)-1):
            cv2.line(vis_frame, tuple(trajectory[i]), tuple(trajectory[i+1]),
                    (0, 255, 0), 2)

        # Draw user-selected points
        for point in user_points:
            cv2.circle(vis_frame, tuple(map(int, point)), 5, (255, 0, 0), -1)

        return vis_frame