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

__doc__ = \
"""
Create one/two figures exemplifiying the correlation between waveforms.
"""


import numpy as np
import matplotlib.pyplot as plt

## Correlating function
def my_correlation(in1, template, lags=None, normalise=True):
    """ From 11_pulsed_timing """
    template_length = len(template)
    in1_length = len(in1)

    if lags is None:
        lags = np.arange(-template_length+1, in1_length + 1)

    # do the correlation jig
    corrs = np.zeros_like(lags, dtype=float)
    for i, l in enumerate(lags):
        if l <= 0: # shorten template at the front
            in1_start = 0
            template_end = template_length

            template_start = -template_length - l
            in1_end = max(0, min(in1_length, -template_start)) # 0 =< l + template_length =< in1_lengt

        elif l > in1_length - template_length:
            # shorten template from the back
            in1_end = in1_length
            template_start = 0

            in1_start = min(l, in1_length)
            template_end = max(0,  in1_length - l)

        else:
            in1_start = min(l, in1_length)
            in1_end   = min(in1_start + template_length, in1_length)

            # full template
            template_start = 0
            template_end   = template_length

        # Slice in1 and template
        in1_slice = in1[in1_start:in1_end]
        template_slice = template[template_start:template_end]

        corrs[i] = np.dot(in1_slice, template_slice)

    if normalise:
        corrs /= np.amax(corrs)

    return corrs, (in1, template, lags)

## Some Functions
def boxcar(x, x_low, x_high, x2=1):
    return np.heaviside(x-x_low, x2) * (1-np.heaviside(x-x_high, x2))

def f1(x, x0=0, x1=1, x2=2, x3=3, h=1):
    """ triangle with boxcar inbetween """
    if x0 == x1:
        slant_up = 0
    else:
        norm = (x1-x0)
        slant_up = (x-x0)/norm*boxcar(x, x0, x1)

    if x2 == x3:
        slant_down = 0
    else:
        norm = (x3-x2)
        slant_down = (x3-x)/norm * boxcar(x, x2, x3)

    y = h*(slant_up + slant_down + boxcar(x, x1, x2))

    return y

def triangle(x, x0=0, w=1, h=1):
    x2 = x1 = x0
    x0 = x1 - w
    x3 = x1 + 2
    return f1(x, x0, x1, x2, x3, h)

def slanted_boxcar(x, x0=0, slant_width = 1, box_width=1, h=1):
    x1 = x0 - box_width/2
    x2 = x0 + box_width/2
    x0 = x1 - slant_width/2
    x3 = x2 + slant_width/2

    return f1(x, x0, x1, x2, x3, h)

def sine(x, x0=0, x1=1, n_periods=0.5, phase=0):
    L = x1-x0

    return np.sin(2*np.pi* x * n_periods/L + phase)*boxcar(x, x0, x1)

def f2(x, x0=0, x1=2, h=1):
    """ logarithm """
    return h*np.log(1+x-x0)*boxcar(x, x0, x1)

## Figures
def plot_waveform(y, x=None, ax=None, fig_kwargs={}, **line_kwargs):
    if ax is None:
        fig, ax = plt.subplots(**fig_kwargs)
        if x is not None:
            ax.set_xlabel("Time")
        else:
            ax.set_xlabel("Sample no.")
        ax.set_ylabel("Amplitude")
    
    if x is not None:
        y, x = x, y
    else:
        x = np.arange(0, len(y))

    ax.plot(x, y, **line_kwargs)

    return ax

def correlation_figure(y1, y2, lags=None, lag_dt=1, ax=None, fig_kwargs={}, **line_kwargs):
    """ figure """
    if ax is None:
        fig, ax = plt.subplots(**fig_kwargs)

        ax.set_xlabel("Delay [samples]")
        ax.set_ylabel("Normalised Correlation")

    corrs, (_, __, lags) = my_correlation(y1, y2, lags=lags, normalise=True)

    ax.plot(lag_dt * lags, corrs, **line_kwargs)

    return ax, (corrs, lags)


def main(N_x = 1e4, x_delay=5, separate_figures=True, fig_kwargs={}):

    x = np.linspace(-2, 8, int(N_x))

    y1 = slanted_boxcar(x, x0=0, slant_width=0, box_width=3) + 1/4*sine(x, x0=-3/2, x1=3/2, phase=np.pi/2)
    y2 = slanted_boxcar(x, x0=0+x_delay, slant_width=1, box_width=0)


    if not separate_figures:
        fig, axs = plt.subplots(2,1, **fig_kwargs)
    else:
        axs = (None, None)
        figs = []

    if True:
        line_kwargs = dict(alpha=0.7)
        ax = plot_waveform(x, y1, ax=axs[0], label="Waveform 1", fig_kwargs=fig_kwargs)
        plot_waveform(x, y2, ax=ax, label="Waveform 2")

        ax.set_xlabel("Time")
        ax.set_ylabel("Amplitude")

        ax.grid()

        ax.get_figure().tight_layout()

        figs.append(ax.get_figure())


    if True:
        lag_dt = x[1] - x[0]

        ax2, (corrs, lags) = correlation_figure(y2, y1, ax=axs[1], lag_dt=lag_dt, fig_kwargs=fig_kwargs)
        ax2.grid()

        ax2.set_xlabel("Delay")
        ax2.set_ylabel("Normalised Correlation")

        best_lag_idx = np.argmax(corrs)
        best_lag_x = lags[best_lag_idx] * lag_dt

        if True:
            wx = 4
            ax2.set_xlim(x_delay - wx, x_delay + wx)

        ax2.axvline(best_lag_x, ls='dashed', alpha=0.7, color='r')
        ax2.get_figure().tight_layout()
        figs.append(ax2.get_figure())

    if not separate_figures:
        return fig
    else:
        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', 'correlation']

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

    pfl.rcParams['font.size'] *= 1.5

    ###
    figs = main(separate_figures=True, fig_kwargs=dict(figsize=(6,4)))
   
    ### Save or show figures
    if not args.fname:
        # empty list, False, None
        plt.show()
    else:

        for i, f in enumerate(figs):
            f.tight_layout()

            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)


