"""
This script was originally written by Tim Huege and Felix Schlüter, and is adapted
to store HDF5 files from origin showers, which additionally applies a bandpass filter
and resamples the traces to 2 GHz sampling.

Additionally, it is possible to only store antennas at the v x v x B axis, which is useful
if we only need to work on vxvxB arm instead.
"""

from __future__ import division

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import numpy as np
import configparser
import sys
import io
import re
import argparse
import h5py
from scipy.signal import resample_poly


VERSION_MAJOR = 0
VERSION_MINOR = 5

def bandpass_filter_trace(trace, trace_sampling, f_min, f_max, sample_axis=0):
    """
    Bandpass filter a trace between `f_min` and `f_max`. The `trace` array can be multidimensional, in
    which case the `sample_axis` parameter indicates which dimension should be taken as the time
    samples (ie this parameter is passed on to the `np.fft.rfft` call).

    NB: taken from smiet.numpy.utilities

    Parameters
    ----------
    trace : np.ndarray
        The array containing the time traces to be filtered
    trace_sampling : float
        The sampling interval of the time traces, in internal units
    f_min : float
        The lower frequency to filter by, in internal units
    f_max : float
        The upper frequency to filter by, in internal units
    sample_axis : int, default=0
        The axis of `trace` which contains the time samples

    Returns
    -------
    filtered_trace : np.ndarray
        The filtered traces, in the same shape as `trace`

    Notes
    -----
    To avoid issues when the maximum of the trace is too close to the edge, all traces are first
    shifted to have their maxima more or less in the middle. After the filter has been applied,
    the traces are rolled back so that they are on same time axis as the input traces.
    """
    # Assuming `trace_sampling` has the correct internal unit, freq is already in the internal unit system
    freq = np.fft.rfftfreq(trace.shape[sample_axis], d=trace_sampling)
    freq_range = np.logical_and(freq > f_min, freq < f_max)

    # Find the median maximum sample number of the traces
    max_index = np.median(np.argmax(trace, axis=sample_axis))
    to_roll = int(trace.shape[sample_axis] / 2 - max_index)

    # Roll all traces such that max is in the middle
    roll_pulse = np.roll(trace, to_roll, axis=sample_axis)

    # FFT, filter, IFFT
    spectrum = np.fft.rfft(roll_pulse, axis=sample_axis)
    spectrum = np.apply_along_axis(
        lambda ax: ax * freq_range.astype("int"), sample_axis, spectrum
    )
    filtered = np.fft.irfft(spectrum, axis=sample_axis)

    return np.roll(filtered, -to_roll, axis=sample_axis)

def get_antenna_position_showerplane(cstrafo, antenna_position : np.ndarray, core : np.ndarray) -> np.ndarray:
    """Get the antenna position in the shower plane"""
    # modify the coordinate system of the antenna positions first
    ant_pos = np.copy(antenna_position)
    ant_pos[0] = -1 * ant_pos[1]  # north -> east
    ant_pos[1] = antenna_position[0]  # west -> north
    ant_pos[2] = antenna_position[2]  # height -> height
    return cstrafo.transform_to_vxB_vxvxB(
        ant_pos, core=core
    )

def get_transformer(f_h5_inputs):
    """
    Get the coordinate transformation object.

    The coordinate transformation object is obtained through radiotools,
    where the inputs from corsika are first converted into the Auger / NRR
    coordinate system.

    Parameters
    ----------
    f_h5_inputs : h5py.Group
        The hdf5 group containing the corsika input parameters.
    """
    Bx, Bz = f_h5_inputs.attrs["MAGNET"]
    zenith = np.deg2rad(f_h5_inputs.attrs["THETAP"][0])
    # shift the azimuth to get to Auger / NRR coordinate system
    azimuth = 3 * np.pi / 2. + np.deg2rad(f_h5_inputs.attrs["PHIP"][0])
    azimuth = rdhelp.get_normalized_angle(azimuth)
    B_inclination = np.arctan2(Bz, Bx)
    B_declination = 0
    # as in Auger / NRR coordinate system, we shift the declination by 90 deg
    magnetic_field_vector = rdhelp.spherical_to_cartesian(B_inclination + np.pi / 2, B_declination + np.pi * 0.5)

    ctrans = coordinatesystems.cstrafo(zenith, azimuth, magnetic_field_vector=magnetic_field_vector)

    return ctrans

