##########################################################################
# program: mdof_initial.py
# author: Tom Irvine
# version: 1.4
# date: December 16, 2011
# description:  multi-degree-of-freedom system subjected to initial
#               condition excitation using odeint
#
##########################################################################

from sys import stdin

from scipy.integrate import odeint
from scipy.linalg import solve

from generalized_eigen import generalized_eigen

from tompy import read_array,GetInteger2,GetInteger3,enter_float
from tompy import WriteArray

import numpy as np

import matplotlib.pyplot as plt

##########################################################################

def calc_deri(z, time, M,C,K):
   n=M.shape[0]
   tmp=np.zeros(2*n,float)
   tmp[:n] = z[n:]

   tmp[n:] = solve(M,-np.dot(C,z[n:])-np.dot(K,z[:n]))
   return tmp


##########################################################################

print("  ")
print(" Select units:      ")
print(" 1=English 2=metric ")

iu=GetInteger2()

print("  ")
if(iu==1):
    print(" units: mass (lbm)")
    print("        damping coefficients (lbf sec/in)")
    print("        stiffness (lbf/in)")
    print("        velocity (in/sec)")
    print("        displacement (in)")
else:
    print(" units: mass (kg)")
    print("        damping coefficients (N sec/m)")
    print("        stiffness (N/m)")
    print("        velocity (m/sec)")
    print("        displacement (m)")



print " "
M=read_array("mass matrix")

if(iu==1):
    M/=386
    

print("  ")

print " Select damping input method "
print "  1=damping coefficient matrix "
print "  2=uniform modal damping "
print "  3=modal damping vector"

idamp=GetInteger3()


if(idamp==1):
    C=read_array("damping coefficient matrix")

if(idamp==2):    
    print " Enter uniform damping ratio "
    damp=enter_float()

if(idamp==3):
    dampV=read_array(" Enter damping ratio vector")



K=read_array("stiffness matrix")


print(" ")


ndof,fn,MS=generalized_eigen(K,M,1)

MST=MS.T

num=ndof


if(idamp==2 or idamp==3):
    ccc=np.zeros((ndof,ndof),float)

    for i in range(0,ndof):
        if(idamp==2):
            ccc[i,i]=2*damp*(2*np.pi*fn[i])
        else:
            ccc[i,i]=2*dampV[i]*(2*np.pi*fn[i])        
        
    s1=np.dot(MST,M)
    s2=np.dot(ccc,s1)
    s3=np.dot(MS,s2)
    C=np.dot(M,s3)
    print(" ")
    print(" C= ")
    print C
    print(" ")
        
##########################################################################


max_fn=max(fn)
min_fn=min(fn)

if(min_fn>0):
    period_max=1/min_fn
else:
    period_min=1/max_fn

period_min=1/max_fn


dur=20*period_max
dt=period_min/50
NT=int(round(dur/dt))


time_vec = np.linspace(0, dur, NT)


v0=read_array("initial velocity vector")
d0=read_array("initial displacement vector")


z_0 = np.zeros(2*num,float) # State vector
z_0[:num] = d0
z_0[num:] = v0


Y = odeint(calc_deri, z_0, time_vec, args=(M,C,K))


YV=np.zeros((NT,num),float)
YD=np.zeros((NT,num),float)

plt.figure(1)
for i in range(0,num):
    dof='dof '+str(i+1)
    plt.plot(time_vec, Y[:,i], label=dof)
    YD[:,i]=Y[:,i]

plt.xlabel('Time (sec)')

plt.grid(True)
plt.title('Displacement')
plt.legend(loc="upper right")

if(iu==1):
    plt.ylabel('Disp (in)')
else:
    plt.ylabel('Disp (m)')

plt.figure(2)
for i in range(0,num):
    dof='dof '+str(i+1)
    plt.plot(time_vec, Y[:,num+i], label=dof)
    YV[:,i]=Y[:,num+i]

plt.xlabel('Time (sec)')

plt.grid(True)
plt.title('Velocity')
plt.legend(loc="upper right")

if(iu==1):
    plt.ylabel('Vel (in/sec)')
else:
    plt.ylabel('Vel (m/sec)')


AT=np.zeros((num,NT),float)
A=np.zeros((NT,num),float)

YVT=YV.T
YDT=YD.T



for i in range (0,NT):
    AT[:,i]=solve(M, -np.dot(C,YVT[:,i]) -np.dot(K,YDT[:,i]))

A=AT.T



plt.figure(3)
for i in range(0,num):
    dof='dof '+str(i+1)

    if(iu==1):
        A[:,i]/=386.
    else:
        A[:,i]/=9.81

    plt.plot(time_vec, A[:,i], label=dof)

plt.xlabel('Time (sec)')

plt.grid(True)
plt.title('Acceleration')
plt.legend(loc="upper right")
plt.ylabel('Accel (G)')


print " "
print " Write output data to files? "
print " 1=yes 2=no"

iwd=GetInteger2()

np1=num+1


if(iwd==1):

    dw=np.zeros((NT,np1),float)
    vw=np.zeros((NT,np1),float)
    aw=np.zeros((NT,np1),float)

    for i in range(0,NT):
        dw[i,0]=time_vec[i]
        vw[i,0]=time_vec[i]   
        aw[i,0]=time_vec[i] 
     
        for j in range(0,num):
            dw[i,j+1]=YD[i,j]
            vw[i,j+1]=YV[i,j]    
            aw[i,j+1]=A[i,j]        
     

    print " "
    print "Enter the output displacement filename: "
    output_file_path =stdin.readline()
    file_path = output_file_path.rstrip('\n') 
    WriteArray(dw,output_file_path)
    

    print " "
    print "Enter the output velocity filename: "
    output_file_path =stdin.readline()
    file_path = output_file_path.rstrip('\n')
    WriteArray(vw,output_file_path)


    print " "
    print "Enter the output acceleration filename: "
    output_file_path =stdin.readline()
    file_path = output_file_path.rstrip('\n')
    WriteArray(aw,output_file_path)


print " view plots"

plt.show()