import geant4_pybind as g4
import os
import matplotlib.pyplot as plt

# Environment variable for Geant4 data files
ensdfstate_data_path = "path_to_G4ENSDFSTATEDATA"
os.environ["G4ENSDFSTATEDATA"] = ensdfstate_data_path

# NIST Manager - Global Scope
nist_manager = g4.G4NistManager.Instance()

# Sensitive Detector Class
class MySensitiveDetector(g4.G4VSensitiveDetector):
    def __init__(self, name):
        super(MySensitiveDetector, self).__init__(name)
        self.energyDeposits = []

    def ProcessHits(self, step, history):
        energy_deposit = step.GetTotalEnergyDeposit()
        if energy_deposit > 0:
            self.energyDeposits.append(energy_deposit)
        return True

# Detector Construction Class
class MyDetectorConstruction(g4.G4VUserDetectorConstruction):
    def __init__(self, sensitive_detector):
        super(MyDetectorConstruction, self).__init__()
        self.sensitive_detector = sensitive_detector
        self.world_logic = None

    def Construct(self):
        air = nist_manager.FindOrBuildMaterial("G4_AIR")

        # Define CsI(Tl) material
        csi_density = 4.51 * g4.g/g4.cm3
        csi = g4.G4Material("CsI", csi_density, 2)
        csi.AddElement(nist_manager.FindOrBuildElement("Cs"), 1)
        csi.AddElement(nist_manager.FindOrBuildElement("I"), 1)

        # World volume
        world_size = 10 * g4.m
        world_solid = g4.G4Box("World", world_size/2, world_size/2, world_size/2)
        self.world_logic = g4.G4LogicalVolume(world_solid, air, "World")
        world_phys = g4.G4PVPlacement(None, g4.G4ThreeVector(0, 0, 0), self.world_logic, "World", None, False, 0, True)

        # Detector volume (CsI(Tl))
        detector_size = g4.G4ThreeVector(1 * g4.cm, 1 * g4.cm, 2 * g4.cm)
        detector_solid = g4.G4Box("Detector", detector_size.x/2, detector_size.y/2, detector_size.z/2)
        detector_logic = g4.G4LogicalVolume(detector_solid, csi, "Detector")
        detector_position = g4.G4ThreeVector(0, 0, 0)
        detector_phys = g4.G4PVPlacement(None, detector_position, detector_logic, "Detector", self.world_logic, False, 0, True)

        # Attach the sensitive detector
        g4.G4SDManager.GetSDMpointer().AddNewDetector(self.sensitive_detector)
        detector_logic.SetSensitiveDetector(self.sensitive_detector)

        return world_phys

# Primary Generator Action Class
class MyPrimaryGeneratorAction(g4.G4VUserPrimaryGeneratorAction):
    def __init__(self):
        super(MyPrimaryGeneratorAction, self).__init__()
        self.particleGun = g4.G4ParticleGun(1)

    def GeneratePrimaries(self, event):
        self.particleGun.SetParticleDefinition(g4.G4Gamma.Definition())
        self.particleGun.SetParticleEnergy(1.17 * g4.MeV)
        self.particleGun.SetParticlePosition(g4.G4ThreeVector(0, 0, -3 * g4.cm))
        self.particleGun.SetParticleMomentumDirection(g4.G4ThreeVector(0, 0, 1))
        self.particleGun.GeneratePrimaryVertex(event)

# Create and register the sensitive detector
sensitive_detector = MySensitiveDetector("MySensitiveDetector")

# Set user-defined detector construction with the sensitive detector
detector = MyDetectorConstruction(sensitive_detector)
run_manager = g4.G4RunManager()
run_manager.SetUserInitialization(detector)

# Physics list
physics_list = g4.FTFP_BERT()
run_manager.SetUserInitialization(physics_list)

# Primary generator
primary_generator = MyPrimaryGeneratorAction()
run_manager.SetUserAction(primary_generator)

# Event Action Class
class MyEventAction(g4.G4UserEventAction):
    def __init__(self, sensitive_detector):
        super(MyEventAction, self).__init__()
        self.sensitiveDetector = sensitive_detector

    def BeginOfEventAction(self, event):
        pass

    def EndOfEventAction(self, event):
        pass

# Set user-defined event action, passing the sensitive detector
event_action = MyEventAction(sensitive_detector)
run_manager.SetUserAction(event_action)

# Visualization (optional)
vis_manager = g4.G4VisExecutive()
vis_manager.Initialize()

# Steel material definition
steel_density = 7.85 * g4.g/g4.cm3  # Typical density of steel
steel = g4.G4Material("Steel", steel_density, 3)
steel.AddElement(nist_manager.FindOrBuildElement("Fe"), 0.9843)
steel.AddElement(nist_manager.FindOrBuildElement("C"), 0.0017)
steel.AddElement(nist_manager.FindOrBuildElement("Mn"), 0.014)

# Initialize and run
run_manager.Initialize()
run_manager.BeamOn(10000)  # Increase the number of events to simulate

# Cobalt-60 source definition
cobalt60_material = nist_manager.FindOrBuildMaterial("G4_Co")

# Source dimensions
source_inner_radius = 0.0 * g4.cm
source_outer_radius = 0.5 * g4.cm  # Radius of the Cobalt-60 core
source_height = 1.0 * g4.cm  # Height of the Cobalt-60 core

# Shield dimensions
shield_thickness = 5 * g4.cm  # Thickness of the steel shield

# Cobalt-60 core
source_solid = g4.G4Tubs("Source", source_inner_radius, source_outer_radius, source_height / 2, 0, 360 * g4.deg)
source_lv = g4.G4LogicalVolume(source_solid, cobalt60_material, "SourceLV")
source_pv = g4.G4PVPlacement(None, g4.G4ThreeVector(), source_lv, "Source", detector.world_logic, False, 0)

# Steel shielding
shield_outer_radius = source_outer_radius + shield_thickness
shield_solid = g4.G4Tubs("Shield", source_outer_radius, shield_outer_radius, source_height / 2, 0, 360 * g4.deg)
shield_lv = g4.G4LogicalVolume(shield_solid, steel, "ShieldLV")
shield_pv = g4.G4PVPlacement(None, g4.G4ThreeVector(), shield_lv, "Shield", detector.world_logic, False, 0)

# After running the simulation, plot the histogram
print("Total energy deposits recorded:", len(sensitive_detector.energyDeposits))
if len(sensitive_detector.energyDeposits) > 0:
    plt.hist(sensitive_detector.energyDeposits, bins=50, range=(0, 2), color='blue', alpha=0.7)
    plt.xlabel('Energy Deposit (MeV)')
    plt.ylabel('Counts')
    plt.title('Energy Deposit Histogram')
    plt.show()
else:
    print("No energy deposits recorded")