def read_input_file(hdf5_file, inp_file):
    if not isinstance(inp_file, io.IOBase):
        inp_file = open(inp_file, "r")

    inp_dict = {}
    for i, line in enumerate(inp_file.readlines()):
        elements = line.strip().split()
        inp_dict[elements[0]] = elements[1:]
    inp_file.close()

    f_h5_inputs = hdf5_file.create_group("inputs")

    # fill general attributes from inp file
    f_h5_inputs.attrs["RUNNR"] = int(inp_dict["RUNNR"][0])
    f_h5_inputs.attrs["EVTNR"] = int(inp_dict["EVTNR"][0])
    f_h5_inputs.attrs["PRMPAR"] = int(inp_dict["PRMPAR"][0])
    f_h5_inputs.attrs["ERANGE"] = np.array(
        [float(inp_dict["ERANGE"][0]), float(inp_dict["ERANGE"][1])]
    )
    f_h5_inputs.attrs["THETAP"] = np.array(
        [float(inp_dict["THETAP"][0]), float(inp_dict["THETAP"][1])]
    )
    f_h5_inputs.attrs["PHIP"] = np.array(
        [float(inp_dict["PHIP"][0]), float(inp_dict["PHIP"][0])]
    )
    f_h5_inputs.attrs["ECUTS"] = np.array(
        [
            float(inp_dict["ECUTS"][0]),
            float(inp_dict["ECUTS"][1]),
            float(inp_dict["ECUTS"][2]),
            float(inp_dict["ECUTS"][3]),
        ]
    )
    try:
        f_h5_inputs.attrs["THIN"] = np.array(
            [
                float(inp_dict["THIN"][0]),
                float(inp_dict["THIN"][1]),
                float(inp_dict["THIN"][2]),
            ]
        )
        f_h5_inputs.attrs["THINH"] = np.array(
            [float(inp_dict["THINH"][0]), float(inp_dict["THINH"][1])]
        )
    except KeyError:
        pass
    f_h5_inputs.attrs["OBSLEV"] = float(inp_dict["OBSLEV"][0])
    f_h5_inputs.attrs["MAGNET"] = np.array(
        [float(inp_dict["MAGNET"][0]), float(inp_dict["MAGNET"][1])]
    )

    try:
        f_h5_inputs.attrs["ATMOD"] = int(inp_dict["ATMOD"][0])
    except KeyError:
        f_h5_inputs.attrs["ATMOD"] = 1  # CORSIKA default, U.S standard by Linsley

    if "ATMFILE" in inp_dict:
        atm_entry = str(inp_dict["ATMFILE"][0])
        print("ATMFILE was set in *inp, store ATMOD = -1")
        f_h5_inputs.attrs["ATMFILE"] = atm_entry
        f_h5_inputs.attrs["ATMOD"] = -1


def read_reas_file(hdf5_file, reas_file):
    if not isinstance(reas_file, io.IOBase):
        reas_file = open(reas_file, "r")

    lreas = reas_file.readlines()
    tmp = "[CoREAS]\n"
    for i, line in enumerate(lreas):
        lreas[i] = line.strip().split(";")[0].strip()
        tmp += lreas[i] + "\n"
    reas_file.close()

    configParser = configparser.RawConfigParser()
    configParser.optionxform = str
    tmp2 = io.StringIO(tmp)
    configParser.read_file(tmp2)
    items = configParser.items("CoREAS")

    f_h5_reas = hdf5_file.create_group("CoREAS")

    # store content of reas file as attributes
    for key, value in items:
        if key in [
            "CoreCoordinateNorth",
            "CoreCoordinateWest",
            "CoreCoordinateVertical",
            "TimeResolution",
            "AutomaticTimeBoundaries",
            "TimeLowerBoundary",
            "TimeUpperBoundary",
            "ResolutionReductionScale",
            "GroundLevelRefractiveIndex",
            "EventNumber",
            "RunNumber",
            "GPSSecs",
            "GPSNanoSecs",
            "CoreEastingOffline",
            "CoreNorthingOffline",
            "CoreVerticalOffline",
            "RotationAngleForMagfieldDeclination",
            "ShowerZenithAngle",
            "ShowerAzimuthAngle",
            "PrimaryParticleEnergy",
            "PrimaryParticleType",
            "DepthOfShowerMaximum",
            "DistanceOfShowerMaximum",
            "MagneticFieldStrength",
            "MagneticFieldInclinationAngle",
        ]:
            if key not in [
                "EventNumber",
                "RunNumber",
                "GPSSecs",
                "GPSNanoSecs",
                "PrimaryParticleType",
            ]:
                f_h5_reas.attrs[key] = float(value)
            else:
                f_h5_reas.attrs[key] = int(value)
        else:
            f_h5_reas.attrs[key] = value


