###########################################################################
# program: srs.py
# author: Tom Irvine
# Email: tom@vibrationdata.com
# version: 1.9
# date: May 20, 2013
# description:
#    
# Calculate the shock response spectrum for an SDOF system
# The file must have two columns: time(sec) & accel(G)
#
# The numerical engine is the Smallwood ramp invariant digital
# recursive filtering relationship
#
###########################################################################

from tompy import read_two_columns,signal_stats,sample_rate_check
from tompy import GetInteger2,GetInteger3,WriteData3
from tompy import enter_damping
from tompy import enter_float
from tompy import time_history_plot
from tompy import srs_plot_pn

from scipy.signal import lfilter
from numpy import zeros,concatenate
from sys import stdin
from math import pi,exp,sqrt,cos,sin

import matplotlib.pyplot as plt


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

#  
#  http://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.lfilter.html
#  
#  a[0]*y[n] = b[0]*x[n] + b[1]*x[n-1] + ... + b[nb]*x[n-nb]
#                        - a[1]*y[n-1] - ... - a[na]*y[n-na]
#


class SRS:

    def __init__(self,b,fn,damp,dt,num_fn):
    
        self.b = b
        self.fn=fn
        self.damp=damp
        self.dt=dt
        self.num_fn=num_fn
   
    @classmethod     
    def a_coeff(cls,omega,damp,dt): 
        ac=zeros(3)    
        ac[0]=1   
        omegad=omega*sqrt(1.-(damp**2))        
        E=exp(-damp*omega*dt)
        K=omegad*dt
        C=E*cos(K)
        ac[1]=-2*C
        ac[2]=+E**2    
        
        return ac
        
        
    def accel_SRS(self):
        
        pos=zeros(self.num_fn)
        neg=zeros(self.num_fn)
        bc=zeros(3)
        
        
        for j in range(0,self.num_fn):
            
            omega=2.*pi*self.fn[j]
            omegad=omega*sqrt(1.-(self.damp**2))
#
#  bc coefficients are applied to the excitation
            
            E=exp(-self.damp*omega*self.dt)
            K=omegad*self.dt
            C=E*cos(K)
            S=E*sin(K)
            Sp=S/K

            bc[0]=1.-Sp
            bc[1]=2.*(Sp-C)
            bc[2]=E**2-Sp
                        
            ac=SRS.a_coeff(omega,self.damp,self.dt)

      
            resp=lfilter(bc, ac, self.b, axis=-1, zi=None)            
#
            pos[j]= max(resp)
            neg[j]= abs(min(resp))    
            
        return pos,neg
        
    def rel_disp_SRS(self):
        
        rd_pos=zeros(self.num_fn)
        rd_neg=zeros(self.num_fn)
        ac=zeros(3)
        bc=zeros(3)
        
        for j in range(0,self.num_fn):
            
            omega=2.*pi*self.fn[j]
            omegad=omega*sqrt(1.-(self.damp**2))            
            
            E =exp(  -self.damp*omega*self.dt)
            E2=exp(-2*self.damp*omega*self.dt)
             
            K=omegad*self.dt
            C=E*cos(K)
            S=E*sin(K)
            
            Omr=(omega/omegad)
            Omt=omega*self.dt
            
            P=2*self.damp**2-1
            
            b00=2*self.damp*(C-1)
            b01=S*Omr*P
            b02=Omt
            
            b10=-2*Omt*C
            b11= 2*self.damp*(1-E2)
            b12=-2*b01   

            b20=(2*self.damp+Omt)*E2
            b21= b01
            b22=-2*self.damp*C               
            
            bc[0]=b00+b01+b02
            bc[1]=b10+b11+b12
            bc[2]=b20+b21+b22
            
            bc=-bc/(omega**3*self.dt)
                        
            ac=SRS.a_coeff(omega,self.damp,self.dt)
            
            resp=lfilter(bc, ac, b, axis=-1, zi=None)

            rd_pos[j]= 386*max(resp)
            rd_neg[j]= 386*abs(min(resp))   
            
        return rd_pos,rd_neg            

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

class READ_DATA:

    def __init__(self):
        pass

    @classmethod     
    def check_data(cls,a,b,num,sr,dt):

        sample_rate_check(a,b,num,sr,dt)
        
        return sr,dt
        
    def read_and_stats(self): 
 
        a,b,num =read_two_columns()

        sr,dt,ave,sd,rms,skew,kurtosis,dur=signal_stats(a,b,num)

        sr,dt=READ_DATA.check_data(a,b,num,sr,dt)

        return a,b,num,sr,dt,dur

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


