#!/usr/bin/env python3
# vim: fdm=indent ts=4

"""
Report best time offset per frequency for each antenna
"""

import matplotlib.pyplot as plt
import numpy as np
from os import path

import aa_generate_beacon as beacon

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

    fname = "ZH_airshower/mysim.sry"

    fig_dir = "./figures" # set None to disable saving
    show_plots = not False

    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)
    time_diffs_fname = 'time_diffs.hdf5' if not True else antennas_fname

    # create fig_dir
    if fig_dir:
        os.makedirs(fig_dir, exist_ok=True)

    # Read in antennas from file
    _, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)

    # For now only implement using one freq_name
    freq_names = antennas[0].beacon_info.keys()
    if len(freq_names) > 1:
        raise NotImplementedError

    freq_name = next(iter(freq_names))
    f_beacon = antennas[0].beacon_info[freq_name]['freq']


    # TODO: redo matrix sweeping for new timing??
    measured_antenna_time_shifts = {}
    for i, ant in enumerate(antennas):
        clock_phase_time = ant.beacon_info[freq_name]['sigma_phase_mean']/(2*np.pi*f_beacon)
        best_k_time = ant.beacon_info[freq_name]['best_k_time']

        total_clock_time = best_k_time + clock_phase_time
        measured_antenna_time_shifts[ant.name] = -1*total_clock_time

    ###
    # Compare actual vs measured time shifts
    ###
    actual_antenna_time_shifts = { a.name: a.attrs['clock_offset'] for a in sorted(antennas, key=lambda a: int(a.name)) }

    N_ant = len(antennas)

    if True:
        # keep dataset in the same ordering
        antenna_names = [int(k)-1 for k,v in actual_antenna_time_shifts.items()]
        actual_time_shifts =   np.array([ v for k,v in actual_antenna_time_shifts.items()])
        measured_time_shifts = np.array([ measured_antenna_time_shifts[k] for k,v in actual_antenna_time_shifts.items() ])

        # remove global shift
        global_shift = actual_time_shifts[0] - measured_time_shifts[0]
        actual_time_shifts -= global_shift

        for i in range(2):
            plot_residuals = i == 1
            colors = ['blue', 'orange']

            fig, axs = plt.subplots(2, 1, sharex=True)

            if True:
                phase2time = lambda x: x/(2*np.pi*f_beacon)
                time2phase = lambda x: 2*np.pi*x*f_beacon
                secax = axs[0].secondary_xaxis('top', functions=(time2phase, phase2time))
                secax.set_xlabel('Phase $2\\pi t f_{beac}$ [rad]')

            if plot_residuals:
                time_shift_residuals = measured_time_shifts - actual_time_shifts
                fig.suptitle("Difference between Measured and Actual clock offsets")
                axs[-1].set_xlabel("Antenna Time Offset Residual $\\Delta_t$ [ns]")
            else:
                fig.suptitle("Comparison Measured and Actual clock offset")
                axs[-1].set_xlabel("Antenna Time Offset $t_c = \\left(\\frac{\\Delta\\varphi}{2\\pi} + k\\right) / f_{beac}$ [ns]")

            i=0
            axs[i].set_ylabel("#")
            if plot_residuals:
                axs[i].hist(time_shift_residuals, bins='sqrt', alpha=0.8, color=colors[0])
            else:
                axs[i].hist(measured_time_shifts, bins='sqrt', density=False, alpha=0.8, color=colors[0], ls='solid' , histtype='step', label='Measured')
                axs[i].hist(actual_time_shifts,   bins='sqrt', density=False, alpha=0.8, color=colors[1], ls='dashed', histtype='step', label='Actual')


            i=1
            axs[i].set_ylabel("Antenna no.")
            if plot_residuals:
                axs[i].plot(time_shift_residuals, np.arange(N_ant), alpha=0.6, ls='none', marker='x', color=colors[0])
            else:
                axs[i].errorbar(measured_time_shifts, np.arange(N_ant), yerr=None, marker='4', alpha=0.7, ls='none', color=colors[0], label='Measured')
                axs[i].plot(actual_time_shifts,  antenna_names, ls='none', marker='3', alpha=0.8, color=colors[1], label='Actual')

                axs[i].legend()
            fig.tight_layout()

            if fig_dir:
                extra_name = "comparison"
                if plot_residuals:
                    extra_name = "residuals"
                fig.savefig(path.join(fig_dir, path.basename(__file__) + f".time.{extra_name}.pdf"))

    if show_plots:
        plt.show()
