#!/usr/bin/env python
# -*- coding: utf-8 -*-

from minicalo import GeometryDescriptor
from minicalo import G4System 
from minicalo_tools import TempFileManager
import uproot
import glob
import pandas as pd
import os


def index_out_of_bounds_workaround(tbranch):
    # ugly workaround for index out of bounds error
    dfs = []
    entries = tbranch['sensor_energy'].num_entries
    
    # check each entry 
    for entry in range(entries):
        try:
            df = tbranch.arrays(entry_start=entry,entry_stop=entry+1, library='pd')
            dfs.append(df)
        except:
            continue
        
    df = pd.concat(dfs)
    return df.reset_index(drop=True)

def run_batch(
        geometry: GeometryDescriptor,
        nEvents: int,
        particleSpec ,
        minEnergy_GeV: float,
        maxEnergy_GeV: float,
        temp_filename: str,
        seed: int
    ):

    if maxEnergy_GeV < 0:
        maxEnergy_GeV = minEnergy_GeV
    
    #check if particleSpec is a list of strings or a single string and assert
    if not isinstance(particleSpec, str):
        assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings"
        for p in particleSpec:
            assert isinstance(p, str), "particleSpec must be a string or a list of strings"
    if isinstance(particleSpec, str):
        particleSpec = [particleSpec]
        
    # filename without file ending(!)
    if not temp_filename.endswith(".root"):
        filename = temp_filename+ ".root"
    else:
        filename = temp_filename

    _G4System = G4System() # create instance of G4System only here
    _G4System.init(geometry, seed)

    try:
        _G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename)
        
        # TO FIX: Geant4 adds "t<threadnumber>" to the filename, circumvent this for one thread, but this is not a good solution
        file = glob.glob(filename.replace(".root",  "*.root"))
        if len(file)>1:
            raise Exception("More than one file found! - Check mulithreading")
        else:
            filename = file[0]
        # conversion from root to pandas dataframe
        
        # TO FIX: circumvent index out of bounds error
        ttree = uproot.open(filename)
        try:
            df = ttree["Hits;1"].arrays(library="pd")
        except:
            df = index_out_of_bounds_workaround(ttree["Hits;1"])
        # save to pickle and delete root file
    except Exception as e:
        raise e
    finally:
        os.remove(filename)
    
    return {'df':df, 'geometry':geometry}

if __name__ == '__main__':
    # do stuff
    # get the first two arguments that correspond to input and output file
    import sys
    assert len(sys.argv) == 3, "Please provide input and output file"

    io = TempFileManager.slave_init(sys.argv[1], sys.argv[2])
    input_data = io.read_input()
    output_data = run_batch(**input_data)
    io.dump_output(output_data)