print(" ")     

print " " 
print "SRS using the Smallwood ramp invariant digital recursive"
print "filtering relationship"
print " "
     
print "The file must have two columns: time(sec) & accel(G)"

a,b,num,sr,dt,dur=READ_DATA().read_and_stats()

print(" ")

damp,Q= enter_damping()

print(" ")
print(" Enter starting frequency (Hz) ")        
f1 = enter_float()

print(" ")
print(" Select octave spacing")
print("   1= one-third")
print("   2= one-sixth")
print("   3= one-twelfth")

noct=GetInteger3()
        
if(noct==1):
    oct=1./3.
if(noct==2):
    oct=1./6.
if(noct==3):
    oct=1./12.
    

print(" ")        
print("  include residual:  1=yes 2=no ")

ires = GetInteger2()
        
   
                
print(" ")
print(" Select output metrics")
print(" ")
print("  acceleration:  1=yes 2=no ")

iacc = GetInteger2()
        
print(" ")        
print("  pseudo velocity:  1=yes 2=no ")

ipv = GetInteger2()

print(" ")        
print("  relative displacement:  1=yes 2=no ")

ird = GetInteger2()

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

nt=int(0.8*(1/f1)/dt)

b_old=b
if(ires==1):
    tz=zeros(nt)
    d=concatenate((b,tz))
    b=d
#*******************************************************************************

resp=zeros(num)
fn=zeros(num)



fn[0]=f1    
for j in range(1,999):
    fn[j]=fn[j-1]*(2.**oct)
    if  fn[j] > sr/8.:
        break

num_fn=j

temp=fn[0:num_fn]
del(fn)
fn=temp
del(temp)

pv_pos=zeros(num_fn)
pv_neg=zeros(num_fn)
rd_pos=zeros(num_fn)
rd_neg=zeros(num_fn)


srs_data=SRS(b,fn,damp,dt,num_fn)


if(iacc==1):
    fp=zeros(num_fn)

    print " "
    print " fn(Hz)  Positive Accel(G)  Negative Accel(G) "

    x_pos,x_neg=srs_data.accel_SRS()

    for j in range(0,num_fn):
        print "%6.4g  %6.4g  %6.4g" % (fn[j],x_pos[j],x_neg[j])
        

if(ipv==1 or ird==1):
        
    rd_pos,rd_neg=srs_data.rel_disp_SRS()        
#    
#*******************************************************************************
#
#   Pseudo Velocity
#

if(ipv==1):
    for j in range(0,num_fn):
        omega=2.*pi*fn[j]
        pv_pos[j]=rd_pos[j]*omega
        pv_neg[j]=rd_neg[j]*omega   

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


if(iacc==1):
    print " "
    print "Enter the output acceleration SRS filename: "
    output_file_path = stdin.readline()
    output_file = output_file_path.rstrip('\n')
    WriteData3(num_fn,fn,x_pos,x_neg,output_file)

if(ipv==1):    
    print " "
    print "Enter the output pseudo velocity SRS filename: "
    output_file_path = stdin.readline()
    output_file = output_file_path.rstrip('\n')
    WriteData3(num_fn,fn,pv_pos,pv_neg,output_file)

if(ird==1):
    print " "
    print "Enter the output relative displacement SRS filename: "
    output_file_path = stdin.readline()
    output_file = output_file_path.rstrip('\n')
    WriteData3(num_fn,fn,rd_pos,rd_neg,output_file)

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

print " "
print " begin plots "
print " "

time_history_plot(a,b_old,1,'Time(sec)','Accel(G)','Base Input','time_history')

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

if(iacc==1):    
    plt.figure(2)
    srs_plot_pn(1,1,fn,x_pos,x_neg,damp,'accel_srs_plot')  

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

if(ipv==1):
    plt.figure(3)
    srs_plot_pn(2,1,fn,pv_pos,pv_neg,damp,'pv_srs_plot')     

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

if(ird==1):
    plt.figure(4)
    srs_plot_pn(3,1,fn,rd_pos,rd_neg,damp,'rd_srs_plot')  

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

plt.show()