import numpy as np
import matplotlib.pyplot as plt
from dataclasses import dataclass
from typing import List, Tuple, Optional

@dataclass
class Motion1D:
    """1D Kinematic motion calculator for point mass"""
    initial_position: float = 0.0  # x0 (m)
    initial_velocity: float = 0.0  # v0 (m/s)
    acceleration: float = 0.0      # a (m/s²)
    
    def position(self, t: float) -> float:
        """Calculate position at time t using x = x0 + v0*t + 0.5*a*t²"""
        return self.initial_position + self.initial_velocity * t + 0.5 * self.acceleration * t**2
    
    def velocity(self, t: float) -> float:
        """Calculate velocity at time t using v = v0 + a*t"""
        return self.initial_velocity + self.acceleration * t
    
    def time_arrays(self, duration: float, dt: float = 0.01) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Generate time arrays for position, velocity, and acceleration"""
        t = np.arange(0, duration + dt, dt)
        x = np.array([self.position(time) for time in t])
        v = np.array([self.velocity(time) for time in t])
        a = np.full_like(t, self.acceleration)
        return t, x, v, a

@dataclass  
class Motion2D:
    """2D Kinematic motion calculator (projectile motion)"""
    launch_speed: float = 0.0          # Initial speed magnitude (m/s)
    launch_angle: float = 0.0          # Launch angle in degrees
    launch_height: float = 0.0         # Launch height above ground (m)
    launch_x: float = 0.0              # Horizontal launch position (m)
    gravity: float = 9.81              # Acceleration due to gravity (m/s²)
    air_resistance: bool = False       # Enable/disable air resistance
    drag_coefficient: float = 0.1      # Simple drag coefficient (1/s)
    is_sphere: bool = False            # Toggle between point mass and sphere
    sphere_radius: float = 0.037       # Sphere radius in meters (default: baseball ~37mm)
    sphere_density: float = 700        # Sphere density in kg/m³ (default: baseball ~700 kg/m³)
    air_density: float = 1.225         # Air density in kg/m³ (sea level)
    sphere_drag_coeff: float = 0.47    # Aerodynamic drag coefficient for smooth sphere
    

    def __post_init__(self):
        """Calculate velocity components and sphere properties from launch parameters"""
        angle_rad = np.radians(self.launch_angle)
        self.initial_velocity_x = self.launch_speed * np.cos(angle_rad)
        self.initial_velocity_y = self.launch_speed * np.sin(angle_rad)
        self.initial_position = (self.launch_x, self.launch_height)
        self.acceleration = (0.0, -self.gravity)
        
        # Calculate sphere properties if using sphere model
        if self.is_sphere:
            self.sphere_volume = (4/3) * np.pi * self.sphere_radius**3
            self.sphere_mass = self.sphere_density * self.sphere_volume
            self.sphere_cross_section = np.pi * self.sphere_radius**2

            # For sphere: F_drag = 0.5 * rho * Cd * A * v²
            # In our linear approximation: F_drag = k * v * m (where k has units 1/s)
            # So: k = (0.5 * rho * Cd * A) / m
            self.effective_drag_coeff = (0.5 * self.air_density * self.sphere_drag_coeff * 
                                       self.sphere_cross_section) / self.sphere_mass
        else:
            # Point mass defaults
            self.sphere_mass = 1.0
            self.sphere_volume = 0.0
            self.sphere_cross_section = 0.0
            self.effective_drag_coeff = self.drag_coefficient
    
    def position(self, t: float) -> Tuple[float, float]:
        """Calculate 2D position at time t"""
        if not self.air_resistance:
            # No air resistance - standard kinematic equations
            x = self.initial_position[0] + self.initial_velocity_x * t
            y = self.initial_position[1] + self.initial_velocity_y * t - 0.5 * self.gravity * t**2
        else:
            # With air resistance - improved model
            k = self.effective_drag_coeff
            
            if k == 0:
                # Fallback to no air resistance
                x = self.initial_position[0] + self.initial_velocity_x * t
                y = self.initial_position[1] + self.initial_velocity_y * t - 0.5 * self.gravity * t**2
            else:
                # Terminal velocity
                v_terminal = self.gravity / k
                exp_kt = np.exp(-k * t)
                
                # Horizontal motion (no forces except drag)
                x = self.initial_position[0] + (self.initial_velocity_x / k) * (1 - exp_kt)
                
                # Vertical motion (gravity + drag)
                # Analytical solution: y = y0 + (v0y + vt)/k * (1 - e^(-kt)) - vt*t
                y = (self.initial_position[1] + 
                     (self.initial_velocity_y + v_terminal) / k * (1 - exp_kt) - 
                     v_terminal * t)
        
        return x, y

    def velocity(self, t: float) -> Tuple[float, float]:
        """Calculate 2D velocity at time t"""
        if not self.air_resistance:
            # No air resistance
            vx = self.initial_velocity_x
            vy = self.initial_velocity_y - self.gravity * t
        else:
            # With air resistance
            k = self.effective_drag_coeff
            exp_kt = np.exp(-k * t)
            
            # Horizontal velocity
            vx = self.initial_velocity_x * exp_kt
            
            # Vertical velocity
            v_terminal = self.gravity / k if k > 0 else float('inf')
            vy = (self.initial_velocity_y + v_terminal) * exp_kt - v_terminal
        
        return vx, vy
    
# Lines 140-152 - Update get_sphere_info method
    def get_sphere_info(self) -> dict:
        """Return sphere physical properties"""
        if not self.is_sphere:
            return {}
        
        # Calculate terminal velocity: vt = mg/(0.5 * rho * Cd * A) = g/k
        terminal_vel = self.gravity / self.effective_drag_coeff if self.effective_drag_coeff > 0 else float('inf')
        
        return {
            'radius_mm': self.sphere_radius * 1000,
            'diameter_mm': self.sphere_radius * 2000,
            'mass_g': self.sphere_mass * 1000,
            'volume_cm3': self.sphere_volume * 1000000,
            'cross_section_cm2': self.sphere_cross_section * 10000,
            'density_kg_m3': self.sphere_density,
            'terminal_velocity_ms': terminal_vel
        }


    def get_launch_info(self) -> dict:
        """Return comprehensive launch information"""
        flight_time = self.calculate_flight_time()
        info = {
            'launch_speed': self.launch_speed,
            'launch_angle': self.launch_angle,
            'launch_height': self.launch_height,
            'initial_velocity_x': self.initial_velocity_x,
            'initial_velocity_y': self.initial_velocity_y,
            'flight_time': flight_time,
            'range': self.calculate_range(),
            'max_height': self.calculate_max_height(),
            'air_resistance': self.air_resistance,
            'is_sphere': self.is_sphere,
            'effective_drag_coeff': self.effective_drag_coeff if self.air_resistance else 0
        }
        
        # Add sphere info if applicable
        if self.is_sphere:
            info.update(self.get_sphere_info())
            
        return info


    # Static method for common sphere presets
    @staticmethod
    def get_sphere_presets() -> dict:
        """Return common sphere presets with realistic properties"""
        return {
            'ping_pong': {
                'name': '🏓 Ping Pong Ball',
                'radius': 0.02,      # 20mm radius (40mm diameter)
                'density': 84,       # Very light
                'drag_coeff': 0.47
            },
            'golf': {
                'name': '⛳ Golf Ball',
                'radius': 0.0214,    # 21.4mm radius (42.8mm diameter)
                'density': 1130,     # Dense
                'drag_coeff': 0.24   # Dimpled surface reduces drag
            },
            'baseball': {
                'name': '⚾ Baseball',
                'radius': 0.037,     # 37mm radius (74mm diameter)
                'density': 700,      # Medium density
                'drag_coeff': 0.35   # Stitched surface
            },
            'tennis': {
                'name': '🎾 Tennis Ball',
                'radius': 0.033,     # 33mm radius (66mm diameter)
                'density': 370,      # Light with hollow core
                'drag_coeff': 0.51   # Fuzzy surface increases drag
            },
            'basketball': {
                'name': '🏀 Basketball',
                'radius': 0.12,      # 120mm radius (240mm diameter)
                'density': 60,       # Very light (hollow)
                'drag_coeff': 0.47
            },
            'bowling': {
                'name': '🎳 Bowling Ball',
                'radius': 0.108,     # 108mm radius (216mm diameter)
                'density': 1400,     # Very dense
                'drag_coeff': 0.47
            }
        }


    def trajectory_data(self, duration: float, dt: float = 0.01) -> dict:
        """Generate complete trajectory data, stopping when projectile hits ground"""
        t = np.arange(0, duration + dt, dt)
        positions = []
        velocities = []
        times = []
        
        for time in t:
            pos = self.position(time)
            vel = self.velocity(time)
            
            # Stop if projectile goes below ground level (y < 0)
            if pos[1] < 0 and len(positions) > 0:
                break
                
            positions.append(pos)
            velocities.append(vel)
            times.append(time)
        
        if len(positions) == 0:
            # Handle edge case
            positions = [(self.launch_x, self.launch_height)]
            velocities = [(self.initial_velocity_x, self.initial_velocity_y)]
            times = [0]
        
        positions = np.array(positions)
        velocities = np.array(velocities)
        times = np.array(times)
        
        return {
            'time': times,
            'x': positions[:, 0],
            'y': positions[:, 1],
            'vx': velocities[:, 0],
            'vy': velocities[:, 1],
            'speed': np.sqrt(velocities[:, 0]**2 + velocities[:, 1]**2)
        }

    def calculate_flight_time(self) -> float:
        """Calculate approximate flight time"""
        if not self.air_resistance:
            # Original calculation without air resistance
            h = self.launch_height
            vy0 = self.initial_velocity_y
            g = self.gravity
            
            if g == 0:
                if vy0 == 0:
                    return float('inf') if h >= 0 else 0
                return h / vy0 if vy0 < 0 else float('inf')
            
            discriminant = vy0**2 + 2*g*h
            if discriminant < 0:
                return 0
                
            t1 = (vy0 + np.sqrt(discriminant)) / g
            t2 = (vy0 - np.sqrt(discriminant)) / g
            
            return max(t1, t2) if max(t1, t2) > 0 else 0
        else:
            # With air resistance, use numerical approach to find landing time
            # This is an approximation - we'll simulate until we hit ground
            max_time = 20.0  # Maximum simulation time
            dt = 0.01
            
            for t in np.arange(0, max_time, dt):
                _, y = self.position(t)
                if y <= 0 and t > 0:
                    return t
            
            return max_time  # Fallback
    
    def calculate_range(self) -> float:
        """Calculate horizontal range of projectile"""
        flight_time = self.calculate_flight_time()
        if flight_time <= 0:
            return 0
        
        # Get final position
        final_x, _ = self.position(flight_time)
        return final_x - self.launch_x

    def get_launch_info(self) -> dict:
        """Return comprehensive launch information"""
        flight_time = self.calculate_flight_time()
        info = {
            'launch_speed': self.launch_speed,
            'launch_angle': self.launch_angle,
            'launch_height': self.launch_height,
            'initial_velocity_x': self.initial_velocity_x,
            'initial_velocity_y': self.initial_velocity_y,
            'flight_time': flight_time,
            'range': self.calculate_range(),
            'max_height': self.calculate_max_height(),
            'air_resistance': self.air_resistance,
            'is_sphere': self.is_sphere,
            'effective_drag_coeff': self.effective_drag_coeff if self.air_resistance else 0
        }
        
        # Add sphere info if applicable
        if self.is_sphere:
            sphere_info = self.get_sphere_info()
            info.update(sphere_info)
            
        return info

    def calculate_max_height(self) -> float:
        """Calculate maximum height reached by projectile"""
        if not self.air_resistance:
            # Analytical solution for no air resistance
            if self.gravity == 0:
                return float('inf') if self.initial_velocity_y > 0 else self.launch_height
            
            # Time to reach maximum height: when vy = 0
            t_max = self.initial_velocity_y / self.gravity
            if t_max <= 0:  # Projectile going downward initially
                return self.launch_height
                
            _, max_y = self.position(t_max)
            return max(max_y, self.launch_height)
        else:
            # Numerical approach for air resistance
            flight_time = self.calculate_flight_time()
            dt = 0.01
            max_height = self.launch_height
            
            for t in np.arange(0, flight_time + dt, dt):
                _, y = self.position(t)
                if y > max_height:
                    max_height = y
                    
            return max_height

class KinematicsVisualizer:
    """Create visualizations for kinematic motion"""
    
    @staticmethod
    def plot_1d_motion(motion: Motion1D, duration: float, title: str = "1D Kinematic Motion"):
        """Create comprehensive 1D motion plots"""
        t, x, v, a = motion.time_arrays(duration)
        
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 12))
        fig.suptitle(title, fontsize=16, fontweight='bold')
        
        # Position vs Time
        ax1.plot(t, x, 'b-', linewidth=2, label='Position')
        ax1.set_ylabel('Position (m)')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        ax1.set_title('Position vs Time')
        
        # Velocity vs Time
        ax2.plot(t, v, 'r-', linewidth=2, label='Velocity')
        ax2.set_ylabel('Velocity (m/s)')
        ax2.grid(True, alpha=0.3)
        ax2.legend()
        ax2.set_title('Velocity vs Time')
        
        # Acceleration vs Time
        ax3.plot(t, a, 'g-', linewidth=2, label='Acceleration')
        ax3.set_xlabel('Time (s)')
        ax3.set_ylabel('Acceleration (m/s²)')
        ax3.grid(True, alpha=0.3)
        ax3.legend()
        ax3.set_title('Acceleration vs Time')
        
        plt.tight_layout()
        return fig
    
    @staticmethod
    def plot_2d_trajectory(motion: Motion2D, duration: float = None, title: str = "2D Projectile Motion"):
        """Create 2D trajectory visualization with launch info"""
        if duration is None:
            duration = motion.calculate_flight_time()
        
        data = motion.trajectory_data(duration)
        info = motion.get_launch_info()
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # Create detailed title with launch parameters
        detailed_title = (f"{title}\n"
                         f"Launch: {info['launch_speed']:.1f} m/s at {info['launch_angle']:.1f}° "
                         f"from {info['launch_height']:.1f}m height")
        fig.suptitle(detailed_title, fontsize=14, fontweight='bold')
        
        # Trajectory plot with key points marked
        ax1.plot(data['x'], data['y'], 'b-', linewidth=2, label='Trajectory')
        
        # Mark launch point
        ax1.plot(motion.launch_x, motion.launch_height, 'go', markersize=8, label='Launch')
        
        # Mark landing point
        if len(data['x']) > 0:
            ax1.plot(data['x'][-1], data['y'][-1], 'ro', markersize=8, label='Landing')
        
        # Mark maximum height
        max_height_idx = np.argmax(data['y'])
        if len(data['x']) > 0:
            ax1.plot(data['x'][max_height_idx], data['y'][max_height_idx], 'mo', markersize=6, label='Max Height')
        
        ax1.set_xlabel('Horizontal Position (m)')
        ax1.set_ylabel('Vertical Position (m)')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        ax1.set_title(f'Trajectory (Range: {info["range"]:.1f}m, Max Height: {info["max_height"]:.1f}m)')
        ax1.set_ylim(bottom=0)  # Ensure ground is at y=0
        
        # Speed vs Time
        ax2.plot(data['time'], data['speed'], 'r-', linewidth=2)
        ax2.set_xlabel('Time (s)')
        ax2.set_ylabel('Speed (m/s)')
        ax2.grid(True, alpha=0.3)
        ax2.set_title(f'Speed vs Time (Flight Time: {info["flight_time"]:.1f}s)')
        
        # Velocity components
        ax3.plot(data['time'], data['vx'], 'g-', linewidth=2, label=f'Vx (const: {info["initial_velocity_x"]:.1f})')
        ax3.plot(data['time'], data['vy'], 'm-', linewidth=2, label=f'Vy (initial: {info["initial_velocity_y"]:.1f})')
        ax3.axhline(y=0, color='k', linestyle='--', alpha=0.5)
        ax3.set_xlabel('Time (s)')
        ax3.set_ylabel('Velocity (m/s)')
        ax3.grid(True, alpha=0.3)
        ax3.legend()
        ax3.set_title('Velocity Components vs Time')
        
        # Position components
        ax4.plot(data['time'], data['x'], 'c-', linewidth=2, label='X position')
        ax4.plot(data['time'], data['y'], 'orange', linewidth=2, label='Y position')
        ax4.axhline(y=0, color='k', linestyle='--', alpha=0.5, label='Ground level')
        ax4.set_xlabel('Time (s)')
        ax4.set_ylabel('Position (m)')
        ax4.grid(True, alpha=0.3)
        ax4.legend()
        ax4.set_title('Position Components vs Time')
        
        plt.tight_layout()
        return fig

# Example usage and test cases
if __name__ == "__main__":
    # Example 1: Constant acceleration (car accelerating)
    car_motion = Motion1D(initial_position=0, initial_velocity=5, acceleration=2)
    
    # Example 2: Free fall
    free_fall = Motion1D(initial_position=100, initial_velocity=0, acceleration=-9.81)
    
    # Example 3: Classic projectile motion (45-degree launch from ground)
    projectile_45 = Motion2D(
        launch_speed=25,      # m/s
        launch_angle=45,      # degrees
        launch_height=0,      # meters (ground level)
        launch_x=0
    )
    
    # Example 4: Projectile launched from height at 30 degrees
    projectile_height = Motion2D(
        launch_speed=20,      # m/s
        launch_angle=30,      # degrees
        launch_height=10,     # meters
        launch_x=0
    )
    
    # Example 5: Horizontal launch (like dropping a ball while moving)
    horizontal_launch = Motion2D(
        launch_speed=15,      # m/s
        launch_angle=0,       # degrees (horizontal)
        launch_height=20,     # meters
        launch_x=0
    )
    
    # Create visualizations
    visualizer = KinematicsVisualizer()
    
    # Plot 1D examples
    visualizer.plot_1d_motion(car_motion, 10, "Car Acceleration")
    visualizer.plot_1d_motion(free_fall, 4.5, "Free Fall Motion")
    
    # Plot 2D examples with different launch conditions
    visualizer.plot_2d_trajectory(projectile_45, title="45° Launch from Ground")
    visualizer.plot_2d_trajectory(projectile_height, title="30° Launch from 10m Height")
    visualizer.plot_2d_trajectory(horizontal_launch, title="Horizontal Launch from 20m")
    
    # Print launch information for educational purposes
    print("\n=== LAUNCH ANALYSIS ===")
    for name, motion in [("45° Ground Launch", projectile_45), 
                        ("30° Height Launch", projectile_height),
                        ("Horizontal Launch", horizontal_launch)]:
        info = motion.get_launch_info()
        print(f"\n{name}:")
        print(f"  Launch Speed: {info['launch_speed']:.1f} m/s at {info['launch_angle']:.1f}°")
        print(f"  Initial Velocity: ({info['initial_velocity_x']:.1f}, {info['initial_velocity_y']:.1f}) m/s")
        print(f"  Flight Time: {info['flight_time']:.2f} s")
        print(f"  Range: {info['range']:.1f} m")
        print(f"  Max Height: {info['max_height']:.1f} m")
    
    plt.show()