##########################################################################
# program: mdof_arbit_force_rk4.py
# author: Tom Irvine
# version: 1.4
# date: December 30, 2011
# description:  Calculate the response of an MDOF system to an applied
#               force using the Runge-Kutta fourth order method.
#
#               The input file must have two columns: time(sec) & force
##########################################################################

from tompy import read_two_columns
from tompy import GetInteger2
from tompy import time_history_plot
from tompy import read_array,enter_int,enter_float

from ode_plots import ode_plots

from scipy.linalg import solve

from generalized_eigen import generalized_eigen

from numpy import zeros,dot,linspace,interp

import matplotlib.pyplot as plt

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

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")
C=read_array("damping coefficient matrix")
K=read_array("stiffness matrix")


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

print(" ")


ndof,fn,NN =generalized_eigen(K,M,2)

num=ndof

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

dtr=period_min/20

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

v0=read_array("initial velocity vector")

d0=read_array("initial displacement vector")

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

if(iu==1):
    print "Each force file must have two columns: time(sec) & force(lbf)"
else:
    print "Each force file must two columns: time(sec) & force(N)"


print " "
print " Enter the number of force files "

nff=enter_int()

MAX=100000


tt=zeros((MAX,nff),float)
ff=zeros((MAX,nff),float)
force_dof=zeros(ndof,int)

for i in range (0,ndof):
    force_dof[i]=-999



print "  "
print " Note: the first dof is 1 "


for i in range (0,nff):
    ii=i+1
    print " "
    print " Enter force file %d " % ii
    a,b,num=read_two_columns()
    L=len(a)
    if(L>MAX):
        L=MAX
    tt[0:L,i]=a[0:L]
    ff[0:L,i]=b[0:L]

    print " "
    print " Enter the number of dofs at which this force is applied"

    nfa=enter_int()


    for j in range (0,nfa):
        print " "

        if(j==0 and nfa==1):
            print " Enter the dof number for this force "

        if(j==0 and nfa>1):
            print " Enter the first dof number for this force "

        if(j>0 and nfa>1):
            print " Enter the next dof number for this force "

        nn=enter_int()
        nn=nn-1
        force_dof[nn]=i


#*******************************************************************************

print " "
print " Enter the duration (sec) "
dur=enter_float()

srr=20*max_fn

print " "
print " Enter the sample rate (recommend > %8.4g) " % srr
sr=enter_float()

dt=1/sr

nt=int(dur/dt)

t = linspace(0., dur, nt)

#*******************************************************************************
#
#  interpolate force
#

print " begin interpolation"

FFI=zeros((nt,nff),float)

for i in range (0,nff):
    tstart=tt[0,i]
    tt[:,i]=tt[:,i]-tstart
#    print max(t),max(tt[:,i])

    last=MAX
    for j in range (1,MAX):
        if(tt[j,i]<=tt[(j-1,i)]):
            last=j
            break

#    print last


    tint=tt[0:last,i]
    fint=ff[0:last,i]

    FFI[:,i] = interp(t,tint,fint)

print " end interpolation"

nrows=FFI.shape[0]
ncolumns=FFI.shape[1]

print nrows,ncolumns

#*******************************************************************************
#

x=zeros((nt,ndof),float)
y=zeros((nt,ndof),float)
a=zeros((nt,ndof),float)


x[i,:]=d0
y[i,:]=v0

nt_m1=nt-1

for i in range (0,nt_m1):

    ta=t[i]
    tb=t[i+1]

    h=tb-ta

    hd2=h/2

    FA=zeros(ndof,float)
    FB=zeros(ndof,float)

    FF1=zeros(ndof,float)
    FF2=zeros(ndof,float)
    FF3=zeros(ndof,float)
    FF4=zeros(ndof,float)


    for j in range (0,ndof):

        j_index=force_dof[j]

        if(j_index!=-999):

            FA[j]=FFI[i,j_index]
            FB[j]=FFI[i+1,j_index]

    FF1= FA
    FF2= (FB + FA)/2.
    FF3= FF2
    FF4= FB

    X1= x[i,:]
    Y1= y[i,:]
    A1=solve(M, FF1 -dot(C,Y1) -dot(K,X1))

    X2= X1+Y1*hd2
    Y2= Y1+A1*hd2
    A2=solve(M, FF2 -dot(C,Y2) -dot(K,X2))

    X3= X1+Y2*hd2
    Y3= Y1+A2*hd2
    A3=solve(M, FF3 -dot(C,Y3) -dot(K,X3))

    X4= X1+Y3*h
    Y4= Y1+A3*h
    A4=solve(M, FF4 -dot(C,Y4) -dot(K,X4))

    x[i+1,:]= X1+(h/6.)*( Y1 + 2*Y2 + 2*Y3 + Y4 )   # displacement
    y[i+1,:]= Y1+(h/6.)*( A1 + 2*A2 + 2*A3 + A4 )   # velocity


    a[i+1,:]=solve(M, FF4 -dot(C,y[i+1,:]) -dot(K,x[i+1,:]))


#*******************************************************************************


if(iu==1):
    a=a/386
else:
    a=a/9.81

fignum=1
ode_plots(fignum,nt,ndof,t,x,y,a,iu)

plt.show()