# -*- coding: utf-8 -*-
"""
Created on Thu Dec  2 19:02:13 2021
@author: dhain460
SDOF - Classic Solution
"""
from numpy import arange, sqrt, exp, cos, sin, array
from numpy.linalg import inv
from matplotlib import pyplot as plt
#import matplotlib.pyplot as plt

# Variables
m = 200                   # mass
k = 1000                  # stiffness
c = 100                   # damping

omega = sqrt(k / m)      # natural frequency
cc = 2 * m * omega       # critical damping
Xi = c / cc              # damping ratio
print('Xi=',Xi)

# Set initial conditions
X0 = 5.0
V0 = 1.0

F0=0

# Time
dt = 0.001
t = arange(0, 20, dt)

# Analytical Solution

if Xi == 1:
    print("Critically damped system")
    gamma = omega * Xi
    c1 = X0
    c2 = V0 + omega * X0
    x = (c1 + c2 * t)*exp(-gamma * t)
    dx = -gamma * exp(-gamma * t) * (c1 + c2*t) + c2 * exp(-gamma * t)
elif Xi > 1:
    print("Overdamped system")
    gamma1 = omega * (Xi + sqrt(Xi **2 - 1))
    gamma2 = omega * (Xi - sqrt(Xi **2 - 1))
    c1 = -(X0*omega*(Xi-sqrt(Xi**2-1)+V0))/(2*omega*sqrt(Xi**2-1))
    c2 = -c1
    x = c1 * exp(-gamma1 * t) +  c2 * exp(-gamma2 * t)
    dx = -gamma1*c1*exp(-gamma1*t)-gamma2*c1*exp(-gamma2*t)
elif Xi < 1:
    print("Underdamped system")
    omega_d = sqrt(1-Xi**2) * omega
    c1 = X0
    c2 = (V0+Xi*omega*X0)/omega_d
    x = exp(-Xi*omega*t)*(c1*cos(omega_d*t)+c2*sin(omega_d*t))
    dx = -Xi*omega*exp(-Xi*omega*t)*(c1*cos(omega_d*t)+c2*sin(omega_d*t)) \
       +exp(-Xi*omega*t)*(-c1*sin(omega_d*t)*omega_d+c2*cos(omega_d*t)*omega_d)
else:
    print("Check code")

#print('x=',len(x)) 
#print('t=',len(t))  

# Numerical Solution - Explicit Solution / Time Stepping

A = array([[m,0],[0,1]])
print(A)
B = array([[c,k],[-1,0]])
print(B)
F = array([F0,0])
print(F)

y = array([V0,X0])
Y = []
V = []

# Time-stepping solution (Explicit Numerical Solution)
# Problem with explicit solution is to find minimal time step
# to get an accurate solution
for time in t:
    F[0] = F0 * cos(omega*time)
    y = y + dt * inv(A).dot(F-B.dot(y))
    Y.append(y[1])
    V.append([y[0]])
    
print('y=',len(Y)) 
print('t=',len(t))  
 
# Plot analytical solutions
plt.clf()
plt.figure(1)
fig,[fig1,fig2] = plt.subplots(1,2,figsize=(12,7))
fig.suptitle('Analytical Solution')
fig1.plot(t,x,'r--')
fig1.set_xlabel('t (s)')
fig1.set_ylabel('x(t)')
fig1.grid('b')

#plt.figure(2)
fig2.plot(x,dx,'r--')
fig2.set_xlabel('x')
fig2.set_ylabel('v = dx/dt')
fig2.grid('b')
plt.show()

# Plot numerical solutions
plt.clf()
plt.figure(2)
fig,[fig1,fig2] = plt.subplots(1,2,figsize=(12,7))
fig.suptitle('Time-Stepping Solution')
fig1.plot(t,Y,'b-')
fig1.set_xlabel('t (s)')
fig1.set_ylabel('x(t)')
fig1.grid('b')

#plt.figure(2)
fig2.plot(Y,V,'b-')
fig2.set_xlabel('x')
fig2.set_ylabel('v = dx/dt')
fig2.grid('b')
plt.show()

# Plotting together analytical and numerical (explicit)
plt.clf()
plt.figure(3)
fig,[fig1,fig2] = plt.subplots(1,2,figsize=(12,7))
fig.suptitle('Analytical & Time-Stepping Solution')
fig1.plot(t,Y,'b-',label='Time-Stepping')
fig1.legend(loc='upper right')
fig1.plot(t,x,'r--',label='Analytical')
fig1.legend(loc='upper right')
fig1.set_xlabel('t (s)')
fig1.set_ylabel('x(t)')
fig1.grid('b')
fig1.legend(loc='upper right')

#plt.figure(2)
fig2.plot(Y,V,'b-')
fig2.plot(x,dx,'r--')
fig2.set_xlabel('x')
fig2.set_ylabel('v = dx/dt')
fig2.grid('b')
plt.show()
