from minicalo import ConstructionWrapper
from minicalo import G4System as _G4System

import plotly.graph_objects as go
import numpy as np
import os
import subprocess
import uproot

from IPython.display import Image, display


class G4System(_G4System):
    def run_visualize(
        self, particleSpec: str, minEnergy_GeV: float, maxEnergy_GeV: float = -1.0
    ):
        if maxEnergy_GeV < 0:
            maxEnergy_GeV = minEnergy_GeV
        # anpassen: run 1 event and get stuff from construction wrapper
        _G4System.run_visualize(self, particleSpec, minEnergy_GeV, maxEnergy_GeV)
        self.displayEvent()

    def run_batch(
        self,
        nEvents: int,
        particleSpec: str,
        minEnergy_GeV: float,
        maxEnergy_GeV: float = -1.0,
    ):
        if maxEnergy_GeV < 0:
            maxEnergy_GeV = minEnergy_GeV
        _G4System.run_batch(self, nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV)

        # conversion from root to pandas dataframe
        ttree = uproot.open("_1234567890_Hits.root")
        df = ttree["Hits;1"].arrays(library="pd")
        return df

    def displayEvent(self, particleSpec, minEnergy_GeV, maxEnergy_GeV=-1, sensor_width=np.array([50])):
        # for loop over all layers
        event = self.run_batch(1, particleSpec, minEnergy_GeV, maxEnergy_GeV)
        to_plot = []
        material_dict = {}

        z0=0
        # loop over materials
        for layer_i, layer in enumerate(reversed(self.cw.getLayers())):
            layer_i = len(self.cw.getLayers()) - layer_i -1

            #
            # plot layers
            #
            layer_hx = sensor_width / 2.
            layer_hy = sensor_width / 2.
            layer_z = layer.thickness
            layer_material = layer.material

            # 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}
                # 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,
                    y = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * layer_hy,
                    z = np.array([0, 0, 0, 0, -layer_z, -layer_z, -layer_z, -layer_z]) + z0,
                    **ijk_cube,
                    **material_dict[layer_material]
                ))
            z0 += layer_z

            #
            # sensors
            #
            
            # if there are sensors in the current layer, add them
            if layer_i in event['sensor_layer'].to_numpy():
                is_in_layer = event['sensor_layer'].to_numpy() == layer_i
                z = event['sensor_dz'].to_numpy()[is_in_layer]
                n_sensors = len(z)
                z=z[0]
                
                xy_centers, hwidth = calculate_sensor_centers(n_sensors, sensor_width)
                
                # loop over all sensors in current layer and add them to plot
                for (x_center, y_center), energy in zip(xy_centers, event['sensor_energy'].to_numpy()[is_in_layer]):
                    to_plot.append(go.Mesh3d(
                        # 8 vertices of a cube
                        x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * hwidth + x_center,
                        y = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * hwidth + y_center,
                        z = np.array([0, 0, 0, 0, z, z, z, z]) + z0,
                        **ijk_cube,
                        flatshading=True,
                        color='black',
                        name='Sensor',
                        opacity= max(0.03, float(energy / event['total_dep_energy'].to_numpy())),
                        showlegend=False,
                    )) 
                    
                z0 += z
        # 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]],  
            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
        #
        to_plot.append(
            go.Scatter3d(
            x=[0, 0],
            y=[0, 0],
            z=[-10, -2],
            mode='lines+text',
            line=dict(color='red', width=3),  # You can change the color and width of the arrow
            text=['Incoming ' + particleSpec],
            textposition='bottom center',
            hoverinfo='text',
            showlegend=False,
        ))

        #
        # finally show plot
        #
        
        fig = go.Figure(data=[
            *to_plot
            ])
        fig.update_layout(legend=dict(x=0))
        # add legend
        fig.show() 





#
# 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 = -25 + (i + 0.5) * sensor_size
            y_center = -25 + (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',
}