import numpy as np
import math

class InverseDynamics:
    def __init__(self, max_steering=0.7, max_acc=8.0, length=4.5):
        """
        :param max_steering: Max steering angle in radians (approx 40 degrees)
        :param max_acc: Max acceleration in m/s^2
        :param length: Vehicle length in meters (Waymo default approx 4.5m)
        """
        self.max_steering = max_steering
        self.max_acc = max_acc
        self.wheelbase = 0.7 * length  # Approximation as per request

    def compute_action(self, current_state, next_state, dt=0.1):
        """
        Compute action [steering, acceleration] from current and next state.
        State format: dictionary or object with keys/attrs: position (x, y), heading, velocity (v_x, v_y)
        or numpy array [x, y, vx, vy, heading]
        
        Using Bicycle Model:
        delta = arctan(L * theta_dot / v)
        acc = (v_next - v_curr) / dt
        """
        
        # Extract state
        # Assume state is dict-like for now, can adapt if needed
        # We need: velocity (scalar), heading
        
        # Helper to get speed
        def get_speed(vel):
            return np.linalg.norm(vel)

        v_curr = get_speed(current_state['velocity'])
        v_next = get_speed(next_state['velocity'])
        
        # 1. Acceleration (longitudinal)
        acc = (v_next - v_curr) / dt
        
        # 2. Steering (lateral)
        # theta_dot = (theta_next - theta_curr) / dt
        theta_curr = current_state['heading']
        theta_next = next_state['heading']
        
        # Handle angle wrapping [-pi, pi]
        diff_theta = theta_next - theta_curr
        if diff_theta > np.pi:
            diff_theta -= 2 * np.pi
        elif diff_theta < -np.pi:
            diff_theta += 2 * np.pi
            
        theta_dot = diff_theta / dt
        
        # Avoid division by zero for stationary vehicles
        if v_curr < 0.1:
            steering = 0.0
        else:
            # delta = arctan(L * theta_dot / v)
            steering = np.arctan(self.wheelbase * theta_dot / v_curr)

        # Normalize actions to [-1, 1]
        norm_acc = np.clip(acc / self.max_acc, -1.0, 1.0)
        norm_steering = np.clip(steering / self.max_steering, -1.0, 1.0)

        return np.array([norm_steering, norm_acc]), {'raw_acc': acc, 'raw_steering': steering}

    def apply_action(self, current_state, action, dt=0.1):
        """
        Forward dynamics: given current_state and action [steering, acc] in [-1, 1], return next_state.
        State format: dict with position (x,y), heading, velocity (vx, vy).
        """
        steering_norm, acc_norm = float(action[0]), float(action[1])
        acc = acc_norm * self.max_acc
        steering = steering_norm * self.max_steering
        pos = np.array(current_state['position'][:2], dtype=np.float64)
        heading = float(current_state['heading'])
        vel = np.array(current_state['velocity'], dtype=np.float64)
        v = np.linalg.norm(vel)
        if v < 0.1:
            v = 0.1
        theta_dot = v * np.tan(steering) / self.wheelbase
        v_next = v + acc * dt
        v_next = max(0.0, v_next)
        heading_next = heading + theta_dot * dt
        heading_next = np.arctan2(np.sin(heading_next), np.cos(heading_next))
        vx_next = v_next * np.cos(heading_next)
        vy_next = v_next * np.sin(heading_next)
        pos_next = pos + dt * np.array([vx_next, vy_next])
        return {
            'position': pos_next,
            'heading': heading_next,
            'velocity': np.array([vx_next, vy_next]),
        }