#!/usr/bin/env python3

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 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 = False

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