#!/usr/bin/env python3


from lib import util
from scipy import signal, interpolate
import matplotlib.pyplot as plt
import numpy as np

try:
    from tqdm import tqdm
except:
    tqdm = lambda x: x

rng = np.random.default_rng()

class Waveform:
    name = None
    signal = None
    dt = None

    _t = None

    def __init__(self,signal=None, dt=None, t=None, name=None):
        self.signal = signal
        self.name = name

        if t is not None:
            assert len(t) == len(signal)

            self._t = t
            self.dt = t[1] - t[0]
        elif dt is not None:
            self.dt = dt

    # Lazy evaluation of time
    @property
    def t(self):
        if self._t is None:
            return self.dt * np.arange(0, len(self.signal))
        return self._t

    @t.setter
    def t(self, value):
        self._t = value

    @t.deleter
    def t(self):
        del self._t

    def __len__():
        return len(self.signal)

def white_noise_realisation(N_samples, noise_sigma=1, rng=rng):
    return rng.normal(0, noise_sigma or 0, size=N_samples)

def antenna_bp(trace, low_bp, high_bp, dt, order=3):
    # order < 6 for stability
    fs = 1/dt
    nyq = 1* fs
    low_bp = low_bp / nyq
    high_bp = high_bp / nyq

    bp_filter = signal.butter(order, [low_bp, high_bp], 'band', fs=fs)
    bandpassed = signal.lfilter(*bp_filter, trace)

    return bandpassed

def my_correlation(in1, template):
    #
    in1_long = np.zeros( (len(in1)+2*len(template)) )
    in1_long[len(template):-len(template)] = in1

    # fill the template with zeros and copy template
    template_long = np.zeros_like(in1_long)
    template_long[len(template):2*len(template)] = template

    lags = np.arange(-len(template), len(in1) ) - len(template)

    # do the correlation jig
    corrs = np.zeros_like(lags, dtype=float)
    for i, l in enumerate(lags):
        lagged_template = np.roll(template_long, l)
        corrs[i] = np.dot(lagged_template, in1_long)

    return corrs, (in1_long, template_long, lags)

def trace_upsampler(template_signal, trace, template_t, trace_t):
    template_dt = template.t[1] - template.t[0]
    trace_dt = trace_t[1] - trace_t[0]

    upsample_factor = trace_dt/template_dt
    upsampled_trace_N = np.ceil(len(trace) * upsample_factor)

    upsample_factor = int(upsample_factor)
    upsampled_trace_N = int(upsampled_trace_N)

    # upsample trace
    upsampled_trace = np.zeros(upsampled_trace_N)
    upsampled_trace[::upsample_factor] = trace

    #upsampled_t = np.arange(trace_t[0], trace_t[-1], template_dt)
    upsampled_t = template_dt * np.arange(len(upsampled_trace)) + trace_t[0]

    return upsampled_trace, upsampled_t

def template_downsampler(template_signal, trace, template_t, trace_t, offset):
    pass


