import multiprocessing
import warnings
import contextlib
import sys

def set_start_method():
    try:
        multiprocessing.set_start_method('spawn', force=True)
    except RuntimeError:
        warnings.warn("Failed to set start method to 'spawn'. It may have already been set.")
        
set_start_method()

from minicalo import GeometryDescriptor
from minicalo import G4System as _G4System

import plotly.graph_objects as go
import pandas as pd
import numpy as np
import os
import glob
import subprocess
import uproot
import time

from IPython.display import Image, display


class __G4System(_G4System):
    

    def run_batch(
        self,
        nEvents: int,
        particleSpec: str,
        minEnergy_GeV: float,
        maxEnergy_GeV: float = -1.0,
        filename: str = "",
    ):

        if maxEnergy_GeV < 0:
            maxEnergy_GeV = minEnergy_GeV

        save_file = len(filename) > 0
        # filename without file ending(!)
        filename = "._" + str(time.perf_counter_ns()) + ".root"
        _G4System.run_batch(self, 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"))
        print(file)
        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
        if save_file:
            df.to_pickle(filename.replace(".root", ".pkl"))
        os.system("rm "+filename)
        
        return df


    def displayEvent(self, logE = False, renderer=None):
        raise NotImplementedError("This method is not implemented anymore. Use the function display_event instead.")



#
# some helpers
#

def calculate_sensor_centers(X, square_size):
    sensor_size = square_size / np.sqrt(X)
    hsensor_size = sensor_size / 2.
    centers = []

    for j in range(int(np.sqrt(X))):
        for i in range(int(np.sqrt(X))):
            x_center = -square_size/2 + (i + 0.5) * sensor_size
            y_center = -square_size/2 + (j + 0.5) * sensor_size
            centers.append((x_center[0], y_center[0]))

    return centers, hsensor_size[0]


ijk_cube = {
    "i": [7, 0, 0, 0, 4, 4, 6, 6, 4, 0, 3, 2],
    "j": [3, 4, 1, 2, 5, 6, 5, 2, 0, 1, 6, 3],
    "k": [0, 7, 2, 3, 6, 7, 1, 1, 5, 5, 7, 6],
    }



col_dict = {
    # active materials
    "G4_POLYSTYRENE": 'red',
    "G4_PLASTIC_SC_VINYLTOLUENE": 'blue',
    "G4_BGO": 'green',
    "G4_LSO": 'yellow',
    "G4_LYSO": 'orange',
    "G4_CESIUM_IODIDE": 'purple',
    "G4_PbWO4": 'pink',
    "G4_Si": 'brown',
    # passive materials
    "G4_Pb": 'darkgrey',
    "G4_Fe": 'lightgrey',
    "G4_W": 'grey',
    "G4_Cu": 'dimgray',
    "G4_BRASS": 'slategrey',
}

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_mini_batch(
        cw : GeometryDescriptor,
        nEvents: int,
        particleSpec: str,
        minEnergy_GeV: float,
        maxEnergy_GeV: float = -1.0,
        batch_seed : int = 0,
        counter : int = 0):

    
    #this is now encapsuled
    from  G4Calo import __G4System
    G4System = __G4System()
    seed = int(batch_seed + counter)
        
    print(f"Running mini batch with seed {seed}")
    G4System.init(cw, seed) #counter gives random seed offset

    df = G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV,"")
    return df

def _run_mini_batch_silent(*args, **kwargs):
    #I tried
    return _run_mini_batch(*args, **kwargs)

_g4calo_threading_lock = None

def init_g4calo_threading_lock():
    global _g4calo_threading_lock
    _g4calo_threading_lock = multiprocessing.Lock()