def read_longitudinal_profile(hdf5_file, long_file):
    if not isinstance(long_file, io.IOBase):
        long_file = io.open(long_file, "r", encoding="UTF-8")

    f_h5_long = hdf5_file.create_group("atmosphere")

    lines = long_file.readlines()
    n_steps = int(lines[0].rstrip().split()[3])

    def my_replace(m):
        return str(m.group(0)[0]) + " -"

    n_data_str = io.StringIO()
    n_data_str.writelines(lines[2 : (n_steps + 2)])
    n_data_str.seek(0)

    n_data = np.genfromtxt(n_data_str)
    for i, line in enumerate(lines):
        lines[i] = re.sub("[0-9]-", my_replace, line)

    dE_data_str = io.StringIO()
    dE_data_str.writelines(lines[(n_steps + 4) : (2 * n_steps + 4)])
    dE_data_str.seek(0)
    dE_data = np.genfromtxt(dE_data_str)
    data_set = f_h5_long.create_dataset("NumberOfParticles", n_data.shape, dtype="f")
    data_set[...] = n_data
    data_set.attrs["comment"] = (
        "The collumns of the data set are: DEPTH, GAMMAS, POSITRONS, ELECTRONS, MU+, MU-, HADRONS, CHARGED, NUCLEI, CHERENKOV"
    )

    data_set = f_h5_long.create_dataset("EnergyDeposit", dE_data.shape, dtype="f")
    data_set[...] = dE_data
    data_set.attrs["comment"] = (
        "The collumns of the data set are: DEPTH, GAMMA, EM IONIZ, EM CUT, MU IONIZ, MU CUT, HADR IONIZ, HADR CUT, NEUTRINO, SUM"
    )

    # read out hillas fit
    hillas_parameter = []
    for line in lines:
        if bool(re.search("PARAMETERS", line)):
            hillas_parameter = [
                float(x) for x in line.split()[2:]
            ]  # strip aways 'PARAMETER', '='
    f_h5_long.attrs["Gaisser-Hillas-Fit"] = hillas_parameter

    long_file.close()


def read_atm_file(hdf5_file, atm_file_path):
    heights, refractive_index = np.genfromtxt(atm_file_path, unpack=True, skip_header=6)
    refractive_index_profile = np.array([heights, refractive_index]).T

    f_h5_atm = hdf5_file.create_group("atmosphere_model")

    data_set = f_h5_atm.create_dataset(
        "RefractiveIndexProfile", refractive_index_profile.shape, dtype="f"
    )
    data_set[...] = refractive_index_profile

    atm_file = open(atm_file_path, "rb")
    lines = atm_file.readlines()

    # skip first entry (0), conversion cm -> m
    h = np.array(lines[1].strip(b"\n").split()[1:], dtype=float) / 100
    a = np.array(lines[2].strip(b"\n").split(), dtype=float) * 1e4
    b = np.array(lines[3].strip(b"\n").split(), dtype=float) * 1e4
    c = np.array(lines[4].strip(b"\n").split(), dtype=float) * 1e-2
    atm_file.close()

    f_h5_atm.attrs["h"] = h
    f_h5_atm.attrs["a"] = a
    f_h5_atm.attrs["b"] = b
    f_h5_atm.attrs["c"] = c