if __name__ == "__main__":
    import os
    import matplotlib
    import sys
    if os.name == 'posix' and "DISPLAY" not in os.environ:
        matplotlib.use('Agg')

    bp_freq = (30e-3, 80e-3) # GHz
    template_dt = 5e-2 # ns
    template_length = 100 # ns
    noise_sigma_factor = 1e0 # keep between 10 and 0.1

    N_residuals = 50*3 if len(sys.argv) < 2 else int(sys.argv[1])

    antenna_dt = 2 # ns
    antenna_timelength = 2048 # ns

    #
    # Create the template
    #
    template = Waveform(None, dt=template_dt, name='Template')
    _deltapeak = util.deltapeak(timelength=template_length, samplerate=1/template.dt, offset=10/template.dt)
    template.signal = antenna_bp(_deltapeak[0], *bp_freq, (np.sqrt(antenna_dt/template_dt))*template_dt) # TODO: fix sqrt constant
    template.peak_sample = _deltapeak[1]
    template.peak_time = template.dt * template.peak_sample


    if True: # show template
        fig, ax = plt.subplots()
        ax.set_title("Deltapeak and Bandpassed Template")
        ax.set_xlabel("Time [ns]")
        ax.set_ylabel("Amplitude")
        ax.plot(template.t, max(template.signal)*_deltapeak[0])
        ax.plot(template.t, template.signal)
        fig.savefig('figures/11_template_deltapeak.pdf')

        if True:
            plt.close(fig)

    #
    # Find difference between true and templated times
    #
    time_residuals = np.zeros(N_residuals)
    for j in tqdm(range(N_residuals)):
        do_plots = j==0

        # receive at antenna
        ## place the deltapeak signal at a random location
        antenna = Waveform(None, dt=antenna_dt, name='Signal')
        antenna_true_signal, antenna_peak_sample = util.deltapeak(timelength=antenna_timelength, samplerate=1/antenna.dt, offset=[0.2, 0.8], rng=rng)
        antenna.signal = antenna_true_signal
        antenna.peak_sample = antenna_peak_sample
        antenna.peak_time = antenna.dt * antenna.peak_sample

        if do_plots:
            print(f"Antenna Peak Time: {antenna.peak_time}")
            print(f"Antenna Peak Sample: {antenna.peak_sample}")

        if False: # bandpass when emitting the signal
            antenna.signal = antenna_bp(antenna.signal, *bp_freq, antenna.dt)

        if False: # flip polarisation
            antenna.signal *= -1

        ## Add noise
        noise_amplitude = max(template.signal) * noise_sigma_factor
        noise_realisation = noise_amplitude * white_noise_realisation(len(antenna.signal))

        antenna_unfiltered_signal = antenna.signal + noise_realisation
        antenna.signal = antenna_bp(antenna_unfiltered_signal, *bp_freq, antenna.dt)

        true_time_offset = antenna.peak_time - template.peak_time

        if do_plots: # show signals
            fig, axs = plt.subplots(2, sharex=True)
            axs[0].set_title("Antenna Waveform")
            axs[-1].set_xlabel("Time [ns]")
            axs[0].set_ylabel("Amplitude")
            axs[0].plot(antenna.t, antenna.signal, label='bandpassed w/ noise', alpha=0.9)
            axs[0].plot(antenna.t, antenna_unfiltered_signal, label='true signal w/ noise', alpha=0.9)
            axs[0].plot(antenna.t, antenna_true_signal, label='true signal w/o noise', alpha=0.9)
            axs[0].legend()

            axs[1].set_title("Template")
            axs[1].set_ylabel("Amplitude")
            axs[1].plot(template.t, template.signal, label='orig')
            axs[1].plot(template.t + true_time_offset, template.signal, label='true moved orig')
            axs[1].legend()

            fig.savefig('figures/11_antenna_signals.pdf')

            if True: # zoom
                wx = 100
                x0 = true_time_offset

                old_xlims = axs[0].get_xlim()
                axs[0].set_xlim( x0-wx, x0+wx)
                fig.savefig('figures/11_antenna_signals_zoom.pdf')

                # restore
                axs[0].set_xlim(*old_xlims)

            if False:
                plt.close(fig)

        axs2 = None
        if True: # upsampled trace
            upsampled_trace, upsampled_t = trace_upsampler(template.signal, antenna.signal, template.t, antenna.t)

            if do_plots: # Show upsampled traces
                fig2, axs2 = plt.subplots(1, sharex=True)
                if not hasattr(axs2, '__len__'):
                    axs2 = [axs2]

                axs2[-1].set_xlabel("Time [ns]")
                axs2[0].set_ylabel("Amplitude")
                axs2[0].plot(antenna.t, antenna.signal, marker='o', label='orig')
                axs2[0].plot(upsampled_t, upsampled_trace, label='upsampled')
                axs2[0].legend(loc='upper right')

                fig2.savefig('figures/11_upsampled.pdf')

                wx = 1e2
                x0 = upsampled_t[0] + wx - 5
                axs2[0].set_xlim(x0-wx, x0+wx)
                fig2.savefig('figures/11_upsampled_zoom.pdf')

                if True:
                    plt.close(fig2)

            # determine correlations with arguments
            lag_dt = upsampled_t[1] - upsampled_t[0]
            corrs, (out1_signal, out2_template, lags) = my_correlation(upsampled_trace, template.signal)

        else: # downsampled template
            raise NotImplementedError

            corrs, (out1_signal, out2_template, lags) = my_downsampling_correlation(template.signal, antenna.signal, template.t, antenna.t)
            lag_dt = upsampled_t[1] - upsampled_t[0]

        # Determine best correlation time
        idx = np.argmax(abs(corrs))
        best_sample_lag = lags[idx]
        best_time_lag = best_sample_lag * lag_dt
        time_residuals[j] = best_time_lag - true_time_offset

        if do_plots and axs2:
            axs2[-1].axvline(best_time_lag, color='r', alpha=0.5, linewidth=2)

        # Show the final signals correlated
        if do_plots:
            fig, axs = plt.subplots(2, sharex=True)
            ylabel_kwargs = dict(
                                #rotation=0,
                                ha='right',
                                va='center'
                            )
            axs[-1].set_xlabel("Time [ns]")

            offset_list = [
                    [best_time_lag, dict(label=template.name, color='orange')],
                    [true_time_offset, dict(label='True offset', color='green')],
                    ]

            # Signal
            i=0
            axs[i].set_ylabel("Amplitude", **ylabel_kwargs)
            axs[i].plot(antenna.t, antenna.signal, label=antenna.name)

            # Put template on an twinned axis (magnitudes are different)
            ax2 = axs[i].twinx()
            for i, offset_args in enumerate(offset_list):
                this_kwargs = offset_args[1]
                offset = offset_args[0]

                ax2.axvline(offset + len(template.signal) * (template.t[1] - template.t[0]), color=this_kwargs['color'], alpha=0.7)
                ax2.axvline(offset, color=this_kwargs['color'], alpha=0.7)
                ax2.plot(offset + template.t, template.signal, **this_kwargs)

            # Correlation
            i=1
            axs[i].set_ylabel("Correlation", **ylabel_kwargs)
            axs[i].plot(lags * lag_dt, corrs)
            for i, offset_args in enumerate(offset_list):
                this_kwargs = offset_args[1]
                offset = offset_args[0]

                axs[i].axvline(offset, ls='--', **this_kwargs)

            if True: # zoom
                wx = len(template.signal) * (template.t[1] - template.t[0])/2
                t0 = best_time_lag

                old_xlims = axs[0].get_xlim()
                axs[i].set_xlim( x0-wx, x0+3*wx)
                fig.savefig('figures/11_corrs_zoom.pdf')

                # restore
                axs[i].set_xlim(*old_xlims)

            fig.tight_layout()
            fig.legend()
            fig.savefig('figures/11_corrs.pdf')

            if False:
                plt.close(fig)

    # Make a plot of the time residuals
    fig, ax = plt.subplots()
    ax.set_title("Template Correlation Lag finding")
    ax.set_xlabel("Time Residual [ns]")
    ax.set_ylabel("#")
    ax.hist(time_residuals, bins='sqrt', density=False)

    ax.legend(title=f"template dt: {template.dt: .1e}ns\nantenna dt: {antenna.dt: .1e}ns")

    fig.savefig("figures/11_time_residual_hist.pdf")

    plt.show()