import numpy as np
import pydrake.symbolic as ps
import torch
import time

from irs_lqr.dynamical_system import DynamicalSystem

class BicycleDynamics(DynamicalSystem):
    def __init__(self, h):
        super().__init__()
        """
        x = [x pos, y pos, heading, speed, steering_angle]
        u = [acceleration, steering_velocity]
        """
        self.h = h
        self.dim_x = 5
        self.dim_u = 2

        """Jacobian computations"""
        self.x_sym = np.array([ps.Variable("x_{}".format(i)) for i in range(self.dim_x)])
        self.u_sym = np.array([ps.Variable("u_{}".format(i)) for i in range(self.dim_u)])
        self.f_sym = self.dynamics_sym(self.x_sym, self.u_sym)

        self.jacobian_xu_sym = ps.Jacobian(self.f_sym, np.hstack((self.x_sym, self.u_sym)))
        
    def dynamics_sym(self, x, u):
        """
        Symbolic expression for dynamics. Used to compute
        linearizations of the system.
        x (np.array, dim: n): state
        u (np.array, dim: m): action        
        """
        heading = x[2]
        v = x[3]
        steer = x[4]
        dxdt = np.array([
            v * ps.cos(heading),
            v * ps.sin(heading),
            v * ps.tan(steer),
            u[0],
            u[1]
        ])
        x_new = x + self.h * dxdt
        return x_new


    def dynamics(self, x, u):
        """
        Numeric expression for dynamics.
        x (np.array, dim: n): state
        u (np.array, dim: m): action
        """
        heading = x[2]
        v = x[3]
        steer = x[4] 
        dxdt = np.array([
            v * np.cos(heading),
            v * np.sin(heading),
            v * np.tan(steer),
            u[0],
            u[1]
        ])
        x_new = x + self.h * dxdt
        return x_new

    def dynamics_batch(self, x, u):
        """
        Batch dynamics. Uses pytorch for 
        -args:
            x (np.array, dim: B x n): batched state
            u (np.array, dim: B x m): batched input
        -returns:
            xnext (np.array, dim: B x n): batched next state
        """
        heading = x[:,2]
        v = x[:,3]
        steer = x[:,4]
        dxdt = np.vstack((
            v * np.cos(heading),
            v * np.sin(heading),
            v * np.tan(steer),
            u[:,0],
            u[:,1]
        )).transpose()
        x_new = x + self.h * dxdt
        return x_new


    def dynamics_batch_torch(self, x, u):
        """
        Batch dynamics. Uses pytorch for 
        -args:
            x (np.array, dim: B x n): batched state
            u (np.array, dim: B x m): batched input
        -returns:
            xnext (np.array, dim: B x n): batched next state
        """
        x = torch.Tensor(x).cuda()
        u = torch.Tensor(u).cuda()

        heading = x[:,2]
        v = x[:,3]
        steer = x[:,4]

        dxdt = torch.vstack((
            v * torch.cos(heading),
            v * torch.sin(heading),
            v * torch.tan(steer),
            u[:,0],
            u[:,1]
        )).T
        x_new = x + self.h * dxdt
        return x_new

    def jacobian_xu(self, x, u):
        """
        Recoever linearized dynamics dfdx as a function of x, u
        """
        env = {self.x_sym[i]: x[i] for i in range(self.dim_x)}
        env.update({self.u_sym[i]: u[i] for i in range(self.dim_u)})
        f_x = ps.Evaluate(self.jacobian_xu_sym, env)
        return f_x 

    def jacobian_xu_batch(self, x, u):
        """
        Recoever linearized dynamics dfd(xu) as a function of x, u
        """ 
        dxdu_batch = np.zeros((
            x.shape[0], x.shape[1], x.shape[1] + u.shape[1]))
        for i in range(x.shape[0]):
            dxdu_batch[i] = self.jacobian_xu(x[i], u[i])
        return dxdu_batch        

[((121, 14, 121, 52), 'pydrake.symbolic.Evaluate', 'ps.Evaluate', ({(121, 26, 121, 46): 'self.jacobian_xu_sym', (121, 48, 121, 51): 'env'}, {}), '(self.jacobian_xu_sym, env)', True, 'import pydrake.symbolic as ps\n'), ((128, 21, 129, 61), 'numpy.zeros', 'np.zeros', ({(128, 30, 129, 60): '(x.shape[0], x.shape[1], x.shape[1] + u.shape[1])'}, {}), '((x.shape[0], x.shape[1], x.shape[1] + u.shape[1]))', True, 'import numpy as np\n'), ((24, 55, 24, 90), 'numpy.hstack', 'np.hstack', ({(24, 65, 24, 89): '(self.x_sym, self.u_sym)'}, {}), '((self.x_sym, self.u_sym))', True, 'import numpy as np\n'), ((98, 12, 98, 27), 'torch.Tensor', 'torch.Tensor', ({(98, 25, 98, 26): 'x'}, {}), '(x)', False, 'import torch\n'), ((99, 12, 99, 27), 'torch.Tensor', 'torch.Tensor', ({(99, 25, 99, 26): 'u'}, {}), '(u)', False, 'import torch\n'), ((37, 16, 37, 31), 'pydrake.symbolic.cos', 'ps.cos', ({(37, 23, 37, 30): 'heading'}, {}), '(heading)', True, 'import pydrake.symbolic as ps\n'), ((38, 16, 38, 31), 'pydrake.symbolic.sin', 'ps.sin', ({(38, 23, 38, 30): 'heading'}, {}), '(heading)', True, 'import pydrake.symbolic as ps\n'), ((39, 16, 39, 29), 'pydrake.symbolic.tan', 'ps.tan', ({(39, 23, 39, 28): 'steer'}, {}), '(steer)', True, 'import pydrake.symbolic as ps\n'), ((57, 16, 57, 31), 'numpy.cos', 'np.cos', ({(57, 23, 57, 30): 'heading'}, {}), '(heading)', True, 'import numpy as np\n'), ((58, 16, 58, 31), 'numpy.sin', 'np.sin', ({(58, 23, 58, 30): 'heading'}, {}), '(heading)', True, 'import numpy as np\n'), ((59, 16, 59, 29), 'numpy.tan', 'np.tan', ({(59, 23, 59, 28): 'steer'}, {}), '(steer)', True, 'import numpy as np\n'), ((106, 16, 106, 34), 'torch.cos', 'torch.cos', ({(106, 26, 106, 33): 'heading'}, {}), '(heading)', False, 'import torch\n'), ((107, 16, 107, 34), 'torch.sin', 'torch.sin', ({(107, 26, 107, 33): 'heading'}, {}), '(heading)', False, 'import torch\n'), ((108, 16, 108, 32), 'torch.tan', 'torch.tan', ({(108, 26, 108, 31): 'steer'}, {}), '(steer)', False, 'import torch\n'), ((79, 16, 79, 31), 'numpy.cos', 'np.cos', ({(79, 23, 79, 30): 'heading'}, {}), '(heading)', True, 'import numpy as np\n'), ((80, 16, 80, 31), 'numpy.sin', 'np.sin', ({(80, 23, 80, 30): 'heading'}, {}), '(heading)', True, 'import numpy as np\n'), ((81, 16, 81, 29), 'numpy.tan', 'np.tan', ({(81, 23, 81, 28): 'steer'}, {}), '(steer)', True, 'import numpy as np\n')]