import h5py as hdf
import math
import numpy as np
from scipy import signal
from scipy.signal import hilbert
import sys


import matplotlib
import matplotlib.pyplot as plt
file = open('/hipa/bd/data/measurements/tina/20240710-223007_2000.txt','r')
#file = open('/hipa/bd/data/measurements/tina/20240710-223158_1175.txt','r')
#file = open('/hipa/bd/data/measurements/tina/20240807-182215.txt','r')
#file = open('/hipa/bd/data/measurements/tina/test.txt','r')
content = file.readlines()
#print(content)
file.close()


t = []
y1 = []
y2 = []
yall1 = []
yall2 = []
idxall = [] 
#t_stepsize =  0.0000000004*50
t_stepsize =  1/(50.6328*10**6)/50 # 0.00000002
print(t_stepsize)
print(1/(2.5*10**9))
#t_stepsize = 1/(2.5*10**9)
t_offset = 2500002/2 
t_inc = 0

'''
for count, entry in enumerate(content[5:]):
    t_inc += t_stepsize
    val=entry.split('\t')
    #print(val)
    #print(val[0])
    #print(val[1])
    y1.append(float(val[0]))
    y2.append(float(val[1]))
    idx.append(t_inc) #count
    if count < 10:
        print(val[0], val[1], count, t_inc )
'''
for count, entry in enumerate(content[5:]):
    entry=entry.replace('\n','')
    val=entry.split('\t')
    t_inc += t_stepsize
    t.append(val[0])
    #if (abs(float(val[1])) > 0.0022):
    #    val[1] = val[1]
    #else:
    #    val[1]=0
    #if (abs(float(val[2])) > 0.015):
    #    val[2] = val[2]
    #else:
    #    val[2]=0

    yall1.append(float(val[1])*(-1))
    yall2.append(float(val[2]))
    idxall.append(t_inc) #count
    #if count < 200:
    #    break
    #    print(val[1], val[2], count, t_inc )
    #if count > 300:
    #    sys.exit(0)

start=12
stop=2500002
_y1 = []
_y2 = []
count = 0
t_inc = 0
idx = []
for val in range(start,stop,50):
     _test1_array = yall1[val : val+50]
     _test2_array = yall2[val : val+50]
     _y1.append(np.array(_test1_array).max())
     _y2.append(np.array(_test2_array).max()) 
     count += 1
     t_inc += t_stepsize
     idx.append(t_inc)

idx = np.array(idx) - t_offset*t_stepsize

     
y1 = _y1
y2 = _y2
print(y1[0:50], y2[0:50] )

'''
    y1 = [ 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3,
            1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3,
            1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3
    ]
    y2 = [ 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1,
            1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3,
            2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1, 2, 2, 1, 4, 5, 6, 2, 1, 1, 1, 3, 1
    ]
    y2=  [ i*5 for i in y2]
'''
f1 = hdf.File('/hipa/bd/data/measurements/Tina_2024-09-13_15:57:07.h5','a')
#f1 = hdf.File('/hipa/bd/data/measurements/Tina_2024-07-03_15:44:14.h5','a') #5Tina_2024-08-07_10:41:04.h5','a')
grp = f1.require_group('Raw_data')
if 't' in grp:
    del grp['t']
if 'y1' in grp:
    del grp['y1']   
if 't_idx' in grp:
    del grp['t_idx']
if 'y2' in grp:
    del grp['y2']
if 'yall1' in grp:
    del grp['yall1']
if 'yall2' in grp:
    del grp['yall2']    
dset_yall1 = grp.create_dataset('yall1', data=yall1)
dset_yall2 = grp.create_dataset('yall2', data=yall2)
dset_t = grp.create_dataset("t", data=t)
dset_y1 = grp.create_dataset('y1', data=y1)
dset_y2 = grp.create_dataset('y2', data=y2)
dset_t_idx = grp.create_dataset("t_idx", data=idx)
print(dset_y1.name)
f1.close()

    
analytic_signal_1 = y1 #hilbert(y1)
amplitude_envelope_1 = np.abs(analytic_signal_1) #analytic_signal_1 #
analytic_signal_2 = y2  #hilbert(y2)
amplitude_envelope_2 = np.abs(analytic_signal_2) #analytic_signal_2 #

