import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from scipy.interpolate import griddata
import math
from scipy.signal import decimate, convolve

ND = np.load("nodelay.npz")['BigImgNDCpsd']
D = np.load("delay.npz")['BigImgCpsd']

plt.plot( D[0][0:1000]-256, label='Enveloppe of the signal')
plt.plot( D[40][0:1000]-256, label='Enveloppe of the signal') 
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
plt.show()

OffSets = []
DD = []
for l in range(60):
    Ref = ND[1][150:1500]-256
    Test = ND[l][400:600]-256
    ND[l][0]= ND[l][4]
    ND[l][1]= ND[l][4]
    ND[l][2]= ND[l][4]
    ND[l][3]= ND[l][4]
    Corr = signal.correlate(Ref/1.0, Test/1.0, mode='same') 
    ACorr = np.argmax(Corr) #
    OffSets.append(ACorr)
print OffSets[3]
OffSets[3] = 286
OffSets[5] = 286+10
OffSets[6] = 286+5
OffSets[8] = 286-80
OffSets[9] = 286+15
for l in range(60):
    zero_array = np.zeros([OffSets[l]])
    #zero_array = np.zeros([0])
    line = np.hstack([zero_array,ND[l]])[0:10000]
    DD.append(line)
OffMax = max(OffSets) 

m = 0
plt.imshow( np.abs(np.abs(np.transpose(DD)[0:4000])[0:1000]),aspect="auto")
plt.show()

166

for i in range(60):
    plt.plot(DD[i][0:800]) 
plt.show()

H = []
SH = []
Reconstructed = []
FFTMap = []
SFFTMap = []
Dec = []
SHDec = []
for ml in range(60): 
    SigTest = DD[ml][0:10000]
    FFT = np.fft.fft(SigTest)
    FStart = 1200
    FStop = 2700
    Ssig = []
    for p in range(10000/2):
        Ssig.append(SigTest[2*p]+SigTest[2*p+1])
    SFFT = np.fft.fft(Ssig)   
    for k in range(len(FFT)/2):
        if (k < FStart or k > FStop): # in (k < 550000 or k > 790000) # 0.068 0.196
            FFT[k] = 0
            FFT[-k] = 0
    tmp = np.fft.fft(SigTest)
    for k in range(50):
        tmp[k] = 0
        tmp[-k] = 0
        SFFT[-k] = 0
        SFFT[k] = 0
    Dec.append(Ssig)
    Filtered = np.real(np.fft.ifft(FFT))
    SFFTMap.append( np.abs(SFFT) )
    FFTMap.append( np.abs(tmp) )
    Reconstructed.append( Filtered[0:6000] )
    H.append ( np.abs(signal.hilbert(Filtered))[0:6000] )
    RawTmp = np.abs(signal.hilbert(np.real(np.fft.ifft(SFFT))))[0:6000/2]
    SH.append ( RawTmp )
    SHDec.append(decimate(RawTmp,15))

plt.plot(Reconstructed[20][2000:4500], label='Enveloppe of the signal') 
plt.plot(H[20][2000:4500], label='Enveloppe of the signal') 
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
plt.show()

plt.imshow(np.sqrt(FFTMap),aspect="auto")
#plt.savefig("D_image.jpg", bbox_inches='tight')
plt.show()

LineNb = 35
plt.plot(range(2500),ND[LineNb,2000:4500]-256, label='Raw signal')
plt.plot(range(2500)-OffSets[LineNb],H[LineNb][2000:4500], label='Enveloppe of signal') 
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
plt.show()

LineNbTmp = [26,27,28]
for k in LineNbTmp:
    plt.plot(H[k][2000:4500], label='Line '+str(k)) 
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
plt.show()

plt.plot(Reconstructed[30][0:6000],"r")
plt.plot(H[30][0:6000],"g")
plt.show()

plt.figure(figsize=(15, 10))
plt.imshow((H),aspect="auto")
plt.savefig("D_image.jpg", bbox_inches='tight')
plt.show()

plt.figure(figsize=(15, 10))
plt.imshow((SHDec),aspect="auto")
plt.savefig("D_image.jpg", bbox_inches='tight')
plt.show()