from collections import deque
import numpy as np
from scipy.optimize import curve_fit
from scipy.special import erf
from scipy.stats import binned_statistic
from scipy.signal import hilbert, chirp, find_peaks, savgol_filter
from scipy import ndimage
from collections import OrderedDict
from cam_server.pipeline.data_processing import functions
from logging import getLogger
import numpy as np


initialized = False
p0 = None

def initialize(parameters, sig):
    global initialized, buffer
    buffer = deque(maxlen=parameters["dark_buffer_length"])
    buffer.append(np.ones(np.shape(sig)))
    buffer.append(np.ones(np.shape(sig)))
    initialized = True

def do_fft(t, tr, lm=False, lx=False, padding=False, bins=None):
    if bins is not None:
        xbin = np.linspace(0,t[-1]+(t[1]-t[0])/2,bins+1)
        tr, _, _ = binned_statistic(t, tr, statistic='mean', bins=xbin, range=None)
        t = np.linspace(0,t[-1],bins)
    if (lm and lx):
        lm = np.argmin(abs(t-lm));
        lx = np.argmin(abs(t-lx))
        tr=tr[lm:lx]
        t=t[lm:lx]
        
    if padding:
        N = t.size
        te = t[1]-t[0]
        t = np.arange(N+pad)*te+t[0]
        trp = np.zeros(N+pad)
        trp[:N]=tr
        tr=trp

    N = t.size
    T = t[N-1]-t[0]
    te = t[1]-t[0]
    fe = 1.0/te
    tfd=np.fft.fft(tr)/N
    freq=np.arange(N)*1.0/T
    tfd = tfd[:int(len(tfd)/2)]
    freq = freq[:int(len(freq)/2)]
    return freq, np.absolute(tfd), np.real(tfd), np.imag(tfd)

def gauss(x, x0, fwhm, a):
    return a*np.exp(-(x-x0)**2/fwhm**2*4*np.log(2))

def gauss_cos( x, x0, fwhm, a, f, phase):
    return a*np.exp(-(x-x0)**2/fwhm**2*4*np.log(2))*np.cos(2*np.pi*f*(x-x0)-phase)

def gauss_sin(x, x0, fwhm, a, f, phase):
    return a*np.exp(-(x-x0)**2/fwhm**2*4*np.log(2))*np.sin(2*np.pi*f*(x-x0)+phase)

def fit_data(x, y, bins=None, diag=False, sigma=None, distance=None):
    bounds=[[x[abs(y).argmax()]-0.3, 0.03, 0.005, 15., -np.inf], [x[abs(y).argmax()]+0.3, 2., 10., 35., np.inf]]
    global p0
    if p0 is None or np.any([~(lb<p<hb) for p, lb, hb in zip(p0, bounds[0], bounds[1])]):
        p0 = find_signal(x, y, distance=distance)
        #p0 = [x[abs(y).argmax()], abs(x[-1]-x[0])/4, abs(y).max(), 24, np.pi]
    
    # assumes that x is in ps
    if bins is not None:
        xbin = np.linspace(0,x[-1]+(x[1]-x[0])/2,bins+1)
        y, _, _ = binned_statistic(x, y, statistic='mean', bins=xbin, range=None)
        x = np.linspace(0,x[-1],bins)
    popt, pcov = curve_fit(gauss_cos, x, y, p0, sigma=sigma, bounds=bounds)
    #x0, fwhm, a, f, phase = p0
    p0 = popt
    x0, fwhm, a, f, phase = popt

    if diag:
        return x0, fwhm, a, f, np.remainder(phase, 2*np.pi), p0, popt
    else:
        return x0, fwhm, a, f, np.remainder(phase, 2*np.pi)

def fit_extrema(time, y, distance, diag=None):
    xp, _ = find_peaks(abs(y), height=abs(y).max()/5, distance=distance)
    yp = y[xp]
    return fit_data(time[xp], yp, bins=None, diag=diag)