def run_batch(
        gd : GeometryDescriptor,
        nEvents: int,
        particleSpec: str,
        minEnergy_GeV: float,
        maxEnergy_GeV: float = -1.0,
        filename: str = "",
        no_mp: bool = False):
    '''
    splits the batch in jobs depending on how many cores are available and runs mini batches in parallel
    '''
    assert nEvents > 0

    nCores = multiprocessing.cpu_count()
    #make sure to adjust cores such that at least 80 events are run per core
    nCores = min(nCores, nEvents // 80 + 1)

    if no_mp:
        nCores = 1

    print(f"Running on {nCores} cores")
    nEventsPerCore = nEvents // nCores
    print(f"Running {nEventsPerCore} events per core")
    nEventsLastCore = nEvents - nEventsPerCore * (nCores - 1)
    print(f"Running {nEventsLastCore} events on last core")
    
    nevents = [nEventsPerCore if i < nCores - 1 else nEventsLastCore for i in range(nCores)]

    with _g4calo_threading_lock if _g4calo_threading_lock is not None else contextlib.nullcontext():
        # check if file exists, if so, read last seed. If not create it
        if os.path.exists(os.path.expanduser("~/.g4calo_seeds.txt")):
            with open(os.path.expanduser("~/.g4calo_seeds.txt"), "r") as f:
                seed = int(f.read())
        else:
            seed = 0
        with open(os.path.expanduser("~/.g4calo_seeds.txt"), "w") as f:
            f.write(str(seed + nCores))

    seed += 1
    print(f"Batch seed: {seed}")
    
    #if no_mp:
    #    dfs = []
    #    for i in range(nCores):
    #        dfs.append(_run_mini_batch_silent(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i))
    #else:
    with multiprocessing.Pool(nCores) as pool:
        dfs = pool.starmap(_run_mini_batch_silent, [(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i) for i in range(nCores)])

    alldf = pd.concat(dfs)
    alldf.reset_index(drop=True, inplace=True)
    #convert alldf['total_dep_energy'] to numpy array and check if all non-zero entries are unique
    non_zero = alldf['total_dep_energy'][alldf['total_dep_energy'] != 0.]
    n_unique = np.unique(non_zero.to_numpy()).shape[0]
    n_ex_non_zero = non_zero.shape[0]

    assert (n_ex_non_zero==n_unique), f"only {n_unique} unique non-zero entries in total_dep_energy out of {n_ex_non_zero}"

    if len(filename):
        alldf.to_pickle(filename)
    
    return alldf

def _fill_event(gd : GeometryDescriptor, 
                 particleSpec: str,
                 energy: float,
                 seed: int = -1):
    
    from  G4Calo import __G4System
    G4System = __G4System()
    G4System.init(gd, seed)
    G4System.run_batch(1, particleSpec, energy, energy,"")
    return gd
        
def display_event(gd : GeometryDescriptor, 
                 particleSpec: str,
                 energy: float,
                 logE = False, renderer=None, seed = -1):
    
    #run _fill_event in forked mode using 1-core multiprocessing to avoid G4 singletons to interfere
    with multiprocessing.Pool(1) as pool:
        gd = pool.apply(_fill_event, (gd, particleSpec, energy))
        
    #use gd to plot
    # for loop over all layers
    to_plot = []
    material_dict = {}
    z0=0
     # sum up total deposited energy
    total_dep_energy = 0
    for layer in gd.getLayers():
        for sensor in layer.sensors:
            total_dep_energy += sensor.getEnergy()
    if total_dep_energy == 0:
        print("No energy deposited in calorimeter!")
        total_dep_energy = 10**-8 # to avoid division by zero
    # loop over layers, invert order
    for layer in gd.getLayers()[::-1]:
        
        layer_width = layer.nx * layer.sens_xwidth * 10. # in mm
        
        #
        # plot layers
        #
        layer_hx = layer_width / 2.
        layer_hy = layer_width / 2.
        layer_dz = layer.thickness * 10. # in mm
        layer_material = layer.material
        z0 = layer.getZ()
        
        # add material to materials if it is not already in there
        if layer_material not in material_dict.keys():
            material_dict[layer_material] = {'name': layer_material, 
                                             'color': col_dict[layer_material], 
                                             'showlegend': False, 
                                             'flatshading': True, 
                                             'opacity': 0.2 * 20. / len(gd.getLayers())}
            # add legend entry
            to_plot.append(go.Mesh3d(x=[None], y=[None], z=[None], i=[0], j=[0], k=[0], 
                                     color=material_dict[layer_material]['color'], 
                                     showlegend=True, name=layer_material))
        to_plot.append(go.Mesh3d(
                # 8 vertices of a cube
                x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * layer_hx,
                z = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * layer_hy,
                #y = np.array([0, 0, 0, 0, layer_dz, layer_dz, layer_dz, layer_dz]) + z0,
                y = (layer_dz / 2) * np.array([-1, -1, -1, -1, 1, 1, 1, 1]) + z0,
                **ijk_cube,
                **material_dict[layer_material]
            ))
        #
        # sensors
        #
        
        # if there are sensors in the current layer, add them
        if layer.sensors != []:
            z = layer.sensors[0].getdz()
            corr = 0. #layer.sensors[0].getX() - layer.sensors[0].getdx()/2. + layer_width/2.
            # loop over all sensors in current layer and add them to plot
            for sensor in layer.sensors:
                x_center = sensor.getX() - corr
                y_center = sensor.getY() - corr
                hwidth = sensor.getdx() /2.
                energy = sensor.getEnergy()
                use_energy = float(energy / total_dep_energy)
                if logE:
                    raise NotImplementedError
                    use_energy = np.log(use_energy+1.) # - np.log(total_dep_energy)
                to_plot.append(go.Mesh3d(
                    # 8 vertices of a cube
                    x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * hwidth + x_center,
                    z = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * hwidth + y_center,
                    y = (layer_dz / 2) * np.array([-1, -1, -1, -1, 1, 1, 1, 1]) + z0,
                    #y = np.array([0, 0, 0, 0, z, z, z, z]) + z0,
                    **ijk_cube,
                    flatshading=True,
                    color='black',
                    name='Sensor',
                    opacity= use_energy, #max(0.03, use_energy),
                    showlegend=False,
                )) 
               

    # add legend entry for sensors
    to_plot.append(go.Mesh3d(x=[None], y=[None], z=[None], i=[0], j=[0], k=[0], 
                                color='black', showlegend=True, name='Sensors'))
    
    # add black-white colorbar for sensor hits
    to_plot.append(go.Surface(
        z=[[0, 0], [0, 0]],  
        x=[[0, 0], [0, 0]],
        y=[[0, 0], [0, 0]],
        colorscale=[[0, 'white'], [1, 'black']],
        showscale=True,
        cmin=0,
        cmax=1,
        colorbar=dict(
            title='Fraction of total deposited Energy', 
            tickvals=[0, 1],  
            ticktext=['0', '1'],
            ticks='outside',  
            ticklen=10,  
        ),
    ))
    #
    # add red arrow for incoming particle
    #
    # Define the start and end points of the line
    start_point = [0, 0, gd.getLayers()[0].getZ()-20.]
    end_point = [0, 0, gd.getLayers()[0].getZ()-10.]
    # Create the line trace
    line_trace = go.Scatter3d(
        x=[start_point[0], end_point[0]],
        z=[start_point[1], end_point[1]],
        y=[start_point[2], end_point[2]],
        mode='lines',
        line=dict(color='red', width=5),
        name='Incoming particle',
        showlegend=True,
    )
    # Calculate the direction vector for the arrow
    direction_vector = [(end_point[0] - start_point[0]), (end_point[1] - start_point[1]), (end_point[2] - start_point[2])]
    # Create the arrowhead at the start point with the opposite direction
    arrowhead_trace = go.Cone(
        x=[end_point[0]],
        z=[end_point[1]],
        y=[end_point[2]],
        u=[direction_vector[0]],
        w=[direction_vector[1]],
        v=[direction_vector[2]],
        sizemode='scaled',
        sizeref=0.8,
        showscale=False,
        colorscale='Reds',
        opacity=1.0,
        anchor='tail',
    )
    # Create the 3D scatter plot with both traces
    to_plot.append(line_trace)
    to_plot.append(arrowhead_trace)
    #
    # finally show plot
    #

    
    
    fig = go.Figure(data=[
        *to_plot
        ])
    fig.update_layout(legend=dict(x=0))

    #rotate standard view point
    # use layer_width as it is constant for all layers and the largest dimension
    fig.update_layout(scene_camera=dict(eye=dict(x=1 * 1.5, y=-1.5, z=1 * 1.5)))
    #name the axes in the HEP way, so y and z switch names
    fig.update_layout(scene=dict(xaxis_title='x [mm]', yaxis_title='z [mm]', zaxis_title='y [mm]'))
    
    #fig.write_html("temp.html")
    #exit()
    if renderer is not None:
        fig.show(renderer=renderer) 
    else:
        fig.show() 



def to_numpy(df_column, dtype=None):
    return np.array(df_column.tolist(), dtype=dtype)