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

import h5py
from itertools import combinations, zip_longest
import matplotlib.pyplot as plt
import numpy as np

import aa_generate_beacon as beacon
import lib

def antenna_true_phases(tx, antennas, freq_name, c_light=3e8):
    """
    Determine true phases from the antenna phases.

    This removes the geometrical phase from the antenna phase.
    """
    if not hasattr(antennas, '__len__'):
        single_ant = True
        antennas = [antennas]

    true_phases = np.empty( (len(antennas)) )
    for i, ant in enumerate(antennas):
        beacon_info = ant.beacon_info[freq_name]
        measured_phase = ant.beacon_info[freq_name]['phase']
        f_beacon = ant.beacon_info[freq_name]['freq']

        geom_time = lib.geometry_time(tx, ant, c_light=c_light)
        geom_phase = geom_time * 2*np.pi*f_beacon

        true_phases[i] =  lib.phase_mod(lib.phase_mod(measured_phase) - lib.phase_mod(geom_phase) )

    return true_phases


if __name__ == "__main__":
    from os import path
    import sys

    import matplotlib
    matplotlib.use('Agg')

    fname = "ZH_airshower/mysim.sry"
    c_light = 3e8*1e-9 # m/ns
    show_plots = True

    remove_absolute_phase_offset_first_antenna = True # takes precedence
    remove_absolute_phase_offset_minimum = True

    ####
    fname_dir = path.dirname(fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)

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

    if not path.isfile(antennas_fname):
        print("Antenna file cannot be found, did you try generating a beacon?")
        sys.exit(1)

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

    # Make sure at least one beacon has been identified
    if not hasattr(antennas[0], 'beacon_info') or len(antennas[0].beacon_info) == 0:
        print(f"No analysed beacon found for {antennas[0].name}, try running the phase analysis script first.")
        sys.exit(1)

    #
    N_beacon_freqs = len(antennas[0].beacon_info)

    for freq_name in antennas[0].beacon_info.keys():
        true_phases = antenna_true_phases(tx, antennas, freq_name, c_light=c_light)

        # Remove the phase from one antenna
        # this is a free parameter
        # (only required for absolute timing)
        if remove_absolute_phase_offset_first_antenna or remove_absolute_phase_offset_minimum:
            if remove_absolute_phase_offset_first_antenna: # just take the first phase
                minimum_phase = -1*true_phases[0]
            else: # take the minimum
                minimum_phase = -1*np.min(true_phases, axis=-1)

            true_phases += minimum_phase
            true_phases = lib.phase_mod(true_phases)

        # Save to antennas in file
        with h5py.File(antennas_fname, 'a') as fp:
            h5group = fp['antennas']

            for i, ant in enumerate(antennas):
                h5ant = fp['antennas'][ant.name]

                h5beacon_freq = h5ant['beacon_info'][freq_name]
                h5beacon_freq.attrs['true_phase'] = true_phases[i]

        # Plot True Phases at their locations
        if show_plots or fig_dir:
            fig, ax = plt.subplots()
            spatial_unit=None
            fig.suptitle('True phases\nf_beacon= {:2.0f}MHz'.format(f_beacon*1e3))

            antenna_locs = list(zip(*[(ant.x, ant.y) for ant in antennas]))
            ax.set_xlabel('x' if spatial_unit is None else 'x [{}]'.format(spatial_unit))
            ax.set_ylabel('y' if spatial_unit is None else 'y [{}]'.format(spatial_unit))
            scatter_kwargs = {}
            scatter_kwargs['cmap'] = 'inferno'
            scatter_kwargs['vmin'] = -np.pi
            scatter_kwargs['vmax'] = +np.pi
            color_label='$\\varphi$'

            sc = ax.scatter(*antenna_locs, c=true_phases, **scatter_kwargs)
            fig.colorbar(sc, ax=ax, label=color_label)

            if fig_dir:
                fig.savefig(path.join(fig_dir, __file__ + f".F{freq_name}.pdf"))

    print(f"True phases written to", antennas_fname)

    if show_plots:
        plt.show()