%pylab inline

from Geant4 import *

from IPython.display import Image

"""
Explanation: An example of GEANT4 in IPython
Version 0.1, released 18/11/2014, alpha
This is a rough example of what can be done within IPython notebook with the GEANT4 python environment. Currently this "linac head" is made of just a vacuum and two scattering foils. The beam distribution currently is just a purely mono-energetic 6 MeV beam projecting out from a point. This is very much a non-physical simulation. It is designed for example purposes only.
As a starting place I recommend fiddling with the geometry class. 
If you want to see an advanced example of Geant4 with its python bindings Christopher Poole's linac is great.
Importing GEANT4
Importing the GEANT4 module is as simple as writing:
from Geant4 import *

The rest of what is written here you might get an idea about from this lecture
End of explanation
"""


class MyDetectorConstruction(G4VUserDetectorConstruction):
    "My Detector Construction"

    def __init__(self):

        G4VUserDetectorConstruction.__init__(self)
        
        
        self.solid = {}
        self.logical = {}
        self.physical = {}
       

        self.create_world(side = 4000,
                          material = "G4_AIR")
        
        self.create_cylinder(name = "vacuum", 
                             radius = 200,
                             length = 320,
                             translation = [0,0,900],
                             material = "G4_Galactic",
                             colour = [1.,1.,1.,0.1],
                             mother = 'world')
        
        self.create_cylinder(name = "upper_scatter", 
                             radius = 10,
                             length = 0.01,
                             translation = [0,0,60],
                             material = "G4_Ta",
                             colour = [1.,1.,1.,0.7],
                             mother = 'vacuum')
        
        self.create_cylinder(name = "lower_scatter", 
                             radius = 30,
                             length = 0.01,
                             translation = [0,0,20],
                             material = "G4_Al",
                             colour = [1.,1.,1.,0.7],
                             mother = 'vacuum')
        
        self.create_applicator_aperture(name = "apature_1",
                                        inner_side = 142,
                                        outer_side = 182,
                                        thickness = 6,
                                        translation = [0,0,449],
                                        material = "G4_Fe",
                                        colour = [1,1,1,0.7],
                                        mother = 'world')
        
        self.create_applicator_aperture(name = "apature_2",
                                        inner_side = 130,
                                        outer_side = 220,
                                        thickness = 12,
                                        translation = [0,0,269],
                                        material = "G4_Fe",
                                        colour = [1,1,1,0.7],
                                        mother = 'world')
        
        self.create_applicator_aperture(name = "apature_3",
                                        inner_side = 110,
                                        outer_side = 180,
                                        thickness = 12,
                                        translation = [0,0,140],
                                        material = "G4_Fe",
                                        colour = [1,1,1,0.7],
                                        mother = 'world')
        
        self.create_applicator_aperture(name = "apature_4",
                                        inner_side = 100,
                                        outer_side = 140,
                                        thickness = 12,
                                        translation = [0,0,59],
                                        material = "G4_Fe",
                                        colour = [1,1,1,0.7],
                                        mother = 'world')
        
        self.create_applicator_aperture(name = "cutout",
                                        inner_side = 100,
                                        outer_side = 120,
                                        thickness = 6,
                                        translation = [0,0,50],
                                        material = "G4_Fe",
                                        colour = [1,1,1,0.7],
                                        mother = 'world')
        
        self.create_cube(name = "phantom",
                         side = 500,
                         translation = [0,0,-250],
                         material = "G4_WATER",
                         colour = [0,0,1,0.4],
                         mother = 'world')
    
 
    
    def create_world(self, **kwargs):
        
        material = gNistManager.FindOrBuildMaterial(kwargs['material'])
        side = kwargs['side']
        
        self.solid['world'] = G4Box("world", side/2., side/2., side/2.)
        
        self.logical['world'] = G4LogicalVolume(self.solid['world'], 
                                                material, 
                                                "world")
        
        self.physical['world'] = G4PVPlacement(G4Transform3D(), 
                                               self.logical['world'], 
                                               "world", None, False, 0)

        visual = G4VisAttributes()
        visual.SetVisibility(False)
        
        self.logical['world'].SetVisAttributes(visual)
        
    
    
    def create_cylinder(self, **kwargs):
        
        name = kwargs['name']
        radius = kwargs['radius']
        length = kwargs['length']
        translation = G4ThreeVector(*kwargs['translation'])
        material = gNistManager.FindOrBuildMaterial(kwargs['material'])
        visual = G4VisAttributes(G4Color(*kwargs['colour']))
        mother = self.physical[kwargs['mother']]
        
        
        self.solid[name] = G4Tubs(name, 0., radius, length/2., 0., 2*pi)
        
        self.logical[name] = G4LogicalVolume(self.solid[name], 
                                             material,
                                             name)
        
        self.physical[name] = G4PVPlacement(None, translation, 
                                            name, 
                                            self.logical[name],
                                            mother, False, 0)

        self.logical[name].SetVisAttributes(visual)
    
    
    def create_cube(self, **kwargs):
        
        name = kwargs['name']
        side = kwargs['side']
        translation = G4ThreeVector(*kwargs['translation'])
        material = gNistManager.FindOrBuildMaterial(kwargs['material'])
        visual = G4VisAttributes(G4Color(*kwargs['colour']))
        mother = self.physical[kwargs['mother']]
        
        self.solid[name] = G4Box(name, side/2., side/2., side/2.)
        
        self.logical[name] = G4LogicalVolume(self.solid[name], 
                                             material,
                                             name)
        
        self.physical[name] = G4PVPlacement(None, translation, 
                                            name, 
                                            self.logical[name],
                                            mother, False, 0)

        self.logical[name].SetVisAttributes(visual)
    
    
    def create_applicator_aperture(self, **kwargs):
        
        name = kwargs['name']
        inner_side = kwargs['inner_side']
        outer_side = kwargs['outer_side']
        thickness = kwargs['thickness']
        translation = G4ThreeVector(*kwargs['translation'])
        material = gNistManager.FindOrBuildMaterial(kwargs['material'])
        visual = G4VisAttributes(G4Color(*kwargs['colour']))
        mother = self.physical[kwargs['mother']]
        
        
        inner_box = G4Box("inner", inner_side/2., inner_side/2., thickness/2. + 1)
        outer_box = G4Box("outer", outer_side/2., outer_side/2., thickness/2.)

        
        self.solid[name] = G4SubtractionSolid(name, 
                                              outer_box, 
                                              inner_box)
        
        self.logical[name] = G4LogicalVolume(self.solid[name], 
                                             material,
                                             name)
        
        self.physical[name] = G4PVPlacement(None, 
                                            translation,
                                            name, 
                                            self.logical[name],
                                            mother, False, 0)

        self.logical[name].SetVisAttributes(visual)
        
        

    # -----------------------------------------------------------------
    def Construct(self): # return the world volume
        
        return self.physical['world'] 

