#!/usr/bin/env python3

"""
Generate a figure showing a beacon (by default 1 period).
"""

import matplotlib.pyplot as plt
import numpy as np

def ttl(t, t_start=0, frequency=1, width=0.5):
    """
    Generate a TTL with width $width$ and starting at $t_start.
    """

    return ( (t%frequency) >= t_start) & ( (t%frequency) < t_start + width)


def main(beacon_type, N_periods=1, overplot_sampling_rates = False, sampling_rate=6, t_start=-0.3, sampling_offset_in_rate=1, ax=None, fig_kwargs={'figsize': (4,2)}, N_t=200, plot_grid=True):

    t_end = N_periods - t_start

    if ax is None:
        fig, ax = plt.subplots(1,1, **fig_kwargs)

    time = np.linspace(t_start, t_end, N_t, endpoint=False)

    if beacon_type == 'ttl':
        sig_func = lambda t: 2*ttl(t, frequency=1) - 1
    elif beacon_type == 'sine':
        sig_func = lambda t: np.sin(2*np.pi*t)
    else:
        raise ValueError("Unknown beacon_type requested.")

    ax.plot(time, sig_func(time))
    ax.grid(plot_grid)
    ax.set_xlim(t_start, t_end)
    ax.set_ylim(-1.5, 1.5)

    if True:
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.spines['left'].set_visible(False)

        ax.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)

    if overplot_sampling_rates:
        sampling_offset = sampling_offset_in_rate * 1/sampling_rate

        sample_time = np.linspace(t_start, t_end, sampling_rate, endpoint=False)

        ax.plot(sample_time, sig_func(sample_time), label='Sampling 1', linestyle='None', marker='x')

        sample_time += sampling_offset
        ax.plot(sample_time, sig_func(sample_time), label='Sampling 2', linestyle='None', marker='+')


        fig.legend(loc="lower center", ncol=2, )


    return fig, ax, dict(
                        sig_func=sig_func,
                        time=time, 
                        sample_time=sample_time,
                        sampling_offset=sampling_offset,
                    )


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

    parser = ArgumentParser(description=__doc__)
    parser.add_argument('type', choices=['ttl','sine'])
    parser.add_argument('--periods', default=1, help='Number of periods to show', type=int)
    parser.add_argument('--with-rates', action='store_true', help='show two samplings with different phase')
    parser.add_argument('--no-grid', dest='grid', action='store_false', help='underlie a grid')
    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()

    if args.fname is not None and path.isdir(args.fname):
        args.fname = path.join(args.fname, path.splitext(path.basename(__file__))[0] + ".pdf")

    ###
    fig, _, _ = main(beacon_type=args.type, overplot_sampling_rates=args.with_rates, N_periods=args.periods, plot_grid=args.grid)

    if args.fname is not None:
        plt.savefig(args.fname)
    else:
        plt.show()