# =============================================================================
# A couple of functions to compute a cross-correlation between the signals from
# pins of double rake probe and then plot the result
# =============================================================================
import numpy as np
from matplotlib import pyplot as plt
from urllib.request import urlopen
plt.ion()
# =============================================================================
# A function for data collection from the pins of the double rake probe
# =============================================================================
def collect(shot, pin):
    print('Data from shot %i, pin %s' % (shot, pin))
    url = f'http://golem.fjfi.cvut.cz/utils/data/{shot}/{pin}'
    ourl = urlopen(url)
    data_all = np.loadtxt(ourl)
    data = np.array([row[1] for row in data_all])
    ourl.close()
    return data


# =============================================================================
# The function will compute cross-correlation between the signals from two pins
# of the double rake probe. 'left' and 'right' are points on x-axis (time)
# that should be taken from 'U_loop' diagnostics and are points indicating a
# plato. Data from that interval gives better results.
# =============================================================================
def cross_corr(shot, pin1, pin2, left, right):
    dt = 1e-6  # time step computed from the frequency
    data1 = collect(shot, pin1)
    data2 = collect(shot, pin2)
    time_axis = np.arange(data1.size)*dt
    mask_ft = (left < time_axis) & (time_axis < right)
    data1 = data1[mask_ft]
    data2 = data2[mask_ft]
    norm = np.sqrt(np.sum(np.correlate(data1, data1)) * np.sum(np.correlate(data2, data2)))
    corr = np.correlate(data1, data2, 'same') / norm 
    N = len(corr)
    time_lag = (np.arange(N) - N//2) * dt  # time lag between two pins
    return time_lag, corr


# =============================================================================
# A function for fitting the cross-correlation for more precise time lag.
# Fitting is done by a polynomial function. 
# =============================================================================
def fit(time_lag, corr):
    center = np.argmax(corr)
    time = time_lag[center - 5:center + 5]
    corr = corr[center - 5:center + 5]
    coefs = np.polyfit(time, corr, 9)
    ffit = np.poly1d(coefs)
    time_new = np.linspace(time[0], time[-1], 100)
    ffit_new = ffit(time_new)
    return time_new, ffit_new


# =============================================================================
# A function for plotting a close-up of the cross-correlation and the fit 
# around the peak
# =============================================================================
def plot_fit():
#    plt.figure(); plt.clf()
    plt.title('Cross-correlation and fit of pins %s and %s, shot %i' % (pin1, pin2, shot))
    plt.xlabel('Time lag, [us]')
    plt.ylabel('Cross-correlation coeficient, [-]')
    c = np.argmax(corr)
    plt.plot(lag[c-5:c+5]*1e6, corr[c-5:c+5], 'o', label = 'corr')
    plt.plot(lag_new*1e6, ffit, label = 'fit')
#    plt.axvline(lag[np.argmax(corr)]*1e6, color = 'k')
    plt.axvline(lag_new[np.argmax(ffit)]*1e6, color='g', label = 'fit lag')
    plt.legend(loc='upper left')
    plt.show()
    
    
# =============================================================================
# Example: plot and results
# =============================================================================
shot, pin1, pin2 = 30375, 'drppin15', 'drppin11'
#shot, pin1, pin2 = 21757, 'a1', 'b1'
lag, corr = cross_corr(shot, pin1, pin2, 10e-3, 12e-3)
lag_new, ffit = fit(lag, corr)

fig = plt.figure()
plot_fit()
fig.set_size_inches(9.5, 6., forward=True)
fig.savefig('Fit, pins %s and %s, shot %i' % (pin1, pin2, shot))

print("Time shift between the pins:", round(lag_new[np.argmax(ffit)]*1e6, 3), "[us]")
print("Poloidal velocity:", round(2.5e-3 / lag_new[np.argmax(ffit)], 3), "[m/s]")