#!/usr/bin/env python3

__doc__ = \
"""
Generate a figure showing the alignment of an external clock
referenced to a WR clock as could be seen on an oscilloscope.
"""

import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as sig

### Functions

def plot_signal( ax, t, t_edge, s_edge, name=None, box_kw={'hatch': '/'}, line_kw={}, annotate_t_edge=True, divide_s_edge=False,**plot_kw):
    """
    Plot a signal directly on an axis.
    Uses t_edge to trigger height
    """
    if divide_s_edge:
        lower = (t > t_edge - s_edge/2)
        upper = (t > t_edge + s_edge/2)
    else:
        lower = (t > t_edge)
        upper = (t > t_edge + s_edge)


    # merge dictionaries correctly
    line_kw = { **plot_kw, **line_kw }

    # plot lower and upper lines
    l = ax.plot(t, lower, **line_kw)
    
    ## update colour when plot_kw was empty
    plot_kw = { **plot_kw, **{'color': l[0].get_color()} }
    line_kw = { **plot_kw, **line_kw }
    box_kw = { **plot_kw, **box_kw }


    l2 = ax.plot(t, upper, **line_kw)

    # plot the shaded box of width t_sigma
    l3 = ax.fill_between(t, upper, lower, **box_kw)

    # annotations
    if name is not None:
        if annotate_t_edge:
            # annotate t_edge
            y = 1
            ax.annotate("$t_\mathrm{{{}}}$".format(name),
                    xy=(t_edge, y), xytext=(t_edge-3, y),
                    va='bottom', ha='center'
            )

        # annotate s_edge
        y = 0.3
        if divide_s_edge:
            annotate_width(ax, "$\sigma_\mathrm{{{}}}$".format(name), t_edge-s_edge/2, t_edge+s_edge/2, y)
        else:
            annotate_width(ax, "$\sigma_\mathrm{{{}}}$".format(name), t_edge, t_edge+s_edge, y)

    return [l, l2, l3]

def plot_diff_time(ax, name, t_1, t_2, y, vline_kw={}, va='bottom'):
    ax.axvline(t_1, **vline_kw)
    ax.axvline(t_2, **vline_kw)

    arrow_kw = {
            'va':va,
    }

    annotate_width(ax, name, t_1, t_2, y)

def annotate_width(ax, name, x1, x2, y, text_kw={}, arrow_kw={}):
    default_arrow_kw = dict(
            xy = (x1, y),
            xytext = (x2,y),
            arrowprops = dict(
                arrowstyle="<->",
                shrinkA=False,
                shrinkB=False
            ),
    )

    default_text_kw = dict(
            va='bottom',
            ha='center',
            xy=((x1+x2)/2, y)
    )

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

    return [an1, an2]


## Main
def main(ref=(10,5), A=(40,10), B=(70,12) ):
    """
    Create a figure with two signals at times t1 and t2 (accuracy s1, s2)
    as compared to a reference timer tr (sr), with annotations.
    """
    
    t = np.linspace(0, 100, 1e3)

    t_A, s_A = A
    t_B, s_B = B
    t_ref, s_ref = ref
   
    box_kw = {
            "alpha": 0.6,
            "hatch": '\\',
    }
    line_kw = {
            }

    vline_kw = {
            "linestyle": (0,(4,2)),
            "color": "k",
            "alpha":0.8,
            }
 
    fig, axs = plt.subplots(3,1,sharex=True, gridspec_kw={'hspace': 0});

    # Overall styling
    axs[-1].set_xticks([])
    axs[-1].set_xticklabels([])
    for ax in axs:
        ax.set_ylim(-0.25, 1.25)
        ax.set_yticks([])
        ax.set_yticklabels([])
        ax.grid()
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.spines['left'].set_visible(False)

    # Create the plots
    i = -1

    # Signal A
    i+=1
    y = 0.6
    axs[i].set_ylabel("Signal A", rotation=0)
    plot_diff_time(axs[i], "$t_\\mathrm{A}}$", t_ref, t_A, y, vline_kw=vline_kw)
    plot_signal(axs[i], t, t_A, s_A, name="A", box_kw=box_kw, line_kw=line_kw, annotate_t_edge=False)

    # Reference
    i+=1
    axs[i].set_ylabel("Reference", rotation=0)
    axs[i].axvline(t_ref, **vline_kw)
    plot_signal(axs[i], t, t_ref, s_ref, name="ref", box_kw=box_kw, line_kw=line_kw, color='g')
    plot_diff_time(axs[i], "$t_\\mathrm{C}$", t_A, t_B, 0.3, vline_kw=vline_kw)

    # Signal B
    i+=1
    axs[i].set_ylabel("Signal B", rotation=0)
    plot_diff_time(axs[i], "$t_\\mathrm{B}}$", t_ref, t_B, y, vline_kw=vline_kw)
    plot_signal(axs[i], t, t_B, s_B, name="B", box_kw=box_kw, line_kw=line_kw, color='purple', annotate_t_edge=False)


    return fig, 0


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.tight_layout()
        plt.savefig(args.fname, transparent=True)
    else:
        plt.show()