import numpy as np

import ROOT
from ROOT import gSystem


""" Script to apply cuts on Allpix-Squared Geant4 data depending
    on the charge deposited inside the detector

"""

# hard-coded detector name
DETECTOR_NAME = "detector"

# Load shared libraries of allpix
# TODO: Check that relative path is correct
gSystem.Load("../lib/libAllpixObjects.so")


def get_deposited_charge_per_event(file):
    """ Calculate the total deposited charge for each event in a ROOT file

    Args:
        file (ROOT.TFile): ROOT File containing a 'DepositedCharge' tree.

    Returns:
        np.array: List of total deposited charge e per event.
    """

    DepositedCharge = file.Get('DepositedCharge')

    deposited_charge_per_event = []

    for i in range(0, DepositedCharge.GetEntries()):
        DepositedCharge.GetEntry(i)

        Charge = getattr(DepositedCharge, DETECTOR_NAME)

        charge_list = [c.getCharge() for c in Charge]

        # Sum to obtain total charge
        deposited_charge_per_event.append(np.sum(charge_list))

    return np.asarray(deposited_charge_per_event)


def get_include_list_from_charge_maximum(deposited_charge_list, max_threshold):
    """ Get list of event indices where the total deposited charge is smaller than a threshold

    Args:
        deposited_charge_list (np.array): Total deposited charge for each event
        max_threshold (float): Threshold (maximum charge) in e

    Returns:
        np.array: Indices of events where the total deposited charge is smaller than the max_treshold
    """
    return np.where(deposited_charge_list <= max_threshold)[0]


def add_tree_from_original_file(original_file, out_file, treename, include_list):
    """ Copy a subset (given by include list) of a tree to a new file

    Args:
        original_file (ROOT.TFile): "Original" ROOT file to copy tree from
        out_file (ROOT.TFile): "Destination" ROOT file to copy tree to
        treename (str): Name of tree to copy
        include_list (np.array): Indices of events which should be copied
    """
    original_tree = original_file.Get(treename)
    # Clone tree structure without copying any events
    out_tree = original_tree.CloneTree(0)

    # Iterate over all events and copy only if the event index is in include_list
    for i in range(0, original_tree.GetEntries()):
        original_tree.GetEntry(i)
        if i in include_list:
            out_tree.Fill()

    out_file.Write()


if __name__ == '__main__':

    # Hardcoded ROOT file containing PixelPulse data from Allpix^2
    original_file = ROOT.TFile('output/data_geant.root')

    # Calculate total deposited charge for each event
    # (this takes some time for large files)
    deposited_charge = get_deposited_charge_per_event(original_file)

    # Get indices of events where the total deposited charge
    # is smaller than 10x the median total deposited charge.
    indices_to_include = get_include_list_from_charge_maximum(
        deposited_charge, 10 * np.median(deposited_charge))

    # Write new ROOT file with subset of events
    out_file = ROOT.TFile('output/data_geant_cut.root', 'RECREATE')
    add_tree_from_original_file(
        original_file, out_file, 'DepositedCharge', indices_to_include)
    add_tree_from_original_file(
        original_file, out_file, 'MCParticle', indices_to_include)
    add_tree_from_original_file(
        original_file, out_file, 'MCTrack', indices_to_include)

    out_file.Close()
    original_file.Close()