import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import welch, hanning
#from autocorr import autocorr
plt.rcParams.update({'font.size': 22})

def autocorr(x, twosided=False, tapered=True):
    import numpy as np
# https://currents.soest.hawaii.edu/ocn_data_analysis/_static/SEM_EDOF.html
    """
    Return (lags, ac), where ac is the estimated autocorrelation 
    function for x, at the full set of possible lags.

    If twosided is True, all lags will be included;
    otherwise (default), only non-negative lags will be included.

    If tapered is True (default), the low-MSE estimate, linearly
    tapered to zero for large lags, is returned.
    """
    nx = len(x)
    xdm = x - x.mean()
    ac = np.correlate(xdm, xdm, mode='full')
    ac /= ac[nx - 1]
    lags = np.arange(-nx + 1, nx)
    if not tapered:  # undo the built-in taper
        taper = 1 - np.abs(lags) / float(nx)
        ac /= taper
    if twosided:
        return lags, ac
    else:
        return lags[nx-1:], ac[nx-1:]


# ***** read u

utime = np.genfromtxt("utime_backstep.dat", dtype=None)

u1=utime[:,0] #x=3.3, y=0.037
u2=utime[:,1] #x=8.7, y=0.037
u3=utime[:,2] #x=18.7, y=0.037
u4=utime[:,3] #x=3.3, y=0.85
u5=utime[:,4] #x=8.7, y=0.85
u6=utime[:,5] #x=18.7, y=0.85
u7=utime[:,6] #x=-3.94, y=1.017
v1=utime[:,7] #x=3.3, y=0.85
v2=utime[:,8] #x=8.7, y=0.85
v3=utime[:,9] #x=18.7, y=0.85
t=utime[:,10]  # time


n=len(u1)

imax=700

u1_fluct=u1-np.mean(u1)
u2_fluct=u2-np.mean(u2)
v1_fluct=v1-np.mean(v1)

u1_fluct_var=np.var(u1_fluct)

two_uu_1=np.correlate(u1_fluct,u1_fluct, mode='full')/u1_fluct_var

# below I compute the two-p corr.
# but since itäs very slow, I've commented the lines
#max=700
#nstep=len(u1_fluct)
#my_two=np.zeros(imax)
#ii=0
#for tau in range (0,imax): # time seperation
   #for it in range (0,nstep-imax-2):
      #ii=ii+1
      #i0=it+tau
      #my_two[tau]=my_two[tau]+u1_fluct[it]*u1_fluct[i0]
#
#my_two=my_two/my_two[0] #normalize _uu_t so that two_uu_t(1)=1
#

two=np.correlate(u1_fluct,u1_fluct,'full')

# find max
nmax=np.argmax(two)
# and its value
two_max=np.max(two)

# two_max is symmwetric. Pick the right half and normalize
two_sym_norm=two[nmax:]/two_max


lags, auto_x = autocorr(u1_fluct)

#%%%%%%%%%%%%%%% plotting section %%%%%%%%%%%%%%%%%%%%%%%%%%
# plot autocorr
fig1 = plt.figure("Figure 1")
plt.plot(t[0:imax],two_sym_norm[0:imax],'b-')
#plt.plot(t[0:imax:10],my_two[0:imax:10],'ro')
plt.plot(t[0:imax],auto_x[0:imax],'g-')
plt.xlabel('separation')
plt.axis([0, 2 ,0 ,1])
plt.title('two-point correlation')
plt.ylabel('2-p')

plt.savefig('two_time_corr_backstep_1-4_python.ps.eps')