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        

[((3342, 3380), 'pydrake.symbolic.Evaluate', 'ps.Evaluate', (['self.jacobian_xu_sym', 'env'], {}), '(self.jacobian_xu_sym, env)\n', (3353, 3380), True, 'import pydrake.symbolic as ps\n'), ((3554, 3613), 'numpy.zeros', 'np.zeros', (['(x.shape[0], x.shape[1], x.shape[1] + u.shape[1])'], {}), '((x.shape[0], x.shape[1], x.shape[1] + u.shape[1]))\n', (3562, 3613), True, 'import numpy as np\n'), ((753, 788), 'numpy.hstack', 'np.hstack', (['(self.x_sym, self.u_sym)'], {}), '((self.x_sym, self.u_sym))\n', (762, 788), True, 'import numpy as np\n'), ((2704, 2719), 'torch.Tensor', 'torch.Tensor', (['x'], {}), '(x)\n', (2716, 2719), False, 'import torch\n'), ((2739, 2754), 'torch.Tensor', 'torch.Tensor', (['u'], {}), '(u)\n', (2751, 2754), False, 'import torch\n'), ((1137, 1152), 'pydrake.symbolic.cos', 'ps.cos', (['heading'], {}), '(heading)\n', (1143, 1152), True, 'import pydrake.symbolic as ps\n'), ((1170, 1185), 'pydrake.symbolic.sin', 'ps.sin', (['heading'], {}), '(heading)\n', (1176, 1185), True, 'import pydrake.symbolic as ps\n'), ((1203, 1216), 'pydrake.symbolic.tan', 'ps.tan', (['steer'], {}), '(steer)\n', (1209, 1216), True, 'import pydrake.symbolic as ps\n'), ((1593, 1608), 'numpy.cos', 'np.cos', (['heading'], {}), '(heading)\n', (1599, 1608), True, 'import numpy as np\n'), ((1626, 1641), 'numpy.sin', 'np.sin', (['heading'], {}), '(heading)\n', (1632, 1641), True, 'import numpy as np\n'), ((1659, 1672), 'numpy.tan', 'np.tan', (['steer'], {}), '(steer)\n', (1665, 1672), True, 'import numpy as np\n'), ((2877, 2895), 'torch.cos', 'torch.cos', (['heading'], {}), '(heading)\n', (2886, 2895), False, 'import torch\n'), ((2913, 2931), 'torch.sin', 'torch.sin', (['heading'], {}), '(heading)\n', (2922, 2931), False, 'import torch\n'), ((2949, 2965), 'torch.tan', 'torch.tan', (['steer'], {}), '(steer)\n', (2958, 2965), False, 'import torch\n'), ((2186, 2201), 'numpy.cos', 'np.cos', (['heading'], {}), '(heading)\n', (2192, 2201), True, 'import numpy as np\n'), ((2219, 2234), 'numpy.sin', 'np.sin', (['heading'], {}), '(heading)\n', (2225, 2234), True, 'import numpy as np\n'), ((2252, 2265), 'numpy.tan', 'np.tan', (['steer'], {}), '(steer)\n', (2258, 2265), True, 'import numpy as np\n')]