#amplitude_envelope_1, peaks1 = signal.find_peaks(y1, height=0.012)
#amplitude_envelope_2, peaks2 = signal.find_peaks(y2, height=0.015)


mean_amplitude_envelope_1 = np.mean(amplitude_envelope_1, keepdims=True)
mean_amplitude_envelope_2 = np.mean(amplitude_envelope_2, keepdims=True)

mean_amplitude_envelope_1 = np.mean(analytic_signal_1, keepdims=True)
mean_amplitude_envelope_2 = np.mean(analytic_signal_2, keepdims=True)

print(mean_amplitude_envelope_1, mean_amplitude_envelope_2)

normalized_amplitude_envelope_1 = (analytic_signal_1 - mean_amplitude_envelope_1) #/ (np.std(amplitude_envelope_1)) # * len(amplitude_envelope_1))
normalized_amplitude_envelope_2 = (analytic_signal_2 - mean_amplitude_envelope_2) #/ (np.std(amplitude_envelope_2))

#normalized_amplitude_envelope_1 = (amplitude_envelope_1 - mean_amplitude_envelope_1) #/ (np.std(amplitude_envelope_1)) # * len(amplitude_envelope_1))
#normalized_amplitude_envelope_2 = (amplitude_envelope_2 - mean_amplitude_envelope_2) #/ (np.std(amplitude_envelope_2))
#rms_amplitude_envelope_1 = np.sqrt(np.mean(normalized_amplitude_envelope_1**2))
#rms_amplitude_envelope_2 = np.sqrt(np.mean(normalized_amplitude_envelope_2**2))
#print(rms_amplitude_envelope_1, rms_amplitude_envelope_2)


#norm_1 = np.linalg.norm(amplitude_envelope_1)
#normalized_amplitude_envelope_1 = amplitude_envelope_1/norm_1
#norm_2 = np.linalg.norm(amplitude_envelope_2)
#normalized_amplitude_envelope_2 = amplitude_envelope_2/norm_2

print (normalized_amplitude_envelope_1[0:10])
print (amplitude_envelope_1[0:10])

'''
y2_peaks = signal.find_peaks(normalized_amplitude_envelope_2, height=0.005)
y1_peaks = signal.find_peaks(normalized_amplitude_envelope_1, height=0.002)



print(y2_peaks)
print(y1_peaks)
print(y1_peaks[0])
print(y1_peaks[1]['peak_heights'])
print("length y1 peaks", len(y1_peaks[0]))
print(y2_peaks[0])
print(y2_peaks[1]['peak_heights'])
print("length y2 peaks", len(y2_peaks[0]))
#sys.exit(1)
normalized_amplitude_envelope_1_p = y1_peaks[1]['peak_heights']
normalized_amplitude_envelope_2_p = y2_peaks[1]['peak_heights']
#normalized_amplitude_envelope_1_p = y1
#normalized_amplitude_envelope_2_p = y2
#sys.exit(0)
'''

y1_array = []
y2_array = []

'''
for i, val in enumerate(y1):
    if (i in y1_peaks[0]):
        _idx = np.where(y1_peaks[0] == i)[0][0]
        y1_array.append(y1_peaks[1]['peak_heights'][_idx])
    else:
        y1_array.append(0)

for i, val in enumerate(y2):
    if i in y2_peaks[0]:
        _idx = np.where(y2_peaks[0] == i)[0][0]
        y2_array.append(y2_peaks[1]['peak_heights'][_idx])
    else:
        y2_array.append(0)                       
'''
#print(len(y1_array))
#print(len(y2_array))
        
#sys.exit(0)
#normalized_amplitude_envelope_1_p = y1_array
#normalized_amplitude_envelope_2_p = y2_array
                        
idx2 = []
idx_neg = []
#for val in reversed(idx[1:]): 
#      idx_neg.append(val * (-1))    
#idx2 = idx + idx_neg
mid_time = idx[-1]/2

