import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt


class TrajectoryFitter:
    def __init__(self, gravity=9.81):
        self.g = gravity

    def trajectory_model(self, t, v0, theta, h0):
        """
        Physics model for projectile motion

        Parameters:
        t: time points
        v0: initial velocity
        theta: launch angle (in degrees)
        h0: initial height

        Returns:
        (x, y) coordinates at each time t
        """
        # Convert angle to radians
        theta_rad = np.radians(theta)

        # Initial velocities in x and y directions
        v0x = v0 * np.cos(theta_rad)
        v0y = v0 * np.sin(theta_rad)

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

        return np.column_stack((x, y))

    def fit_trajectory(self, points, times=None):
        """
        Fit trajectory to user-selected points

        Parameters:
        points: array of (x, y) coordinates
        times: optional array of timestamps for each point

        Returns:
        (v0, theta, h0): initial velocity, launch angle, initial height
        trajectory: predicted points along entire path
        """
        points = np.array(points)

        # If times not provided, estimate based on x positions
        if times is None:
            times = (points[:, 0] - points[0, 0]) / np.linalg.norm(
                points[1] - points[0]
            )

        # Initial guesses
        initial_height = points[0, 1]

        # Estimate initial angle and velocity from first two points
        if len(points) >= 2:
            dx = points[1, 0] - points[0, 0]
            dy = points[1, 1] - points[0, 1]
            initial_theta = np.degrees(np.arctan2(dy, dx))
            initial_v0 = np.sqrt(dx**2 + dy**2) / (times[1] - times[0])
        else:
            initial_theta = 45
            initial_v0 = 50

        # Fit physics model to points
        try:
            params, covariance = curve_fit(
                lambda t, v0, theta: self.trajectory_model(
                    t, v0, theta, initial_height
                ),
                times,
                points,
                p0=[initial_v0, initial_theta],
                bounds=([0, -90], [1000, 90]),  # Reasonable bounds for golf
            )

            v0_fit, theta_fit = params

            # Generate smooth trajectory for visualization
            t_smooth = np.linspace(min(times), max(times), 100)
            trajectory = self.trajectory_model(
                t_smooth, v0_fit, theta_fit, initial_height
            )

            return {
                "initial_velocity": v0_fit,
                "launch_angle": theta_fit,
                "initial_height": initial_height,
                "trajectory": trajectory,
                "times": t_smooth,
                "covariance": covariance,
            }

        except RuntimeError as e:
            print(f"Fitting failed: {e}")
            return None

    def calculate_metrics(self, fit_results):
        """Calculate additional metrics from the fit"""
        v0 = fit_results["initial_velocity"]
        theta = fit_results["launch_angle"]
        h0 = fit_results["initial_height"]

        # Maximum height
        theta_rad = np.radians(theta)
        v0y = v0 * np.sin(theta_rad)
        max_height = h0 + (v0y**2) / (2 * self.g)

        # Total distance (range)
        t_total = (v0y + np.sqrt(v0y**2 + 2 * self.g * h0)) / self.g
        total_distance = v0 * np.cos(theta_rad) * t_total

        # Flight time
        flight_time = t_total

        return {
            "max_height": max_height,
            "total_distance": total_distance,
            "flight_time": flight_time,
            "initial_velocity_mph": v0 * 2.237,  # Convert m/s to mph
        }

    def plot_trajectory(self, points, fit_results):
        """Visualize the fitted trajectory and original points"""
        plt.figure(figsize=(12, 6))

        # Plot original points
        points = np.array(points)
        plt.scatter(points[:, 0], points[:, 1], color="red", label="Selected Points")

        # Plot fitted trajectory
        trajectory = fit_results["trajectory"]
        plt.plot(trajectory[:, 0], trajectory[:, 1], "b-", label="Fitted Trajectory")

        plt.grid(True)
        plt.xlabel("Distance (m)")
        plt.ylabel("Height (m)")
        plt.title("Golf Ball Trajectory Fit")
        plt.legend()
        plt.axis("equal")
        plt.show()


# Example usage:
if __name__ == "__main__":
    # Sample points (x, y) in meters
    points = [
        (0, 0),  # Starting point
        (50, 20),  # Some point during flight
        (100, 0),  # Landing point
    ]

    fitter = TrajectoryFitter()
    results = fitter.fit_trajectory(points)

    if results:
        metrics = fitter.calculate_metrics(results)
        print("\nFitted Parameters:")
        print(
            f"Initial Velocity: {results['initial_velocity']:.1f} m/s ({metrics['initial_velocity_mph']:.1f} mph)"
        )
        print(f"Launch Angle: {results['launch_angle']:.1f} degrees")
        print(f"Max Height: {metrics['max_height']:.1f} m")
        print(f"Total Distance: {metrics['total_distance']:.1f} m")
        print(f"Flight Time: {metrics['flight_time']:.1f} s")

        fitter.plot_trajectory(points, results)