'''2nd Order step response
    my'' +cy' +ky = f
'''
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

def eq_of_motion(w, t, p):
    """
    Defines the differential equations for the direct-force mass-spring-damper system.

    Arguments:
        w :  vector of the state variables:
        t :  time
        p :  vector of the parameters:
    """
    x, x_dot = w
    m, k, c, StartTime, F_amp = p

    # Create sysODE = (x', x_dot')
    sysODE = [x_dot,
              -k/m * x - c/m * x_dot + f(t, p)/m]
    return sysODE


def f(t, p):
    """
    defines the disturbance force input to the system
    """
    m, k, c, StartTime, F_amp = p
    
    # Select one of the two inputs below
    # Be sure to comment out the one you're not using
    
    # Input Option 1: 
    #    Just a step in force beginning at t=DistStart
    # f = F_amp * (t >= DistStart)
    
    # Input Option 2:
    #    A pulse in force beginning at t=StartTime and ending at t=(StartTime + 0.5)
    f = F_amp * (t >= StartTime) * (t <= StartTime + 0.5) + 5 * (t >= StartTime+0.25) * (t <= StartTime + 0.5)
    
    return f

# Define the System Parameters
m = 1.0                     # kg
k = (2.0 * np.pi)**2        # N/m (Selected to give an undamped natrual frequency of 1Hz)
wn = np.sqrt(k / m)         # Natural Frequency (rad/s)

z = 0.1                     # Define a desired damping ratio
c = 2 * z * wn * m          # calculate the damping coeff. to create it (N/(m/s))

# Set up simulation parameters

# ODE solver parameters
abserr = 1.0e-9
relerr = 1.0e-9
max_step = 0.01
stoptime = 10.0
numpoints = 10001

# Create the time samples for the output of the ODE solver
t = np.linspace(0.0, stoptime, numpoints)

# Initial conditions
x_init = 0.0                        # initial position
x_dot_init = 0.0                    # initial velocity

# Set up the parameters for the input function
StartTime = 0.5              # Time the f(t) input will begin
F_amp = 10.0                # Amplitude of Disturbance force (N)

# Pack the parameters and initial conditions into arrays 
p = [m, k, c, StartTime, F_amp]
x0 = [x_init, x_dot_init]

# Call the ODE solver.
resp = odeint(eq_of_motion, x0, t, args=(p,), atol=abserr, rtol=relerr,  hmax=max_step)

# Make the figure pretty, then plot the results
#   "pretty" parameters selected based on pdf output, not screen output
#   Many of these setting could also be made default by the .matplotlibrc file

# Set the plot size - 3x2 aspect ratio is best
fig = plt.figure(figsize=(6, 4))
ax = plt.gca()
plt.subplots_adjust(bottom=0.17, left=0.17, top=0.96, right=0.96)

# Change the axis units to serif
plt.setp(ax.get_ymajorticklabels(),family='serif',fontsize=18)
plt.setp(ax.get_xmajorticklabels(),family='serif',fontsize=18)

ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')

ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')

# Turn on the plot grid and set appropriate linestyle and color
ax.grid(True,linestyle=':',color='0.75')
ax.set_axisbelow(True)

# Define the X and Y axis labels
plt.xlabel('Time (s)', family='serif', fontsize=22, weight='bold', labelpad=5)
plt.ylabel('Position', family='serif', fontsize=22, weight='bold', labelpad=10)

# Plot the first element of resp for all time. It corresponds to the position.
plt.plot(t, resp[:,0], linewidth=2, linestyle = '-', label=r'Response')

# uncomment below and set limits if needed
# xlim(0,5)
# ylim(0,10)

# # Create the legend, then fix the fontsize
# leg = plt.legend(loc='upper right', fancybox=True)
# ltext  = leg.get_texts()
# plt.setp(ltext,family='serif',fontsize=18)

# Adjust the page layout filling the page using the new tight_layout command
plt.tight_layout(pad = 0.5)

# save the figure as a high-res pdf in the current folder
# It's saved at the original 6x4 size
# plt.savefig('MCHE485_DirectForcePulseWithDamping.pdf')

fig.set_size_inches(9, 6) # Resize the figure for better display in the notebook
plt.show()