def get_max(c, dpx_poly=None, offset=0):
    """getting maximum from a correlation curve (optionally using polynomial fit)"""
    im = c.argmax()
    mx = c[im]

    if dpx_poly:
        try:
            poly = np.polyfit(np.arange(-dpx_poly//2,dpx_poly//2),c[im-dpx_poly//2:im+dpx_poly//2], 2)
            root = np.roots(np.polyder(poly,1))
            im+= np.real(root[np.argmin(np.abs(root))])
        except:
            im=im
    return im + offset, mx

def find_signal(time, y, distance):
    dt = (time[1]-time[0])
    xp, _ = find_peaks(abs(y), height=abs(y).max()/5, distance=distance)
    yp = abs(y)[xp]
    p0 = [xp[yp.argmax()], (xp[-1]-xp[0])/3, yp.max()]
    popt, pcov = curve_fit(gauss, xp, yp, p0)
    x0, fwhm, a = popt
    xp_sign = xp * np.sign(y[xp])
    # idx_closest = np.argmin(np.abs(xp_sign-x0))
    # f = 1/(dt*np.mean(np.diff(xp_sign[idx_closest-1:idx_closest+1]))) Metric that doesnt make sense but is sensitive to the phase
    xp_pos = xp_sign[xp_sign>0]
    idx_closest = np.argmin(np.abs(xp_pos-x0))
    weights=abs(y)[xp[:-1]]+abs(y)[xp[1:]]
    f = 1/(dt*np.average(np.diff(xp), weights=weights)*2)
    #f = 1/(dt*np.mean(np.diff(xp_pos[idx_closest-2:idx_closest+1]))) works but uses only two points
    #f = 1/(dt*np.mean(np.diff(xp_pos))) <-- rather noisy
    phase = f*(np.min(xp_pos[xp_pos>x0]-x0)*dt)*2*np.pi
    x0 = x0*dt
    fwhm = fwhm*dt
    return x0, fwhm, a, f, phase
    
def normstep(step):
    """normalizing a test signal for np.correlate"""
    step = step - np.mean(step)
    step = step / np.sum(step**2)
    return step

def get_reference_function(sigma_px=30, reflen=150):
    rng = reflen / np.sqrt(2) / sigma_px / 2
    ref = -erf(np.linspace(-rng, rng, reflen)) / 2
    ref = normstep(ref)
    return ref

def process(bsdata, pid, timestanp, parameters):
    global initialized
    global p0
    ## initialize all values sent to BS stream to None
    x=None
    sig=None
    bg=None
    time=None
    ratio=None
    freq=None
    fft_frequency = None
    fft_amplitude = None
    fft_amplitude_real = None
    fft_amplitude_imag = None
    x0=None 
    fwhm=None
    a=None 
    f=None
    phase=None
    c=None


    ## params used for evaluation    
    calibration = parameters["calibration"]
    padding = parameters["padding"]
    lm = parameters["time_min"]
    lx = parameters["time_max"]
    
    
    ## get data and events for sorting
    events = bsdata[parameters["events"]]
    sig = bsdata[parameters["sig"]]
    x = np.arange(0,len(sig))

    if not initialized:
        initialize(parameters, sig)

    is_midir_dark = parameters["is_midir_dark"]
    is_midir_dark_evt = False
    
    # add to dark reference
    if is_midir_dark:
        is_midir_dark_evt = events[parameters["event_code_midir_dark"]]
        if is_midir_dark_evt:
            buffer.append(sig)
    # or analyze
    bg = np.mean(buffer, axis=0)
    if not is_midir_dark_evt:
        ratio = sig/bg
        time = np.polyval(calibration, x)
        #fft_frequency, fft_amplitude, fft_amplitude_real, fft_amplitude_imag = do_fft(time, ratio-1, lm=False, lx=False, padding=False, bins=parameters["fft_bins"])
        ref=get_reference_function(parameters["sigma_px"], parameters["reflen"])
        c = np.correlate(ratio-1, ref, "same")
        if parameters["fit_data"]:
            x0, fwhm, a, f, phase = fit_data(time, c, bins=parameters["bins"], distance=parameters["distance"])
        else:
            x0, fwhm, a, f, phase = find_signal(time, c, distance=parameters["distance"])
        
    data = {
        "SARES20-CEP02:spectrometer_x": x,
        "SARES20-CEP02:spectrometer_signal": sig,
        "SARES20-CEP02:spectrometer_background": bg,
        "SARES20-CEP02:spectrometer_time": time,
        "SARES20-CEP02:spectrometer_ratio": ratio,
        "SARES20-CEP02:spectrometer_correlation": c,        
        "SARES20-CEP02:fft_frequency": fft_frequency,
        "SARES20-CEP02:fft_amplitude_abs": fft_amplitude,
        "SARES20-CEP02:fft_amplitude_real": fft_amplitude_real,
        "SARES20-CEP02:fft_amplitude_imag": fft_amplitude_imag,
        "SARES20-CEP02:fit_x0": x0,
        "SARES20-CEP02:fit_fwhm": fwhm,
        "SARES20-CEP02:fit_amplitude": a,
        "SARES20-CEP01:fit_frequency": f,
        "SARES20-CEP01:fit_phase": phase,
        }

        #data.update({"p_est": pinit, "popt": popt})
    return data