#!/usr/bin/env python3
# vim: fdm=marker fmr=<<<,>>>

__doc__ = \
"""
Create one/two figures exemplifiying the fourier transform of noisy sine wave.
"""


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

rng = np.random.default_rng()

def fft_spectrum( signal, sample_rate, fft=None, freq=None):
    N_samples = len(signal)
    real_signal = np.isrealobj(signal)

    if fft is None:
        if real_signal:
            fft = ft.rfft
            freq = ft.rfftfreq
        else:
            fft = ft.fft
            freq = ft.fftfreq

    if freq is None:
        freq = ft.fftfreq

    spectrum = fft(signal) / len(signal)
    if real_signal:
        spectrum *= 2
    freqs = freq(N_samples, 1/sample_rate)

    return spectrum, freqs

def dtft_spectrum(signal, time, freqs):
    freqtime = np.outer(freqs, time)

    # determine all coefficients in one go
    c_k =  2*np.cos(2*np.pi*freqtime) / len(signal)
    s_k = -2*np.sin(2*np.pi*freqtime) / len(signal)

    # dot product gives the dtft
    result = np.dot(c_k, signal) + 1j*np.dot(s_k, signal)

    return result, freqs

## From https:
def multiple_formatter(denominator=2, number=np.pi, latex='\pi'):
    def gcd(a, b):
        while b:
            a, b = b, a%b
        return a

    def _multiple_formatter(x, pos):
        den = denominator
        num = int(np.rint(den*x/number))
        com = gcd(num,den)
        (num,den) = (int(num/com),int(den/com))
        if den==1:
            if num==0:
                return r'$0$'
            if num==1:
                return r'$%s$'%latex
            elif num==-1:
                return r'$-%s$'%latex
            else:
                return r'$%s%s$'%(num,latex)
        else:
            if num==1:
                return r'$\frac{%s}{%s}$'%(latex,den)
            elif num==-1:
                return r'$\frac{-%s}{%s}$'%(latex,den)
            else:
                return r'$\frac{%s%s}{%s}$'%(num,latex,den)
    return _multiple_formatter
#

def phase_plot(ax, phase_on_xaxis=False, limits=(-1*np.pi, +1*np.pi), major_ticks_div=1, minor_ticks_div=4):
    set_label = ax.set_ylabel
    set_lims = ax.set_ylim
    axis = ax.yaxis
    if phase_on_xaxis:
        set_label = ax.set_xlabel
        set_lims = ax.set_ylim
        axis = ax.xaxis

    set_label("Phase")
    if limits is not None:
        set_lims(*limits)
        ax.margins(y=0.2)

    if major_ticks_div:
        axis.set_major_locator(plt.MultipleLocator(np.pi / major_ticks_div))
    if minor_ticks_div:
        axis.set_minor_locator(plt.MultipleLocator(np.pi / minor_ticks_div))
    axis.set_major_formatter(plt.FuncFormatter(multiple_formatter()))

    return ax

def main(
        f_sine=0.05153, # GHz
        noise_sigma = 0.5,
        sine_amplitude = 1,
        phase=0.4,
        ):

    t_fine = np.linspace(0, 200, 500) # ns
    t_sampled = np.linspace(0, 200, 50) # ns
    sine_fine = sine_amplitude * np.cos( 2*np.pi * f_sine * t_fine + phase)

    sine_sampled = sine_amplitude * np.cos( 2*np.pi * f_sine * t_sampled + phase)
    combined = noise_sigma*rng.normal(size=len(sine_sampled)) + sine_sampled

    figs = []
    # time domain
    if True:
        fig, ax = plt.subplots()
        ax.set_xlabel("Time [ns]")
        ax.set_ylabel("Amplitude")

        ax.plot(t_fine, sine_fine, label="Clean Signal")
        ax.plot(t_sampled, combined, label="+ Noise", marker='o')

        ax.legend()
        ax.grid()

        figs.append(ax.get_figure())

    # frequency domain
    if True:
        fft_spec, fft_freqs = fft_spectrum(combined, 1/(t_sampled[1] - t_sampled[0]))

        dtft_freqs = np.linspace(fft_freqs[0], fft_freqs[-1], 500)

        dtft_spec, dtft_freqs = dtft_spectrum(combined, t_sampled, dtft_freqs)
        
        fig2 = plt.figure()
        gs = gridspec.GridSpec(2, 1, figure=fig2, height_ratios=[3,1], hspace=0)
        
        ax1 = fig2.add_subplot(gs[:-1, -1])
        ax2 = fig2.add_subplot(gs[-1, -1], sharex=ax1)
        
        axes = np.array([ax1, ax2])

        ax2.set_xlabel("Frequency [GHz]")
        if True:
            dtft_freqs *= 1e3
            fft_freqs *= 1e3
            f_sine *= 1e3
            ax2.set_xlabel("Frequency [MHz]")

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


        # indicate f_sine
        for ax in axes:
            ax.axvline(f_sine, label="$f_\mathrm{sin}$", ls='dashed', color='red')

        # plot amplitudes
        ax1.set_ylabel("Amplitude")
        ax1.set_ylim(0, None)
        ax1.yaxis.get_major_ticks()[0].set_visible(False)# suppress 0.0
        ax1.grid()
        ax1.plot(fft_freqs,  np.abs(fft_spec), marker='o', label='DFT')
        ax1.plot(dtft_freqs, np.sqrt(np.abs(dtft_spec)**2), label='DTFT')

        ax1.legend()

        # plot phases
        phase_plot(ax2, limits=[-1.1*np.pi, +1.1*np.pi])
        ax2.grid()

        ax2.plot(fft_freqs,  np.angle(fft_spec), marker='o', label='DFT')
        ax2.plot(dtft_freqs, np.angle(dtft_spec), label='DTFT')

        figs.append(fig2)

        return figs

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

    import os
    import sys
    # Append parent directory to import path so pyfiglib can be found
    sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
    import pyfiglib as pfl


    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.", default=os.getcwd())

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

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

    pfl.rcParams['font.size'] = 20
    pfl.rcParams['figure.figsize'] = (8,6)
    pfl.rcParams['figure.constrained_layout.use'] = True

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

        for i, f in enumerate(figs):

            if len(args.fname) == 1 and len(figs) != 1:
                for ext in default_extensions:
                    f.savefig(path.join(args.fname[0], default_names[i]) + "." + ext, transparent=True)
            else:
                f.savefig(args.fname[i], transparent=True)