#!/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

    max_clock_offset = 100 if len(sys.argv) < 2 else int(sys.argv[1]) # ns
    remake_clock_offsets = True

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

    fname = "ZH_airshower/mysim.sry"

    ####
    fname_dir = path.dirname(fname)
    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)

    print(f"Modifying clocks upto {max_clock_offset}ns.")
    
    # 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 True: # uniform
            clock_offsets = max_clock_offset * (2*rng.uniform(size=N_antennas) - 1)
        else: # normal
            clock_offsets = max_clock_offset * rng.normal(0, 1, size=N_antennas)

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

            if 'traces' 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['traces'][0, :] -= tmp_offset
                else:
                    clock_offsets[i] = tmp_offset
                    continue

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

    # save to simple csv
    np.savetxt(clocks_fname, clock_offsets)

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