#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jul 22 19:55:56 2017

@author: Sean
"""
import time
import numpy as np
import matplotlib.pyplot as plt
import para_fit as pf

def cut(x,n):
    x[x<n] = 0
    return x

def d_thrsd(a,n = 0.15):
    a = np.array(a)
    x = a.shape
    if len(x) == 1:
        thrsd = a.max() - n
        a = cut(a,thrsd)
    else:
        for i in range(x[0]):
            thrsd = a[i,:].max() - n
            a[i,:] = cut(a[i,:],thrsd)
    return a, thrsd

def avg(a,num):
    a = np.array(a)
    m = 0
    [prd, smp_pt] = a.shape
    x_pt = int(prd//num)
    a = a.reshape(num,x_pt,smp_pt)
    for i in range(num):
        m = m + d_thrsd(a[i,:,:],0.1)[0]
    m = m/num
    return m

def findpeak(a):
    m = a.argmax()
    return m


t0 = time.time()

smp_rate = 200e6            # sampling rate of DAQ
frm_rate = 10               # Galvo frequency
vco_frq = 320e3               # VCO frequency
line_rate = 200             # 
fft_num = 2048
fft = np.fft.fft

prd = int(smp_rate//vco_frq)
x_pt = int(vco_frq//line_rate)
y_pt = line_rate//frm_rate

data = np.fromfile('a_filter 0', dtype = '>d')
grp = len(data)//prd
data = data.reshape([grp,prd])

if prd < fft_num:
    pad_zeros = np.zeros([grp,fft_num - prd])
    data = np.concatenate((data,pad_zeros),axis =1)

a = np.abs(fft(data))
m = np.zeros(x_pt)
st_pt = 265
end_pt = 950
#
xs = a[0:1600,st_pt:end_pt].copy()
#xs = avg(xs,y_pt)
xs = list(xs)

xss = a[0:1600,st_pt:end_pt].copy()
xss = list(xss)

#m = m.reshape([y_pt,x_pt])
#plt.plot(xs[0],'o')

xdata = np.array(range(len(xss[100])))
p = pf.para_fit()
l = len(xss)
m = np.zeros([l,2]).astype('uint16')
i = 0
for x in xs:
    x,thrsd = d_thrsd(x,0.2)
    p.thrsd = thrsd
    xss[i] = x.copy()
    m[i,0] = int(x.argmax())
    x[m[i,0] -10: m[i,0] +10] = 0
    if len(x[x>0]) == 0:
        x = xs[i].copy()
        x,thrsd = d_thrsd(x,0.3)
        xss[i] = x.copy()
        x[m[i,0] -10: m[i,0] +10] = 0
        if len(x[x>0]) == 0:
            x = xss[i].copy()
            x[m[i,0]] = 0
            m[i,1] = x.argmax()
        else:
            m[i,1] = x.argmax()
    else:
        m[i,1] = x.argmax()
    i = i+1

plt.plot(((m[:,0] + m[:,1])/2 + st_pt) * 0.7, '.')
plt.plot(m[:,0],'.')
plt.plot(m[:,1],'.')

t = time.time() - t0
print(t)