#!/usr/bin/env python3

"""
Test the functions in lib concerning
beacon generation and phase measuring
work correctly together.
"""


import matplotlib.pyplot as plt
import numpy as np

import lib

seed = 12345
dt = 1 # ns
frequency = 52.12345e-3 # GHz
N = 5e2
t_clock = 0
beacon_amplitude = 1

t = np.arange(0, 10*int(1e3), dt, dtype=float)
rng = np.random.default_rng(seed)


# Vary both the base time and the phase
t_extra = 0
phase_res = np.zeros(int(N))
amp_res = np.zeros(int(N))
for i in range(int(N)):
    # Change timebase
    t -= t_extra
    t_extra = (2*rng.uniform(size=1) - 1) *1e3
    t += t_extra

    # Randomly phased beacon
    phase = lib.phase_mod(np.pi*(2*rng.uniform(size=1) -1)) # rad
    beacon = beacon_amplitude * lib.sine_beacon(frequency, t, t0=0, phase=phase, peak_at_t0=False)

    if False: # blank part of the beacon
        blank_low, blank_high = int(0.2*len(t)), int(0.4*len(t))
        t_mask = np.ones(len(t), dtype=bool)
        t_mask[blank_low:blank_high] = False

        t = t[t_mask]
        beacon = beacon[t_mask]

    # Introduce clock errors
    t += t_clock
    _, measured_phases, measured_amplitudes = lib.find_beacon_in_traces([beacon], t, frequency, frequency_fit=False)
    t -= t_clock

    # Save residuals
    phase_res[i] = lib.phase_mod(measured_phases[0] - phase)
    amp_res[i] = measured_amplitudes[0] - beacon_amplitude

###
## Present figures
###
phase2time = lambda x: x/(2*np.pi*frequency)
time2phase = lambda x: 2*np.pi*x*frequency

fig, ax = plt.subplots()
ax.set_title("Sine beacon phase determination\nwith random time shifts")
ax.set_xlabel("$\\varphi_{meas} - \\varphi_{true}$ [rad]")
ax.set_ylabel("#")
ax.hist(phase_res, bins='sqrt')
if True:
    ax.axvline( -1*lib.phase_mod(t_clock*frequency*2*np.pi), color='red', lw=5, alpha=0.8, label='true t_clock')
if True:
    secax = ax.secondary_xaxis(1.0, functions=(phase2time, time2phase))
    secax.set_xlabel('Time [ns]')
ax.legend(title='N={:.1e}'.format(len(t)))
phase_xlims = ax.get_xlim()
fig.tight_layout()

amp_xlims = None
if True:
    fig, ax = plt.subplots()
    ax.set_title("Amplitude")
    ax.set_xlabel("$A_{meas} - 1$")
    ax.set_ylabel("#")
    ax.legend(title='N={:.1e}'.format(len(t)))
    ax.hist(amp_res, bins='sqrt')
    fig.tight_layout()

    amp_xlims = ax.get_xlim()

if True:
    fig, ax = plt.subplots()
    ax.grid()
    ax.set_xlabel("Phase Res [rad]")
    ax.set_ylabel("Amplitude Res")
    sc = ax.scatter( phase_res, amp_res )
    #fig.colorbar(sc, ax=ax)

    ax.set_xlim(phase_xlims)
    if amp_xlims is not None:
        ax.set_ylim(amp_xlims)
    if True:
        ax.axvline( -1*lib.phase_mod(t_clock*frequency*2*np.pi), color='red', lw=5, alpha=0.8, label='true t_clock')
    if True:
        secax = ax.secondary_xaxis(1.0, functions=(phase2time, time2phase))
        secax.set_xlabel('Time [ns]')
    ax.legend(title='N={:.1e}'.format(len(t)))
    fig.tight_layout()

plt.show()