"""
Explanation: Setting the requirements for a simulation
Now, before we start up our simulation we need to define a few classes. Geant4 absolutely requires the following three classes:

A detector geometry, this is where stuff is and what it is made of
A physics list, which is what particles to use and how to simulate them
and a primary generator, which is how do you generate your particles, what energy do they have, direction, position, type etc.

This small example only includes those three classes, however things such as outputting where energy was deposited requires further elements to the simulation such as scoring
A good overview of the base GEANT4 simulation requirements is found in the Geant4 documentation
Creating the geometry class
Everything written in the next cell defines what exists in your world. To get an idea of what is going on here is the palce to start. A list of Geant4 materials will be useful to know what you can quickly make your world out of. My recommendation would be to have a fiddle around and see what happens. That's the way I tend to learn best. Once you desire to extend yourself you might want to have a look at some of the Geant4 documentation:

learn how Geant4 handles "solids", "logicals", and "physicals"
then delve deeper into various available Geant4 solids
End of explanation
"""


# set geometry
detector = MyDetectorConstruction()
gRunManager.SetUserInitialization(detector)

"""
Explanation: Now that you have made your geometry class, time to load it up
End of explanation
"""


# set physics list
physics_list = FTFP_BERT()
gRunManager.SetUserInitialization(physics_list)

