from sympy import symbols
from sympy.physics.mechanics import *
from sympy import Dummy, lambdify
from numpy import array, hstack, zeros, linspace, pi, ones
from numpy.linalg import solve
from scipy.integrate import odeint

from pydy.viz.shapes import Cylinder, Sphere, Cube
import pydy.viz
from pydy.viz.visualization_frame import VisualizationFrame
from pydy.viz.scene import Scene

######################
# Equation of motion #
######################

n = 6

q = dynamicsymbols('q:' + str(n + 1))  # Generalized coordinates
u = dynamicsymbols('u:' + str(n + 1))  # Generalized speeds
f = dynamicsymbols('f')                # Force applied to the cart
    
m = symbols('m:' + str(n + 1))         # Mass of each bob
l = symbols('l:' + str(n))             # Length of each link
c_f = symbols('c_f:' + str(n + 1))
g, t = symbols('g t')                  # Gravity and time

I = ReferenceFrame('I')                # Inertial reference frame
O = Point('O')                         # Origin point
O.set_vel(I, 0)                        # Origin's velocity is zero

P0 = Point('P0')                       # Hinge point of top link
P0.set_pos(O, q[0] * I.x)              # Set the position of P0    
P0.set_vel(I, u[0] * I.x)              # Set the velocity of P0
Pa0 = Particle('Pa0', P0, m[0])        # Define a particle at P0

frames = [I]                              # List to hold the n + 1 frames
points = [P0]                             # List to hold the n + 1 points
particles = [Pa0]                         # List to hold the n + 1 particles
forces = [(P0, f * I.x - u[0] * c_f[0] * I.x - m[0] * g * I.y)] # List to hold the n + 1 applied forces, including the input force, f
kindiffs = [q[0].diff(t) - u[0]]          # List to hold kinematic ODE's

for i in range(n):
    Bi = I.orientnew('B' + str(i), 'Axis', [q[i + 1], I.z])   # Create a new frame
    Bi.set_ang_vel(I, u[i + 1] * I.z)                         # Set angular velocity
    frames.append(Bi)                                         # Add it to the frames list

    Pi = points[-1].locatenew('P' + str(i + 1), l[i] * Bi.x)  # Create a new point
    Pi.v2pt_theory(points[-1], I, Bi)                         # Set the velocity
    points.append(Pi)                                         # Add it to the points list
    
    Pai = Particle('Pa' + str(i + 1), Pi, m[i + 1])           # Create a new particle
    particles.append(Pai)                                     # Add it to the particles list

    forces.append((Pi, -u[i + 1]*c_f[i + 1]*I.x - m[i + 1] * g * I.y))                  # Set the force applied at the point
        
    kindiffs.append(q[i + 1].diff(t) - u[i + 1])              # Define the kinematic ODE:  dq_i / dt - u_i = 0

kane = KanesMethod(I, q_ind=q, u_ind=u, kd_eqs=kindiffs) # Initialize the object
fr, frstar = kane.kanes_equations(particles, forces)     # Generate EoM's fr + frstar = 0

##############
# Simulation #
##############

arm_length = 3. / n                          # The maximum length of the pendulum is 1 meter
bob_mass = 1. / n                          # The maximum mass of the bobs is 10 grams
coeff_fric_bob = 0.20 # friction coefficient
parameters = [g, m[0], c_f[0]]                       # Parameter definitions starting with gravity and the first bob
cart_mass = 1.0
coeff_fric_cart = 0.18
# parameter_vals = [9.81, 0.01 / n, coeff_fric]            # Numerical values for the first two
parameter_vals = [9.81, cart_mass, coeff_fric_cart]
for i in range(n):                           # Then each mass and length
    parameters += [l[i], m[i + 1], c_f[i + 1]]            
    parameter_vals += [arm_length, bob_mass, coeff_fric_bob]

dynamic = q + u                                                # Make a list of the states
dynamic.append(f)                                              # Add the input force
dummy_symbols = [Dummy() for i in dynamic]                     # Create a dummy symbol for each variable
dummy_dict = dict(zip(dynamic, dummy_symbols))                 
kindiff_dict = kane.kindiffdict()                              # Get the solved kinematical differential equations
M = kane.mass_matrix_full.subs(kindiff_dict).subs(dummy_dict)  # Substitute into the mass matrix 
F = kane.forcing_full.subs(kindiff_dict).subs(dummy_dict)      # Substitute into the forcing vector
M_func = lambdify(dummy_symbols + parameters, M)               # Create a callable function to evaluate the mass matrix 
F_func = lambdify(dummy_symbols + parameters, F)               # Create a callable function to evaluate the forcing vector 

def right_hand_side(x, t, args):
    """Returns the derivatives of the states.

    Parameters
    ----------
    x : ndarray, shape(2 * (n + 1))
        The current state vector.
    t : float
        The current time.
    args : ndarray
        The constants.

    Returns
    -------
    dx : ndarray, shape(2 * (n + 1))
        The derivative of the state.
    
    """
    u = 0.0                              # The input force is always zero     
    arguments = hstack((x, u, args))     # States, input, and parameters
    dx = array(solve(M_func(*arguments), # Solving for the derivatives
        F_func(*arguments))).T[0]
    
    return dx

x0 = hstack(( 0, pi / 2 * ones(len(q) - 1), 1e-3 * ones(len(u)) )) # Initial conditions, q and u
t = linspace(0, 10, 1000)                                          # Time vector
y = odeint(right_hand_side, x0, t, args=(parameter_vals,))         # Actual integration

#===========#
# Animation #
#===========#

cart_shape = Cube(color='lightcyan', length=0.3)
cart_viz_frame = VisualizationFrame(I, points[0], cart_shape)
visualization_frames = [cart_viz_frame]

for i in range(n):
    bob_tmp_shape = Sphere(color='lightcyan', radius=0.1)
    bob_tmp_viz_frame = VisualizationFrame(I, points[i + 1], bob_tmp_shape)
    visualization_frames.append(bob_tmp_viz_frame)

# bob1_shape = Sphere(color='lightcyan', radius=0.1)
# bob2_shape = Sphere(color='lightcyan', radius=0.1)

# bob1_viz_frame = VisualizationFrame(I, points[1], bob1_shape)
# bob2_viz_frame = VisualizationFrame(I, points[2], bob2_shape)

parameter_dict = dict(zip(parameters, parameter_vals))

# link1_shape = Cylinder(radius=0.08, length=arm_length, color='blue')
# link1_viz_frame = VisualizationFrame('Link 1', frames[1], points[1], link1_shape)
# link2_shape = Cylinder(radius=0.08, length=arm_length, color='blue')
# link2_viz_frame = VisualizationFrame('Link 2', frames[2], points[2], link2_shape)

scene = Scene(I, O)

# scene.visualization_frames = [cart_viz_frame, bob1_viz_frame, bob2_viz_frame] #, link1_viz_frame, link2_viz_frame]

scene.visualization_frames = visualization_frames

scene.states_symbols = q + u
scene.constants = parameter_dict
scene.states_trajectories = y

scene.display()
