#!/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 os.path as path

from copy import deepcopy as copy

from earsim import REvent, Antenna

import aa_generate_beacon as beacon

import lib

clocks_fname = 'clocks.csv'


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

    max_clock_offset = 100# 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)
    
    # 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())

        clock_offsets = max_clock_offset * (2*rng.uniform(size=N_antennas) - 1)

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

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

            ant_attrs = ant_group.attrs
            if 'clock_offset' in ant_attrs:
                tmp_offset = ant_attrs['clock_offset']
                if remake_clock_offsets:
                    ant_group['traces'][0, :] -= tmp_offset 
                else:
                    clock_offsets[i] = tmp_offset
                    continue

            ant_attrs['clock_offset'] = clk_offset
            ant_group['traces'][0, :] += clk_offset 

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

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