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

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')

class UI:
    def __init__(self) -> None:
        self._size_F = 256
        self._size_H = 256
        self._sigma  = 1.1
        self._seed   = 0
        self._shift  = 0
        self._fig, self._axs = self.init_ui()

        self.update()

    def init_ui(self) -> None:
        fig, (ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8) = plt.subplots(8, 1, gridspec_kw={'height_ratios':[4,4,4,1,1,1,1,1]})

        ax1.set_title('Terrain')
        self.plt_f, = ax1.plot([])
        self.plt_h, = ax1.plot([])

        ax2.set_title('MOSSE response signal')
        ax2.set_ylim([0, 1.2])
        self.line_r  = ax2.axvline()
        self.plt_r,  = ax2.plot([])
        self.plt_r2, = ax2.plot([])

        ax3.set_title('Gaussian')
        self.plt_g, = ax3.plot([])

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

        self._slider1 = Slider(ax4, 'sigma',  0.1, 5,    valinit=self._sigma)
        self._slider2 = Slider(ax5, 'shift', -180, 180,  valinit=self._shift,  valstep=1)
        self._slider3 = Slider(ax6, 'seed',   0,   50,   valinit=self._seed,   valstep=1)
        self._slider4 = Slider(ax7, 'Size f', 64,  1024, valinit=self._size_F, valstep=8)
        self._slider5 = Slider(ax8, 'Size h', 64,  1024, valinit=self._size_H, valstep=8)

        self._slider1.on_changed(self.update_sigma)
        self._slider2.on_changed(self.update_shift)
        self._slider3.on_changed(self.update_seed)
        self._slider4.on_changed(self.update_size_f)
        self._slider5.on_changed(self.update_size_h)

        fig.subplots_adjust(hspace=0.75)

        return fig, [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8]

    def update(self):
        np.random.seed(self._seed)
        g  = signal.windows.gaussian(self._size_F, std=self._sigma,sym=True)
        f  = generate_signal(self._size_F) * signal.windows.hamming(self._size_F)
        h  = np.zeros((self._size_F))
        s = self._size_F//2-self._size_H//2
        h[s:s+self._size_H]  = (f[s:s+self._size_H] + np.random.normal(0,0.5, self._size_H)) * signal.windows.hamming(self._size_H)
        h = np.roll(h,int(self._shift/180*self._size_F//2))
        F  = fft(f)
        G  = fft(g)
        H = fft(h)
        R = H*G/F
        r = ifft(R)
        s = np.argmax(abs(r))
        r2 = np.copy(r)
        r2[s-5:s+5] = 0

        x = np.linspace(-180,180,self._size_F)
        self.plt_g.set_data(x, g)
        self.plt_r.set_data(x,abs(r))
        self.plt_r2.set_data(x,abs(r2))
        self.plt_f.set_data(x,f)
        self.plt_h.set_data(x,h)
        self._axs[0].set_ylim([min(np.min(h),np.min(f))-1, max(np.max(h),np.max(f))+1])
        self._axs[1].set_ylim([0, np.max(r)+0.2])
        self._axs[2].set_ylim([0, np.max(g)*1.1])
        self.line_r.set_xdata([round((np.argmax(abs(r))/self._size_F-0.5)*360)])
        self._fig.canvas.draw_idle()

    def update_sigma(self, val):
        self._sigma = val
        self.update()
    
    def update_size_f(self, val):
        if val < self._size_H:
            self._slider4.eventson = False
            self._slider4.set_val(max(val, self._size_H))
            self._slider4.eventson = True
    
        self._size_F = max(val, self._size_H)
        self.update()

    def update_size_h(self, val):
        if val > self._size_F:
            self._slider5.eventson = False
            self._slider5.set_val(min(val, self._size_F))
            self._slider5.eventson = True
    
        self._size_H = min(val, self._size_F)
        self.update()

    def update_seed(self, val):
        self._seed = val
        self.update()

    def update_shift(self, val):
        self._shift = val
        self.update()

    def show(self) -> None:
        plt.show()

if __name__ == '__main__':
    ui = UI()
    ui.show()