#!/usr/bin/env python3

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


import matplotlib.pyplot as plt
import numpy as np

from earsim import Antenna
import lib

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

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

tx = Antenna(x=0,y=0,z=0,name='tx')
rx = Antenna(x=tx.x,y=tx.y,z=tx.z,name='rx')

# 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

    # Randomise Antenna Location
    if True:
        rx.x, rx.y, rx.z = (2*rng.uniform(size=3) -1) * 1e4

    # Randomly phased beacon
    # at Antenna
    phase = lib.phase_mod(np.pi*(2*rng.uniform(size=1) -1)) # rad
    beacon = beacon_amplitude * lib.beacon_from(tx, rx, frequency, t, t0=0, phase=phase, peak_at_t0=False, c_light=c_light, radiate_rsq=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

    calculated_phases = lib.remove_antenna_geometry_phase(tx, rx, frequency, measured_phases, c_light=c_light)

    # Save residuals
    phase_res[i] = lib.phase_mod(calculated_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("Measured phase at Antenna - geometrical phase")
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()