#!/usr/bin/env python3

__doc__ = \
"""
Show the curve for signal-to-noise ratio vs N_samples
"""

from collections import namedtuple

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import scipy.fftpack as ft

rng = np.random.default_rng()

passband = namedtuple("Band", ['low', 'high'], defaults=[0, np.inf])

def get_freq_spec(val,dt):
    """From earsim/tools.py"""
    fval = np.fft.fft(val)[:len(val)//2]
    freq =  np.fft.fftfreq(len(val),dt)[:len(val)//2]
    return fval, freq


def ft_spectrum( signal, sample_rate=1, ftfunc=None, freqfunc=None, mask_bias=False, normalise_amplitude=False):
    """Return a FT of $signal$, with corresponding frequencies"""

    if True:
        return get_freq_spec(signal, 1/sample_rate)

    n_samples = len(signal)

    if ftfunc is None:
        real_signal = np.isrealobj(signal)
        if False and real_signal:
            ftfunc = ft.rfft
            freqfunc = ft.rfftfreq
        else:
            ftfunc = ft.fft
            freqfunc = ft.fftfreq

    if freqfunc is None:
        freqfunc = ft.fftfreq

    normalisation = 2/len(signal) if normalise_amplitude else 1

    spectrum = normalisation * ftfunc(signal)
    freqs = freqfunc(n_samples, 1/sample_rate)

    if not mask_bias:
        return spectrum, freqs
    else:
        return spectrum[1:], freqs[1:]


def plot_spectrum( spectrum, freqs, plot_complex=False, plot_power=False, plot_amplitude=None, ax=None, freq_unit="Hz", freq_scaler=1):
    """ Plot a signal's spectrum on an Axis object"""
    plot_amplitude = plot_amplitude or (not plot_power and not plot_complex)
    alpha = 1

    if ax is None:
        ax = plt.gca()

    ax.set_title("Spectrum")
    ax.set_xlabel("f" + (" ["+freq_unit+"]" if freq_unit else "" ))
    ylabel = ""
    if plot_amplitude or plot_complex:
        ylabel = "Amplitude"
    if plot_power:
        if ylabel:
            ylabel += "|"
        ylabel += "Power"
    ax.set_ylabel(ylabel)

    if plot_complex:
        alpha = 0.5
        ax.plot(freqs/freq_scaler, np.real(spectrum), '.-', label='Real', alpha=alpha)
        ax.plot(freqs/freq_scaler, np.imag(spectrum), '.-', label='Imag', alpha=alpha)

    if plot_power:
        ax.plot(freqs/freq_scaler, np.abs(spectrum)**2, '.-', label='Power', alpha=alpha)

    if plot_amplitude:
        ax.plot(freqs/freq_scaler, np.abs(spectrum), '.-', label='Abs', alpha=alpha)

    ax.legend()

    return ax

def plot_phase( spectrum, freqs, ylim_epsilon=0.5, ax=None, freq_unit="Hz", freq_scaler=1):
    if ax is None:
        ax = plt.gca()

    ax.set_ylabel("Phase")
    ax.set_xlabel("f" + (" ["+freq_unit+"]" if freq_unit else "" ))

    ax.plot(freqs/freq_scaler, np.angle(spectrum), '.-')
    ax.set_ylim(-1*np.pi - ylim_epsilon, np.pi + ylim_epsilon)

    return ax

def plot_signal( signal, sample_rate = 1, ax=None, time=None, time_unit="s", **kwargs):
    if ax is None:
        ax = plt.gca()

    if time is None:
        time = np.arange(len(signal))/sample_rate

    ax.set_title("Signal")
    ax.set_xlabel("t" + (" ["+time_unit+"]" if time_unit else "" ))
    ax.set_ylabel("A(t)")

    ax.plot(time, signal, **kwargs)

    return ax

def plot_combined_spectrum(spectrum, freqs,
                           spectrum_kwargs={}, fig=None, gs=None, freq_scaler=1, freq_unit="Hz"):
    """Plot both the frequencies and phase in one figure."""

    # configure plotting layout
    if fig is None:
        fig = plt.figure(figsize=(8, 16))

    if gs is None:
        gs = gridspec.GridSpec(2, 1, figure=fig, height_ratios=[3,1], hspace=0)

    ax1 = fig.add_subplot(gs[:-1, -1])
    ax2 = fig.add_subplot(gs[-1, -1], sharex=ax1)

    axes = np.array([ax1, ax2])

    # plot the spectrum
    plot_spectrum(spectrum, freqs, ax=ax1, freq_scaler=freq_scaler, freq_unit=freq_unit, **spectrum_kwargs)

    # plot the phase
    plot_phase(spectrum, freqs, ax=ax2, freq_scaler=freq_scaler, freq_unit=freq_unit)

    ax1.xaxis.tick_top()
    [label.set_visible(False) for label in ax1.get_xticklabels()]

    return fig, axes


def phasemod(phase, low=np.pi):
    """
    Modulo phase such that it falls within the
    interval $[-low, 2\pi - low)$.
    """
    return (phase + low) % (2*np.pi) - low

def save_all_figs_to_path(fnames, figs=None, default_basename=__file__, default_extensions=['.pdf', '.png']):
    if figs is None:
        figs = [plt.figure(i) for i in plt.get_fignums()]

    default_basename = path.basename(default_basename)

    # singular value
    if isinstance(fnames, (str, True)):
        fnames = [fnames]

    if len(fnames) == len(figs):
        fnames_list = zip(figs, fnames, False)
    elif len(fnames) == 1:
        tmp_fname = fnames[0] #needed for generator
        fnames_list = ( (fig, tmp_fname, len(figs) > 1) for fig in figs)
    else:
        # outer product magic
        fnames_list = ( (fig,fname, False) for fname in fnames for fig in figs )
    del fnames
    # format fnames
    pad_width = max(2, int(np.floor(np.log10(len(figs))+1)))

    fig_fnames = []
    for fig, fnames, append_num in fnames_list:
        if not hasattr(fnames, '__len__') or isinstance(fnames, str):
            # single name
            fnames = [fnames]

        new_fnames = []
        for fname in fnames:
            if path.isdir(fname):
                fname = path.join(fname, path.splitext(default_basename)[0]) # leave off extension
            if append_num is True:
               fname += ("_fig{:0"+str(pad_width)+"d}").format(fig.number)

            if not path.splitext(fname)[1]: # no extension
                for ext in default_extensions:
                    new_fnames.append(fname+ext)
            else:
                new_fnames.append(fname)

        fig_fnames.append(new_fnames)

    # save files
    for fnames, fig in zip(fig_fnames, figs):
        for fname in fnames:
            fig.savefig(fname, transparent=True)

def sine_fitfunc(t, amp=1, freq=1, phase=0, off=0):
    """Simple sine wave for fitting purposes"""
    return amp*np.sin( 2*np.pi*freq*t + phase) + off

def sampled_time(sample_rate=1, start=0, end=1, offset=0):
    return offset + np.arange(start, end, 1/sample_rate)


def bandpass_mask(freqs, band=passband()):
    low_pass = abs(freqs) <= band[1]
    high_pass = abs(freqs) >= band[0]

    return low_pass & high_pass

def bandsize(band = passband()):
    return band[1] - band[0]

def bandlevel(samples, samplerate=1, band=passband(), normalise_bandsize=True, **ft_kwargs):
    fft, freqs = ft_spectrum(samples, samplerate, **ft_kwargs)

    bandmask = bandpass_mask(freqs, band=band)

    if normalise_bandsize:
        bins = np.count_nonzero(bandmask, axis=-1)
    else:
        bins = 1

    level = np.sum(np.abs(fft[bandmask])**2)

    return level/bins


def noisy_sine_sampling(time, init_params, noise_sigma=1, rng=rng):
    if init_params[2] is None:
        init_params[2] = phasemod(2*np.pi*rng.random())

    samples = sine_fitfunc(time, *init_params)
    noise = rng.normal(0, noise_sigma, size=len(samples))


    return samples, noise

def main(
        N = 1,
        f_sample = 250e6, # Hz
        t_length = 1e4 * 1e-9, # s

        noise_band = passband(30e6, 80e6),
        noise_sigma = 1,

    # signal properties
        f_sine = 50e6,
        signal_band = passband(50e6 - 1e6, 50e6 + 1e6),
        sine_amp = 0.2,
        sine_offset = 0,
        return_ranges_plot = False,
        cut_signal_band_from_noise_band = False
    ):
    N = int(N)

    init_params = np.array([sine_amp, f_sine, None, sine_offset])

    axs = None
    snrs = np.zeros( N )
    time = sampled_time(f_sample, end=t_length)
    for j in range(N):
        samples, noise = noisy_sine_sampling(time, init_params, noise_sigma)


        # determine signal to noise
        noise_level = bandlevel(noise, f_sample, noise_band)
        if cut_signal_band_from_noise_band:
            lower_noise_band = passband(noise_band[0], signal_band[0])
            upper_noise_band = passband(signal_band[1], noise_band[1])

            noise_level = bandlevel(noise, f_sample, lower_noise_band)
            noise_level += bandlevel(noise, f_sample, upper_noise_band)

        signal_level = bandlevel(samples, f_sample, signal_band)

        snrs[j] = np.sqrt(signal_level/noise_level)

        # make a nice plot showing what ranges were taken
        # and the bandlevels associated with them
        if return_ranges_plot and j == 0:
            combined_fft, freqs = ft_spectrum(samples+noise, f_sample)

            # plot the original signal
            if False:
                _, ax = plt.subplots()
                ax = plot_signal(samples+noise, sample_rate=f_sample/1e6, time_unit='us', ax=ax)

            # plot the spectrum
            if True:
                freq_scaler=1e6
                _, axs = plot_combined_spectrum(combined_fft, freqs, freq_scaler=freq_scaler, freq_unit='MHz')

                # indicate band ranges and frequency
                for ax in axs:
                    ax.axvline(f_sine/freq_scaler, color='r', alpha=0.4)
                    ax.axvspan(noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, color='purple', alpha=0.3, label='noiseband')
                    ax.axvspan(signal_band[0]/freq_scaler, signal_band[1]/freq_scaler, color='orange', alpha=0.3, label='signalband')

                # indicate initial phase
                axs[1].axhline(init_params[2], color='r', alpha=0.4)

                # plot the band levels
                levelax = axs[0].twinx()
                levelax.set_ylabel("Bandlevel")
                levelax.hlines(signal_level, noise_band[0]/freq_scaler, signal_band[1]/freq_scaler, colors=['orange'])
                levelax.hlines(noise_level, noise_band[0]/freq_scaler, noise_band[1]/freq_scaler, colors=['purple'])
                levelax.set_ylim(bottom=0)

            axs[0].legend()

            # plot signal_band pass signal
            if False:
                freqs = np.fft.fftfreq(len(samples), 1/f_sample)
                bandmask = bandpass_mask(freqs, band=signal_band)
                fft = np.fft.fft(samples)
                fft[ ~bandmask ] = 0
                bandpassed_samples = np.fft.ifft(fft)

                _, ax3 = plt.subplots()
                ax3 = plot_signal(bandpassed_samples, sample_rate=f_sample/1e6, time_unit='us', ax=ax3)
                ax3.set_title("Bandpassed Signal")


    return snrs, axs


if __name__ == "__main__":
    from argparse import ArgumentParser
    import os.path as path

    rng = np.random.default_rng(1)

    parser = ArgumentParser(description=__doc__)
    parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")

    args = parser.parse_args()
    default_extensions = ['.pdf', '.png']

    if args.fname == 'none':
        args.fname = None

    ###
    t_lengths = np.linspace(1e3, 5e4)* 1e-9 # s
    N = 10e1
    f_sine = 53.3e6 # Hz
    f_sample = 250e6 # Hz

    if False:
        N = 1 # Note: keep this low, N figures will be displayed!
        N_t_length = 10
        for t_length in t_lengths[-N_t_length-1:-1]:
            snrs = np.zeros( int(N))
            for i in range(int(N)):
                delta_f = 1/t_length
                snrs[i], axs = main(
                        N=1,
                        t_length=t_length,
                        f_sample=f_sample,

                        # signal properties
                        f_sine = f_sine,
                        sine_amp = 1,
                        noise_sigma = 1,

                        noise_band = passband(30e6, 80e6),
                        signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f),

                        return_ranges_plot=True
                        )

                axs[0].set_title("SNR: {}, N:{}".format(snrs[i], t_length*f_sample))
                axs[0].set_xlim(
                        (f_sine - 20*delta_f)/1e6,
                        (f_sine + 20*delta_f)/1e6
                        )

            print(snrs, "M:",np.mean(snrs))

            plt.show(block=False)

    else:
        #original code
        my_snrs = np.zeros( (len(t_lengths), int(N)) )
        for j, t_length in enumerate(t_lengths):
            return_ranges_plot = ((j==0) and True) or ( (j==(len(t_lengths)-1)) and True)

            delta_f = 1/t_length

            my_snrs[j], axs = main(
                    N=N,
                    t_length=t_length,
                    f_sample = f_sample,

                    # signal properties
                    f_sine = f_sine,
                    sine_amp = 1,
                    noise_sigma = 1,

                    noise_band = passband(30e6, 80e6),
                    signal_band = passband(f_sine- 3*delta_f, f_sine + 3*delta_f),

                    return_ranges_plot=return_ranges_plot,
                    )

            if return_ranges_plot:
                ranges_axs = axs

        fig, axs2 = plt.subplots()
        axs2.set_xlabel("N = T*$f_s$")
        axs2.set_ylabel("SNR")

        for j, t_length in enumerate(t_lengths):
            t_length = t_length * f_sample
            axs2.plot(np.repeat(t_length, my_snrs.shape[1]), my_snrs[j], ls='none', color='blue', marker='o', alpha=max(0.01, 1/my_snrs.shape[1]))
        # plot the means
        axs2.plot(t_lengths*f_sample, np.mean(my_snrs, axis=-1), color='green', marker='*', ls='none')

    ### Save or show figures
    if not args.fname:
        # empty list, False, None
        plt.show()
    else:
        save_all_figs_to_path(args.fname, default_basename=__file__)