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

tx_fname = 'tx.json'
antennas_fname = 'antennas.hdf5'

def write_tx_file(fname, tx, f_beacon):
    with open(fname, 'w') as fp:
        return json.dump(
                dict(
                    f_beacon=f_beacon,
                    tx=dict(
                        x=tx.x,
                        y=tx.y,
                        z=tx.z,
                        name=tx.name
                        )
                    ),
                fp
                )

def read_tx_file(fname):
    with open(fname, 'r') as fp:
        data = json.load(fp)

    f_beacon = data['f_beacon']
    tx = Antenna(**data['tx'])

    return tx, f_beacon

def read_beacon_hdf5(fname):
    with h5py.File(fname, 'r') as h5:
        tx_attrs = h5['tx'].attrs
        f_beacon = tx_attrs.get('f_beacon')

        mydict = { k:tx_attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
        tx = Antenna(**mydict)

        antennas = []
        for k, ant in h5['antennas'].items():
            mydict = { k:ant.attrs.get(k) for k in ['x', 'y', 'z', 'name'] }
            antenna = Antenna(**mydict)
            antenna.t = ant['traces'][0]
            antenna.Ex = ant['traces'][1]
            antenna.Ey = ant['traces'][2]
            antenna.Ez = ant['traces'][3]
            if len(ant['traces']) > 4:
                antenna.beacon = ant['traces'][4]

            antennas.append(antenna)

    return f_beacon, tx, antennas

def init_antenna_hdf5(fname, tx = None, f_beacon = None):
    with h5py.File(fname, 'w') as fp:
        if tx is not None or f_beacon is not None:
            tx_group = fp.create_group('tx')
            tx_attrs = tx_group.attrs

            if f_beacon is not None:
                tx_attrs['f_beacon'] = f_beacon

            if tx is not None:
                tx_attrs['x'] = tx.x
                tx_attrs['y'] = tx.y
                tx_attrs['z'] = tx.z
                tx_attrs['name'] = tx.name

    return fname

def append_antenna_hdf5(fname, antenna, columns = [], name='traces', prepend_time=True, overwrite=True):
    if not overwrite:
        raise NotImplementedError

    with h5py.File(fname, 'a') as fp:
        if 'antennas' in fp.keys():
            if not overwrite:
                raise NotImplementedError
            group = fp['antennas']
        else:
            group = fp.create_group('antennas')


        if antenna.name in group:
            if not overwrite:
                raise NotImplementedError
            ant_group = group[antenna.name]
        else:
            ant_group = group.create_group(antenna.name)

        ant_attrs = ant_group.attrs
        ant_attrs['x'] = antenna.x
        ant_attrs['y'] = antenna.y
        ant_attrs['z'] = antenna.z
        ant_attrs['name'] = antenna.name

        if name in ant_group:
            if not overwrite:
                raise NotImplementedError
            del ant_group[name]

        dset = ant_group.create_dataset(name, (len(columns) + 1*prepend_time, len(columns[0])), dtype='f')

        if prepend_time:
            dset[0] = antenna.t

        for i, col in enumerate(columns, 1*prepend_time):
            dset[i] = col
    
if __name__ == "__main__":
    from os import path

    remake_tx = False

    fname = "ZH_airshower/mysim.sry"
    tx = Antenna(x=-500,y=0,z=0,name='tx')
    f_beacon = 50e-3 # GHz

    beacon_amplitudes = 1e-6*np.array([1e2, 0, 0]) # mu V/m

    ####
    fname_dir = path.dirname(fname)
    tx_fname = path.join(fname_dir, tx_fname)
    antennas_fname = path.join(fname_dir, antennas_fname)

    if not path.isfile(tx_fname) or remake_tx:
        write_tx_file(tx_fname, tx, f_beacon)
    else:
        tx, f_beacon = read_tx_file(tx_fname)

    # read in antennas
    ev = REvent(fname)
    N_antennas = len(ev.antennas)

    # initialize hdf5 file
    init_antenna_hdf5(antennas_fname, tx, f_beacon)

    # make beacon per antenna
    for i, antenna in enumerate(ev.antennas):
        beacon = lib.beacon_from(tx, antenna, f_beacon, antenna.t)

        E = np.array([antenna.Ex, antenna.Ey, antenna.Ez, beacon])
        append_antenna_hdf5( antennas_fname, antenna, E, name='orig_traces', prepend_time=True)

        # add to relevant polarisation
        for i, _ in enumerate(beacon_amplitudes):
            E[i] += beacon_amplitudes[i]*beacon

        append_antenna_hdf5( antennas_fname, antenna, E, name='traces', prepend_time=True)

    print("Antenna HDF5 file written as " + str(antennas_fname))