#!/usr/bin/env python

from time import sleep
from collections import deque
import numpy as np
from scipy import signal

from zoetrope import aniplot as plt
from bstrd import BS, bsstream
plt.blit = False

# config
#cam_name = "SARES11-XMI125-C4P1.jet_projection"
cam_name = "SARES11-SPEC125-M2.jet_projection"
#xmin = 0
#xmax = 400
#ymin = 400
#ymax = 700
#proj_axis = 0
#um_per_px = 2.28
um_per_px = 3.33 
buffer_length = 1000
shots = 10

def get_data(cam):
    jet = np.empty((shots, 550)) 
    for i in range(shots):
        jet[i] = cam.get()
        next(bsstream)
    proj = jet.max(axis=0)
    return proj 


#def get_data(cam):
#    jet = np.empty(550)
#    jet = cam.get()
#    next(bsstream)
#    proj = jet 
#    return proj


def find_spacing(py, extend=5):
    py -= py.mean()
    #tukey_filter = np.concatenate((signal.tukey(120)[0:30],np.ones(291),signal.tukey(120)[-30:-1]))
    #tukey_filter = signal.gaussian(350, 100)
    #py *= tukey_filter
    #py = signal.savgol_filter(py, 39, polyorder=2)

    if extend > 1:
        py = np.concatenate((py, np.zeros(len(py) * extend)))

    spectrum = np.abs(np.fft.fft(py))**2
    freq = np.fft.fftfreq(len(spectrum))

    which = (freq > 0.005)
    spectrum = spectrum[which]
    freq = freq[which]

    index_peaks = signal.find_peaks(spectrum)[0]
    peaks_pos = freq[index_peaks]
    peak_height = spectrum[index_peaks]

    cut_freq = peaks_pos[np.argmax(peak_height)]
    return freq, spectrum, cut_freq


def make_xs(arr):
    return np.arange(len(arr))


# create channel
cam = BS(cam_name)
proj = get_data(cam)
pixel_proj  = make_xs(proj)

plt.fig.set_figheight(9)
plt.fig.set_figwidth(9)
plt.suptitle(cam_name)

plt.subplot(221)
plt.title("projection")
pln_proj = plt.plot(proj, pixel_proj)
plt.xlim(-45,45)
plt.grid()

plt.subplot(222)
plt.title("FFT")
plt.xlabel("frequency in 1/pixel")
pln_fft_spec = plt.plot([0])
plt.xlim(0,0.25)
plt.grid()
#plt.ylim(0,500000)
pln_fft_cut  = plt.plot([0])

plt.subplot(212)
plt.title(f"spacing assuming {um_per_px} μm per pixel")
pln_spac = plt.plot([0], 'o')
pln_spac.ax.axhline(y=120, color='r') # 120um spacing set for Neutze cytco
plt.xlabel('1000 (*10) most recent images')
plt.ylabel('ladder spacing in $\mu$m')
plt.ylim(0,250)
plt.xlim(0,buffer_length)

plt.grid()
plt.tight_layout()

spacings = deque(maxlen=buffer_length)

for counter in plt.show():
    print(counter)
    next(bsstream)

    proj = get_data(cam)

    pln_proj.set(proj, pixel_proj)

    freq, spectrum, cut_freq = find_spacing(proj, extend=10)
    spac = um_per_px / cut_freq
    spacings.append(spac)

    pln_spac.set(make_xs(spacings), spacings)
    pln_fft_spec.set(freq, spectrum)
    pln_fft_cut.set([cut_freq]*2, [spectrum.min(), spectrum.max()])
    pln_fft_spec.ax.set_title(f"FFT\ncurrent spacing: {spac:.1f} μm")

    # this, I need to move into the library
    ps = [pln_proj, pln_spac, pln_fft_spec]
    for p in ps:
        p.ax.relim()
        p.ax.autoscale_view()

bsstream.close()