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

"""
Add a uniformly sampled time offset
to the clock of each antenna.
"""

import numpy as np
import json
import h5py

import aa_generate_beacon as beacon

clocks_fname = 'clocks.csv'


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

    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument('max_clock_offset', nargs='?', type=float, default=25, help='(Default: %(default)d)')
    parser.add_argument('-s', '--seed', type=int, nargs='?', default=None, help='Fix seed if supplied.')
    parser.add_argument('--uniform', action='store_const', const='uniform', dest='dist_type')
    parser.add_argument('--gaussian',   action='store_const', const='gauss', dest='dist_type')
    parser.add_argument('-r','--read-clocks-file', action='store_true', dest='read_clocks_file')
    parser.add_argument('--no-save-clocks', action='store_false', dest='save_clocks')
    parser.set_defaults(dist_type='gauss')

    parser.add_argument('--data-dir', type=str, default="./data", help='Path to Data Directory. (Default: %(default)s)')

    args = parser.parse_args()

    max_clock_offset = args.max_clock_offset # ns
    remake_clock_offsets = True

    seed = args.seed
    rng = np.random.default_rng(seed)

    ####
    fname_dir = args.data_dir
    clocks_fname = path.join(fname_dir, clocks_fname)
    antennas_fname = path.join(fname_dir, beacon.antennas_fname)


    if path.isfile(clocks_fname) and not remake_clock_offsets:
        print("Clock offsets previously generated")
        sys.exit()

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

    # read in antennas
    with h5py.File(antennas_fname, 'a') as fp:
        if 'antennas' not in fp.keys():
            print("Antenna file corrupted? no antennas")
            sys.exit(1)

        group = fp['antennas']

        N_antennas = len(group.keys())

        if args.read_clocks_file and path.isfile(clocks_fname): # read clock deviations from file
            print(f"Reading clocks from {clocks_fname}.")
            clock_offsets = np.loadtxt(clocks_fname)

        elif True: # random clock deviations
            print(f"Modifying clocks upto {max_clock_offset}ns.")
            clock_offsets = np.zeros( N_antennas )
            if args.dist_type == 'uniform': # uniform
                print("Uniform distribution")
                clock_offsets = max_clock_offset * (2*rng.uniform(size=N_antennas) - 1)
            else: # normal
                print("Gaussian distribution")
                clock_offsets = max_clock_offset * rng.normal(0, 1, size=N_antennas)
        else: # Hardcoded offsets
            print("Hardcoded offsets")

            f_beacon = 51.53e-3 # GHz

            clock_offsets = np.zeros( N_antennas )
            if False:
                clock_offsets[59] = np.pi/2 / (2*np.pi*f_beacon)
            elif True:
                clock_offsets += 1/8 * 1/f_beacon
            print(clock_offsets)


        # modify time values of each antenna
        trace_key = 'filtered_traces'
        for i, name in enumerate(group.keys()):
            h5ant = group[name]
            clk_offset = clock_offsets[i]


            if trace_key not in h5ant.keys():
                print(f"Antenna file corrupted? no 'traces' in {name}")
                sys.exit(1)

            h5ant_attrs = h5ant.attrs
            if 'clock_offset' in h5ant_attrs:
                tmp_offset = h5ant_attrs['clock_offset']
                if remake_clock_offsets:
                    h5ant[trace_key][0, :] -= tmp_offset
                    h5ant['E_AxB'][0, :] -= tmp_offset
                else:
                    clock_offsets[i] = tmp_offset
                    continue

            h5ant_attrs['clock_offset'] = clk_offset
            h5ant[trace_key][0, :] += clk_offset
            h5ant['E_AxB'][0, :] += clk_offset

    # save to simple csv
    if args.save_clocks:
        np.savetxt(clocks_fname, clock_offsets)

    print("Antenna clocks modified in " + str(antennas_fname))