def read_antenna_data(
    hdf5_file, list_file, antenna_folder, cstrafo, core_pos, store_only_vxvxB: bool = False, trace_params : dict = {
        "f_min": 30e6,
        "f_max": 500e6,
        "dt_resample" : 1e-9,
        "upsampling_factor": 5
    },
):
    if not isinstance(list_file, io.IOBase):
        list_file = open(list_file, "r")

    lines = [item for item in list_file.readlines() if item != "\n"]  # skip empty lines
    lines = np.array(
        [item for item in lines if not item.startswith("#")]
    )  # skip comments
    list_file.close()

    dt_original = hdf5_file.attrs["TimeResolution"]
    dt_resample = trace_params["dt_resample"]  
    upsampling_factor = trace_params["upsampling_factor"] 
    downsampling_factor = int(upsampling_factor / (dt_original / dt_resample))
    print("Resampling traces from %.1f to dt = %.1f ns (up: %d, down: %d)" % (dt_original * 1e9, dt_resample * 1e9, upsampling_factor, downsampling_factor))

    observers = hdf5_file.create_group("observers")

    for line in lines:
        ll = line.strip().split()
        antenna_position = np.array(ll[2:5], dtype=float)
        antenna_label = ll[5]

        # get position in shower plane
        ant_pos_showerplane = get_antenna_position_showerplane(cstrafo, antenna_position / 100, core_pos)  # convert to m

        if store_only_vxvxB:
            # find the vxvxB coordinate
            # it should be ~ 0 in vxB axis and > 0 in vxvxB axis
            ants_at_vxvxB = (np.abs(ant_pos_showerplane[0]) < 1) & (np.sign(ant_pos_showerplane[1]) > 0)
            if not ants_at_vxvxB:
                # print("Skipping antenna %s, not at vxvxB axis." % antenna_label)
                continue

            # print("Reading antenna %s at position %s m (shower plane: %s m)" % (antenna_label, antenna_position / 100, ant_pos_showerplane))

        antenna_file = os.path.join(antenna_folder, "raw_%s.dat" % antenna_label)

        data = np.genfromtxt(antenna_file)

        # at this point, already bandpass and downsample
        trace_time, trace_val = data[:, 0], data[:, 1:4]

        # NB: everything is in seconds, so we convert all relevant
        # quantities to seconds
        filtered_trace = bandpass_filter_trace(
            trace_val,
            trace_sampling=dt_original,
            f_min=trace_params["f_min"],  # 30 MHz
            f_max=trace_params["f_max"],  # 500 MHz
            sample_axis=0
        )
        
        downsampled_trace = resample_poly(filtered_trace, up=upsampling_factor, down=downsampling_factor, axis=0)

        # also convert the timing
        n_resampled = downsampled_trace.shape[0]
        trace_time_resampled = np.arange(n_resampled) * dt_resample + trace_time[0]
        # combine the data back
        filtered_resampled_data = np.concatenate(
            (trace_time_resampled[:, np.newaxis], downsampled_trace), axis=-1
        )
        data_set = observers.create_dataset(antenna_label, filtered_resampled_data.shape, dtype=float, compression='gzip')
        data_set[...] = filtered_resampled_data

        data_set.attrs["position"] = np.array(antenna_position, dtype=float)
        data_set.attrs["name"] = antenna_label

        # there seems to be a problem when storing a list of unicodes in an attribute (is the case for python3). followed the fix from:
        # https://github.com/h5py/h5py/issues/289
        if len(ll) > 6:
            data_set.attrs["additional_arguments"] = [a.encode("utf8") for a in ll[6:]]

    # modify the timing resolution
    hdf5_file.attrs["TimeResolution"] = dt_resample