"""
Explanation: The physics list
There are a range of physics options available, generally one would define their own physics class. Just for ease at the moment I have used a standard one (and because I am still learning myself). If someone really wanted to use this for realistic results you would need to customise this yourself, especially a complicated concept called cuts. I have found in one case that the default option gave awkward PDDs.
End of explanation
"""


class MyPrimaryGeneratorAction(G4VUserPrimaryGeneratorAction):
    "My Primary Generator Action"

    def __init__(self):
        
        G4VUserPrimaryGeneratorAction.__init__(self)
        
        particle_table = G4ParticleTable.GetParticleTable()

        electron = particle_table.FindParticle(G4String("e-"))
        positron = particle_table.FindParticle(G4String("e+"))
        gamma = particle_table.FindParticle(G4String("gamma"))
        
        beam = G4ParticleGun()
        beam.SetParticleEnergy(6*MeV)
        beam.SetParticleMomentumDirection(G4ThreeVector(0,0,-1))
        beam.SetParticleDefinition(electron)
        beam.SetParticlePosition(G4ThreeVector(0,0,1005))
        
        self.particleGun = beam

        
    def GeneratePrimaries(self, event):
        
        self.particleGun.GeneratePrimaryVertex(event)


"""
Explanation: Generating the beam
Given here is a very simplistic beam. I have made it so that mono-energetic electrons are shot all in the same direction, all from a single point. This of course is not realistic.
End of explanation
"""


primary_generator_action = MyPrimaryGeneratorAction()
gRunManager.SetUserAction(primary_generator_action)

"""
Explanation: And this is now loading up the generator we have just made
End of explanation
"""


# Initialise
gRunManager.Initialize()

"""
Explanation: Initialise the simulation
Now that we have set all of the requirements for our simulation we can initialise it
End of explanation
"""


%%file macros/raytrace.mac

/vis/open RayTracer
/vis/rayTracer/headAngle 340.
/vis/rayTracer/eyePosition 200 200 250  cm
/vis/rayTracer/trace images/world.jpg

gUImanager.ExecuteMacroFile('macros/raytrace.mac')

# Show image
Image(filename="images/world.jpg")

"""
Explanation: Seeing the geometry
To see the beautiful geometry we have made we can use a raytracer macro. This first cell creates the macro file. The second cell runs the file and then displays the created png image.
End of explanation
"""


%%file macros/dawn.mac

/vis/open DAWNFILE

/vis/scene/create
/vis/scene/add/volume

/vis/scene/add/trajectories smooth
/vis/modeling/trajectories/create/drawByCharge
/vis/modeling/trajectories/drawByCharge-0/default/setDrawStepPts true
/vis/modeling/trajectories/drawByCharge-0/default/setStepPtsSize 2

/vis/scene/endOfEventAction accumulate 1000

/vis/scene/add/hits

/vis/sceneHandler/attach
#/vis/scene/add/axes 0. 0. 0. 10. cm

/vis/viewer/set/targetPoint 0.0 0.0 300.0 mm
/vis/viewer/set/viewpointThetaPhi 90 0

/vis/viewer/zoom 1

gUImanager.ExecuteMacroFile('macros/dawn.mac')

"""
Explanation: Seeing the particle tracks
We can use a dawn macro in order to print out the particle tracks that a given number of generated particles porduced.
End of explanation
"""


gRunManager.BeamOn(50)

!mv g4_00.prim images/world.prim

"""
Explanation: Once we have defined how we want to see the tracks we can beam on our pretend linac with 50 electrons.
End of explanation
"""


!dawn -d images/world.prim
!convert images/world.eps images/world.png

"""
Explanation: The beam on created a prim file which needs to be converted to png for viewing.
End of explanation
"""


Image("images/world.png")

!G4VRML_DEST_DIR=.
!G4VRMLFILE_MAX_FILE_NUM=1
!G4VRMLFILE_VIEWER=echo

gApplyUICommand("/vis/open VRML2FILE")
gRunManager.BeamOn(1)

!mv g4_00.wrl images/world.wrl

"""
Explanation: And here is our wonderful simulation
End of explanation
"""


%load_ext version_information

%version_information matplotlib, numpy

"""
Explanation: Versions
End of explanation
"""


