import numpy as np
from autograd import jacobian
from autograd.numpy import cos, sin

class ControlSpace:
    """
    A generic state space with states x and control variable u.  
    The primary function of this class is to define the transformation from the current 
    state x to the next state xnext by applying a control u.
    """
    def __init__(self):
        pass
    def state_dimension(self):
        """Returns the dimension of the state space associated with this."""
        raise NotImplementedError()
    def control_dimension(self):
        """Returns the set from which controls should be sampled"""
        raise NotImplementedError()
    def next_state(self,x,u):
        """Produce the next state after applying control u to state x"""
        raise NotImplementedError()
    def interpolator(self,x,u):
        """Returns the interpolator (e.g., a trajectory) that goes from x to
        self.nextState(x,u)."""
        raise NotImplementedError()
    def connection(self,x,y):
        """Returns a sequence of controls that connects x to y, if
        applicable.  Return None if no such sequence exists."""
        return None
    def linearize(self, x_bar, u_bar, dt):
        """Compute linearized dynamics around (x_bar, u_bar). Returns A_lin, B_lin, C_lin."""
        raise NotImplementedError()

class Roomba(ControlSpace):
    """
    Class that describes the dynamics of a roomba-like robot.
    """
    def __init__(self, settings):
        self.max_speed = settings['Roomba']['max_speed']
        self.max_acc = settings['Roomba']['max_acc']
        self.max_d_acc = settings['Roomba']['max_d_acc']
        self.max_steer = np.radians(settings['Roomba']['max_steer'])
        self.max_d_steer = np.radians(settings['Roomba']['max_d_steer'])
        self.radius = settings['Roomba']['radius']  

    def state_dimension(self):
        # State is [x, y, theta]
        return 3

    def control_dimension(self):
        # Control is [v, steering_angle]
        return 2

    def derivative(self, x, u):
        """
        Compute the derivative of the state x with respect to control u.
        State x: [x, y, theta]
        Control u: [v, steering_angle]
        """

        x_pos, y_pos, theta = x
        v, steering_angle = u


        # Dynamics equations
        dx_pos = v * cos(theta)
        dy_pos = v * sin(theta)
        dtheta = steering_angle

        return np.array([dx_pos, dy_pos, dtheta])

    def next_state(self, x, u, dt):
        """
        Compute the next state given the current state and control input.
        """
        dxdt = u[0] * np.cos(x[2])
        dydt = u[0] * np.sin(x[2])
        dthetadt = u[1]
        
        # Update state using Euler integration
        new_x = x[0] + dxdt * dt
        new_y = x[1] + dydt * dt
        new_theta = x[2] + dthetadt * dt

        # Return the predicted next state
        return np.array([new_x, new_y, new_theta])

    def linearize(self, x_bar, u_bar, dt):
        """
        Linearize the system dynamics around (x_bar, u_bar) and discretize them.

        Returns:
            A_lin, B_lin, C_lin: Discrete-time linearized state-space matrices
        """
        x_bar = np.array(x_bar, dtype=np.float64)
        u_bar = np.array(u_bar, dtype=np.float64)

        # Compute continuous-time Jacobians
        A_continuous, B_continuous = self.compute_jacobian(x_bar, u_bar)
        
        # Discretize the system using Euler approximation
        A_lin = np.eye(self.state_dimension()) + A_continuous * dt
        B_lin = B_continuous * dt

        # Compute C_lin
        f_xu = self.derivative(x_bar, u_bar)
        C_lin = dt * (f_xu - A_continuous @ x_bar - B_continuous @ u_bar)

        return A_lin, B_lin, C_lin

    # def compute_jacobian(self, x_bar, u_bar):
    #     """
    #     Compute the Jacobian matrices A and B at (x_bar, u_bar).

    #     Returns:
    #         A: Jacobian of dynamics with respect to the state x
    #         B: Jacobian of dynamics with respect to the control input u
    #     """
    #     # Use autograd's jacobian function
    #     A_func = jacobian(lambda x: self.derivative(x, u_bar))
    #     B_func = jacobian(lambda u: self.derivative(x_bar, u))

    #     # Evaluate Jacobians at (x_bar, u_bar)
    #     A = A_func(x_bar)
    #     B = B_func(u_bar)

    #     return A, B
    
    def compute_jacobian(self, x, u):
        """
        Compute the Jacobian matrices A and B based on analytical derivatives.
        
        Parameters:
        - x: state vector [x, y, theta]
        - u: control input vector [v, omega]
        
        Returns:
        - A: Jacobian of dynamics with respect to the state x
        - B: Jacobian of dynamics with respect to the control input u
        """
        x_pos, y_pos, theta = x
        v, omega = u

        # Compute A (Jacobian with respect to state)
        A = np.zeros((self.state_dimension(), self.state_dimension()))
        A[0, 2] = -v * np.sin(theta)
        A[1, 2] = v * np.cos(theta)

        # Compute B (Jacobian with respect to control)
        B = np.zeros((self.state_dimension(), self.control_dimension()))
        B[0, 0] = np.cos(theta)
        B[1, 0] = np.sin(theta)
        B[2, 1] = 1.0

        return A, B