import numpy as np
import scipy.signal as signal
from scipy.fftpack import fft, fftshift, ifft
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

N = 256

def gaussian_filter1d(size,sigma):
    filter_range = np.linspace(-int(size/2),int(size/2),size)
    gaussian_filter = [1 / (sigma * np.sqrt(2*np.pi)) * np.exp(-x**2/(2*sigma**2)) for x in filter_range]
    return gaussian_filter

def generate_signal(N:int) -> np.ndarray:
    x = np.arange(1, N)
    y = np.zeros((N))
    for i in range(len(x)):
        y[i] = np.random.normal(scale=1) + (y[i-1] if i > 1 else 0)
    return np.convolve(y,gaussian_filter1d(N,1),'same')

if __name__ == '__main__':

    np.random.seed(42)
    f_full = generate_signal(2048)
    x = np.arange(-180,180,360/N)
    f = f_full[1024:1024+N]
    h = f_full[1024:1024+N]
    g = signal.gaussian(N, std=10,sym=True)

    F  = fft(f)
    F_ = np.conjugate(F)
    G  = fft(g)

    K_ = (G*F_)/(F*F_)

    H = fft(h)
    r = ifft(H*K_)

    # ==========================================

    fig, (ax1, ax2, ax3, ax4, ax5, ax6, ax7) = plt.subplots(7, 1, gridspec_kw={'height_ratios':[4,4,4,1,1,1,1]})

    ax1.set_title('Terrain')
    plt_f, = ax1.plot(x,f)
    plt_h, = ax1.plot(x,h)

    ax2.set_title('MOSSE response signal')
    ax2.set_ylim([0, 1.2])
    line_r = ax2.axvline(x=-N//2+np.argmax(abs(r)), color='r')
    plt_r, = ax2.plot(x,abs(r))
    plt_r2, = ax2.plot(x,abs(r))

    ax3.set_title('Gaussian')
    plt_g, = ax3.plot(x,g)

    ax1.set_xlim([-180,180])
    ax2.set_xlim([-180,180])
    ax3.set_xlim([-180,180])

    slider1 = Slider(ax4, 'sigma', 0.3, 10, valinit=0.1)
    slider2 = Slider(ax5, 'shift', -N//2  , N//2, valinit=0, valstep=1)
    slider3 = Slider(ax6, 'seed', 0  , 50, valinit=0, valstep=1)
    slider4 = Slider(ax7, 'N', 128  , 1024, valinit=256, valstep=8)

    sigma = 0.3
    shift = 0

    def update():
        # K_ = (G*F_)/(F*F_)
        window = np.ones((N)) #signal.windows.hamming(N)
        H = fft(h*window)
        F = fft(f*window)
        R = H*G/F
        r = ifft(R)
        s = np.argmax(abs(r))
        r2 = np.copy(r)
        r2[s-5:s+5] = 0

        plt_g.set_data(x,g)
        plt_r.set_data(x,abs(r))
        plt_r2.set_data(x,abs(r2))
        plt_h.set_data(x,h)
        plt_f.set_data(x,f)
        ax1.set_ylim([min(np.min(h),np.min(f))-1, max(np.max(h),np.max(f))+1])
        ax2.set_ylim([0, np.max(r)+0.2])
        line_r.set_xdata(round((np.argmax(abs(r))/N-0.5)*360))
        fig.canvas.draw_idle()

    def update_sigma(val):
        global g, G, sigma, N
        sigma = val
        g = signal.gaussian(N, std=sigma,sym=True)
        G  = fft(g)
        update()

    def update_shift(val):
        global shift, H, h
        shift = -val
        h = f_full[1024+round(shift*(N/360)):1024+N+round(shift*(N/360))]
        noise = np.random.normal(0,0.5, N)
        h = h+noise
        update()

    def update_seed(val):
        global f_full, f, h, F, H, F_, shift
        np.random.seed(val)
        f_full = generate_signal(2048)
        f = f_full[1024:1024+N]
        h = f_full[1024+round(shift*(N/360)):1024+N+round(shift*(N/360))]
        noise = np.random.normal(0,0.5, N)
        h = h+noise
        update()

    def update_n(val):
        global g, G, N, f_full, f, h, F, H, F_, shift, x, sigma, slider2
        N = val
        x = np.arange(-180,180,360/N)
        g = signal.gaussian(N, std=sigma,sym=True)
        G  = fft(g)
        f = f_full[1024:1024+N]
        h = f_full[1024+shift:1024+N+shift]
        noise = np.random.normal(0,0.5, N)
        h = h+noise

        update()


    slider1.on_changed(update_sigma)
    slider2.on_changed(update_shift)
    slider3.on_changed(update_seed)
    slider4.on_changed(update_n)

    plt.subplots_adjust(hspace=0.5)
    plt.show()