def write_coreas_hdf5_file(
    reas_filename, output_filename, f_h5=None, atm_file=None, ignore_atm_file=False, only_vxvxB=False
):
    # create hdf5 file
    if f_h5 is None:
        f_h5 = h5py.File(
            output_filename, "w", driver="core"
        )  # "core": will only write file to disk ones it is closed
    f_h5.attrs["version_major"] = VERSION_MAJOR
    f_h5.attrs["version_minor"] = VERSION_MINOR

    # read all parameter from the SIM??????.reas file
    read_reas_file(f_h5, reas_filename)

    # print error message in case CorsikaParameterFile was not specified in the .reas file
    inp_filename = f_h5["CoREAS"].attrs["CorsikaParameterFile"]
    if len(inp_filename) == 0:
        sys.exit("CorsikaParameterFile was not specified in the .reas file, aborting!")
    finp = open(os.path.join(os.path.dirname(reas_filename), inp_filename), "r")

    # read cards from .inp file and stores them into "inputs" group in f_h5 file
    read_input_file(f_h5, finp)
    print("Read input parameters from: %s" % inp_filename)

    if "ATMFILE" in f_h5["inputs"].attrs:
        print("Found following path for ATMFILE: %s" % f_h5["inputs"].attrs["ATMFILE"])
        atm_file_full_path = f_h5["inputs"].attrs["ATMFILE"]
        atm_file_name = os.path.basename(atm_file_full_path)

        if atm_file is not None:
            print("Use user specified ATMFILE: %s" % atm_file)
            read_atm_file(f_h5, atm_file)
        elif os.path.exists(atm_file_full_path):
            print("Read ATMFILE from: %s" % atm_file_full_path)
            read_atm_file(f_h5, atm_file_full_path)
        elif os.path.exists(atm_file_name):
            print("Read ATMFILE from: %s" % atm_file_name)
            read_atm_file(f_h5, atm_file_name)
        elif ignore_atm_file:
            print("Do not read any ATMFILE")
            pass
        else:
            sys.exit(
                "Could not find the atm file in %s or %s. Full stop!"
                % (atm_file_full_path, atm_file_name)
            )

    # read in long file
    long_filename = (
        "DAT" + os.path.splitext(os.path.basename(reas_filename))[0][-6:] + ".long"
    )
    longinp = io.open(
        os.path.join(os.path.dirname(reas_filename), long_filename),
        "r",
        encoding="UTF-8",
    )
    read_longitudinal_profile(f_h5, longinp)
    print("Read longitudinal profile from: %s" % long_filename)

    # read in antenna data
    listfile = open(os.path.splitext(reas_filename)[0] + ".list", "r")
    number = os.path.splitext(os.path.basename(reas_filename))[0][3:]
    antenna_folder = os.path.join(
        os.path.dirname(reas_filename), "SIM%s_coreas" % number
    )
    cstrafo = get_transformer(f_h5["inputs"])
    core_pos = np.array(
        [
            -1 * f_h5["CoREAS"].attrs["CoreCoordinateWest"],
            f_h5["CoREAS"].attrs["CoreCoordinateNorth"],
            f_h5["inputs"].attrs["OBSLEV"],
        ]
    ) / 100  # convert to m
    read_antenna_data(f_h5["CoREAS"], listfile, antenna_folder, cstrafo, core_pos, only_vxvxB)
    print("Read antenna data from folder: %s" % antenna_folder)

    return f_h5
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="coreas hdf5 converter for origin showers")
    parser.add_argument("input_file", type=str, help="path to input reas file")

    parser.add_argument(
        "-o",
        "--outputDirectory",
        dest="output_directory",
        type=str,
        default=None,
        help="path to output hdf5 directory",
    )
    parser.add_argument(
        "-of",
        "--outputFileName",
        dest="output_filename",
        type=str,
        default=None,
        help="output filename",
    )

    parser.add_argument(
        "--only_vxvxB",
        action="store_true",
        help="Store only the v x v x B polarization component of the electric field (Default: false).",
    )

    parser.add_argument(
        "--ignoreATMfile",
        action="store_true",
        help="If a ATMFILE was used for the simulation but can"
        " not be found choose to ignore it (Default: false).",
    )

    parser.add_argument(
        "--atmfile",
        type=str,
        default=None,
        help='Path to a GDAS atm file. Only required if "ATMFILE" specified in *inp file.'
        " If not provided check for file and location specified in *inp file.",
    )

    args = parser.parse_args()

    reas_filename = args.input_file
    output_basename = os.path.splitext(os.path.basename(reas_filename))[0]
    if args.only_vxvxB:
        output_basename += "_vxvxB"
    output_filename = output_basename + ".hdf5"

    if args.output_filename is not None:
        output_filename = args.output_filename
    if args.output_directory is not None:
        output_filename = os.path.join(args.output_directory, output_filename)

    try:
        from radiotools import helper as rdhelp
        from radiotools import coordinatesystems
    except ModuleNotFoundError as e:
        sys.exit("Could not find the radiotools module: '{}'\n"
                    "You can get this module from https://github.com/nu-radio/radiotools.\n"
                    "Make sure to add it to your enviourment, e.g., PYTHONPATH too. Stopping ...".format(e))

    f_h5 = write_coreas_hdf5_file(
        reas_filename,
        output_filename,
        atm_file=args.atmfile,
        ignore_atm_file=args.ignoreATMfile,
        only_vxvxB=args.only_vxvxB,
    )
    f_h5.close()
