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

import aa_generate_beacon as beacon

import lib
from lib.earsim import REvent, Antenna


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()):
            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))