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

__doc__ = \
"""
Two figures showing synchronising on a sine beacon and a pulse.
"""


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 _annotate_width(
        ax, name,
        x1, x2, y1=None, y2=None,
        text_dx=(0,0),
        text_kw={}, arrow_kw={}
):
    """
    Annotate a width between two points, with both an arrow between
    the points, and a text between them.

    Parameters:
    -----------
    ax: Axes
     the Axes to plot on
    name: str
     text to put on top of the arrow
    x1: float or tuple
     (horizontal) location of the first point
    x2: float or tuple
     (horizontal) location of the first point
    y1: float
     vertical location of the first point
    y2: float
     vertical location of the first point

    """
    if hasattr(x1, '__len__'):
        if y1 is None:
            y1 = x1[1]
        x1 = x1[0]

    if hasattr(x2, '__len__'):
        if y2 is None:
            y2 = x2[1]
        x2 = x2[0]

    y1 = 0 if y1 is None else y1
    y2 = y1 if y2 is None else y2

    default_arrow_kw = dict(
            xy = (x1, y1),
            xytext = (x2,y2),
            arrowprops = dict(
                arrowstyle="<->",
                shrinkA=False,
                shrinkB=False,
            ),
    )

    default_text_kw = dict(
            va='bottom',
            ha='center',
            xy=((x1+x2)/2 + text_dx[0], (y1+y2)/2 + text_dx[1])
    )

    an1 = ax.annotate("", **{**default_arrow_kw, **arrow_kw})
    an2 = ax.annotate(name, **{**default_text_kw, **text_kw})

    return [an1, an2]

def main(
        f_sine = 0.05153, # GHz
        timelength = 80, # ns
        samplerate = 1, # GHz
        phase_diff = 1.2*np.pi, # rad
        t_beacon_offset = 4.4, # ns
        ):

    t_sampled = np.arange(0, timelength, 1/samplerate)
    t_impulse = np.arange(0, timelength, 1/4 * 1/samplerate)

    beacons = [0.1 * np.cos(2*np.pi*f_sine*(t_sampled - t_beacon_offset) + phase_diff*i) for i in range(2)]

    t_beacon_delay = phase_diff/(2*np.pi*f_sine)
    beacon_ticks = np.array([ n/f_sine for n in range(1+int((t_sampled[-1] - t_sampled[0])*f_sine)) ])
    t_beacon_ticks = [ beacon_ticks + t_beacon_offset - i*t_beacon_delay for i in range(2) ]

    def gaussian(x, mu=0, sigma=1 ):
        return 1/(sigma*np.sqrt(2*np.pi)) * np.exp(- (x-mu)**2 / sigma**2 )

    impulse_time_diff = 2.1/f_sine + (1/f_sine - t_beacon_delay)
    impulse_func = lambda x, mu=0: gaussian(x, mu, 2)

    impulse_timing = t_beacon_offset + .5/f_sine + np.array([0, impulse_time_diff])
    impulses = [ impulse_func(t_impulse, impulse_time) for impulse_time in impulse_timing ]

    figs = []
    fig_kwargs = {}
    tick_kwargs = dict(color='k', alpha=0.2, ls=(0, (3,2)))
    arrow_kwargs = dict(arrowprops=dict(arrowstyle="->" ))
    arrow_text_kwargs = dict()
    arrow_y = 0.12

    if True:
        fig, axes = plt.subplots(2, 1, **{**dict(sharex=True, gridspec_kw=dict(hspace=0)), **fig_kwargs})
        text_dx = (1, 0.005)

        if False:
            for _ax in axes:
                _ax.spines[:].set_visible(False)

        axes[-1].set_xlabel("Time")
        axes[-1].set_xticks([], [])

        axes[0].set_ylabel("Reference")
        axes[1].set_ylabel("Antenna")

        for i in range(0, 2):
            axes[i].set_yticks([], [])
            axes[i].plot(t_impulse, impulses[i])
            axes[i].plot(t_sampled, beacons[i], marker='.')

            # indicate timing of ticks
            [axes[i].axvline(tick, **tick_kwargs) for tick in t_beacon_ticks[i] if tick > t_sampled[0] and tick < t_sampled[-1] ]

        # get the first ticks
        first_ticks = [ min(ticks) for ticks in t_beacon_ticks ]
        first_ticks = [ tick if tick > 0 else tick + 1/f_sine for tick in first_ticks ]
        _annotate_width(axes[1], '$t_\\varphi$', first_ticks[0], first_ticks[1], text_dx=text_dx, y1=arrow_y, text_kw=arrow_text_kwargs, arrow_kw=arrow_kwargs)

        figs.append(fig)

    if True:
        fig, axes = plt.subplots(2, 1, **{**dict(sharex=True, gridspec_kw=dict(hspace=0)), **fig_kwargs})
        text_dx = (0, 0.005)

        if False:
            for _ax in axes:
                _ax.spines[:].set_visible(False)

        axes[-1].set_xlabel("Time")
        axes[-1].set_xticks([], [])

        if not True:
            axes[0].set_ylabel("Reference")
            axes[1].set_ylabel("Antenna")

        for i in range(0, 2):
            t_delta = (i == 1) * ( t_beacon_delay - 1/f_sine )
            axes[i].set_yticks([], [])
            axes[i].plot(t_impulse + t_delta , impulses[i])
            axes[i].plot(t_sampled + t_delta , beacons[i], marker='.')

            # indicate timing of ticks
            t_delta = (i == 1) * ( t_beacon_delay - 1/f_sine )
            [axes[i].axvline(tick+t_delta, **tick_kwargs) for tick in t_beacon_ticks[i] if tick > t_sampled[0] and tick < t_sampled[-1] ]

        # get the tick before the impulse
        _diffs = np.array([ impulse_timing[i] - t_beacon_ticks[i] for i in range(0,2) ])
        _diffs[_diffs < 0 ] = np.inf # only early ticks
        impulse_ticks_idx = np.argmin(abs(_diffs), axis=1)
        impulse_ticks = [ t_beacon_ticks[i][idx] + (i==1)*t_delta for i, idx in enumerate(impulse_ticks_idx) ]

        _annotate_width(axes[1], '$kT$', impulse_ticks[0], impulse_ticks[1], text_dx=text_dx, y1=arrow_y, text_kw=arrow_text_kwargs, arrow_kw=arrow_kwargs)


        figs.append(fig)

    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.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=None)

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

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

    pfl.rcParams['font.size'] = 20
    pfl.rcParams['figure.figsize'] = (6,4)
    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)