for val in idx:
    idx_neg.append(val-mid_time)

#print(normalized_amplitude_envelope_1_p)
#print(normalized_amplitude_envelope_2_p)
#normalized_amplitude_envelope_1_p = y1
#normalized_amplitude_envelope_2_p = y2
    
corr = signal.correlate(normalized_amplitude_envelope_2, normalized_amplitude_envelope_1, mode='same', method='auto')/len(amplitude_envelope_1)
corr_full = signal.correlate(normalized_amplitude_envelope_2, normalized_amplitude_envelope_1, mode='full', method='auto') #/len(amplitude_envelope_1)
#corr_valid = signal.correlate(normalized_amplitude_envelope_1, normalized_amplitude_envelope_2, mode='valid', method='auto')
lags = signal.correlation_lags(len(normalized_amplitude_envelope_2), len(normalized_amplitude_envelope_1), mode='same')
lags_full = signal.correlation_lags(len(normalized_amplitude_envelope_2), len(normalized_amplitude_envelope_1), mode='full')

lags_full = lags_full #* t_stepsize
lags =  lags #* t_stepsize

#lags_valid =   signal.correlation_lags(len(normalized_amplitude_envelope_2), len(normalized_amplitude_envelope_1), mode='valid')                            
lag_full = lags_full[np.argmax(corr_full)]
lag = lags[np.argmax(corr)]
#lag_valid =  lags_valid[np.argmax(corr_valid)]                             
#corr = corr/rms_amplitude_envelope_1/rms_amplitude_envelope_2
print("lag same/full", lag, lag_full)
print("lag same/full", lag * t_stepsize, lag_full * t_stepsize)
print("length lags", len(lags))

print("VALID", signal.correlate(normalized_amplitude_envelope_1, normalized_amplitude_envelope_2, mode='valid'))

#lags = signal.correlation_lags(len(amplitude_envelope_2), len(amplitude_envelope_1))
#corr4 = signal.correlate(analytic_signal_1, analytic_signal_2, mode='full')
print(len(y2))
print(len(amplitude_envelope_2))
print("max correlation", np.max(corr))
print("arg of max corrrelation", np.argmax(corr))
print("time index", idx[np.argmax(corr)])
#print(y1)

#y1 = [0.000231, -0.000273, 0.000206, 0.000408,
#0.000609, 0.000231, 0.000508, 0.000760,
#0.000433, 0.000206, 0.000231, 0.000130]

subplots = 4
ln=500
ln2=1239770
ln3=1260230


if subplots == 4:
    fig, (ax) = plt.subplots(nrows=2, ncols=2, figsize=(18,9))
    
    #ax1.yticks(np.arange(min(y1), max(y1), 1))
    ax[0,0].ticklabel_format(useOffset=False, style='plain')
    #ax[0,0].plot(idx[0:ln], amplitude_envelope_1[0:ln], 'b')
    #y1
    ax[0,0].plot(idxall[0:ln], yall1[0:ln], 'ro')
    ###ax[0,0].plot(idx[0:ln2], normalized_amplitude_envelope_1[0:ln2], 'ro')
    #ax1.set_ylim([-0.05, 0.05])
    #y2
    ###ax[1,0].plot(idx[0:ln2], normalized_amplitude_envelope_2[0:ln2], 'ro')
    ax[1,0].plot(idxall[0:ln], yall2[0:ln], 'ro')
    #ax[1,0].plot(idx[0:ln], amplitude_envelope_2[0:ln], 'g')
#else:
#    fig, (ax3, ax4) = plt.subplots(2, figsize=(20,12))



ax[0,1].plot(lags[:], corr[:])
#ax3.plot(corr)
ax[1,1].plot(lags_full[:], corr_full[:], 'yo')
#start, end = ax1.get_ylim()
#start2, end2 = ax2.get_ylim()
#print(start, end)
#print(start2, end2)
#ax1.yaxis.set_ticks(np.arange(start, end, 10))
#ax2.yaxis.set_ticks(np.arange(start2, end2, 10))


plt.show()