#!/usr/bin/env python3

"""Generate some figures showing the alignment of clocks
in a White Rabbit system with GrandMaster setup.
"""

import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as sig
rng = np.random.default_rng(12345)

### Functions
def pps(t, t_start, width=0.5):
    """
    Generate a PPS with width $width$ and starting at $t_start.
    """

    return (t > t_start) & (t < t_start + width)


def detect_rising_edges(threshold, data):
    """
    Detect rising edges in data.

    https://stackoverflow.com/a/50365462
    """
    return np.flatnonzero((data[:-1] < threshold) & (data[1:] > threshold))+1



def first_shared_edge(x1, x2, threshold=0.3):
    try:
        length = len(x2)
    except TypeError:
        length = 1

    x1_edges = detect_rising_edges(threshold, x1)
    if length > 1:
            x2_edges = detect_rising_edges(threshold, x2)
    else:
        x2_edges = x2

    start_edge = x1_edges[x1_edges > x2_edges][0]

    return start_edge


def aligned_pps(t, clock_in, pps_in, width=None):
    t_start = t[first_shared_edge(clock_in, pps_in)]

    if width is not None:
        return pps(t, t_start, width)
    else:
        return pps(t, t_start)

## Main
def main(time_base = 10e-9):
    """
    Generate a figure showing the required GrandMaster inputs
    with an aligned PPS out and DIO clock,
    and a random input event and its timestamp.
    """

    clock_freq = 10e6 # Hz
    dio_freq = 12.5*clock_freq # Hz
    
    pps_in_early = -1.7/clock_freq #s
    pps_in_width = 1e2/clock_freq #s
    pps_out_width = 1e2/clock_freq #s
    
    t = np.linspace(-2.25*1/clock_freq, 0.8*1/clock_freq, 5000) #
    
    random_event_idx = rng.integers(len(t)*2/3, len(t)) # Somewhere within the time space
    random_event = t[random_event_idx]
    
    
    ## Create Grandmaster input signals
    clock_in = (sig.square(2*np.pi*clock_freq*t)+1)/2
    pps_in = pps(t, pps_in_early, pps_in_width)
    
    ## Determine output signal
    clock_alignment = t[first_shared_edge(clock_in, pps_in)]
    pps_out = aligned_pps(t, clock_in, pps_in, width=pps_out_width)
    dio = (sig.square(2*np.pi*dio_freq*(t - clock_alignment))+1)/2
    
    ## Random event timestamp
    timestamped_event = t[first_shared_edge(dio, random_event_idx)]


    # Create the figure
    fig, axs = plt.subplots(4,1, sharex=True, gridspec_kw={'hspace': 0}, figsize=(16,4))
    
    ## Plot signals
    i=0
    axs[i].set_ylabel("$\mathrm{PPS}_\mathrm{in}$", rotation='horizontal', ha='right', va='center')
    axs[i].plot(t, pps_in, 'purple', label="PPS in")
    
    i+=1
    axs[i].set_ylabel("GM Clock\n($10\,\mathrm{MHz}$)", rotation='horizontal', ha='right', va='center')
    axs[i].plot(t, clock_in, label='10MHz in')
    
    i+=1
    axs[i].set_ylabel("FMC DIO\n($125\,MHz$)", rotation='horizontal', ha='right', va='center')
    axs[i].plot(t, dio, 'y', label='DIO')
    axs[i].plot(random_event, 0.5, 'r*')
    axs[i].axvline(timestamped_event, color='b')
    
    i+=1
    axs[i].set_ylabel("$\mathrm{PPS}_\mathrm{out}$", rotation='horizontal', ha='right', va='center')
    axs[i].plot(t, pps_out, 'g', label="PPS out")
    
    
    ## Styling
    for ax in axs:
        ax.axvline(clock_alignment, color='r', linestyle='--')
        ax.set_ylim(-0.2, 1.2)
        ax.set_yticks([])
        ax.set_yticklabels([])
        ax.grid()

    if time_base == 10e-9:
        axs[-1].set_xlabel("Time (ns)")
    
        ticks = axs[-1].get_xticks()*10/time_base # 10 was experimentally determined
        axs[-1].set_xticklabels(np.floor(ticks))
    else:
        axs[-1].set_xlabel("Time (s)")

    return fig, (clock_alignment, random_event, timestamped_event)

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

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

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

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