import openmc

###############################################################################
#                             MATERIAL DEFINITION
###############################################################################
# Define materials.
# Oxide Fuel with (U,Pu)O2 20 at% enrichment Pu vector
fuel1 = openmc.Material(1, "fuel1")
fuel1.set_density('g/cm3', 10.50)
fuel1.add_nuclide('Pu238', 0.008, 'ao')
fuel1.add_nuclide('Pu239', 0.1, 'ao')
fuel1.add_nuclide('Pu240', 0.052, 'ao')
fuel1.add_nuclide('Pu241', 0.026, 'ao')
fuel1.add_nuclide('Pu242', 0.014, 'ao')
fuel1.add_nuclide('U238', 0.8, 'ao')
fuel1.add_nuclide('O16', 2, 'ao')
fuel1.temperature = 1500.0

# New fuel with 25 at% of plutonium
fuel2 = openmc.Material(2, "fuel2")
fuel2.set_density('g/cm3', 10.50)
fuel2.add_nuclide('Pu238', 0.01, 'ao')
fuel2.add_nuclide('Pu239', 0.125, 'ao')
fuel2.add_nuclide('Pu240', 0.065, 'ao')
fuel2.add_nuclide('Pu241', 0.0325, 'ao')
fuel2.add_nuclide('Pu242', 0.0175, 'ao')
fuel2.add_nuclide('U238', 0.75, 'ao')
fuel2.add_nuclide('O16', 2, 'ao')
fuel2.temperature = 1500.0

# Blanket material - depleted uranium
blanket = openmc.Material(3, "blanket")
blanket.set_density('g/cm3', 10)
blanket.add_nuclide('U238', 1, 'ao')
blanket.add_nuclide('O16', 2, 'ao')
blanket.temperature = 700.0

# Cladding  -  HT-9 ferritic-martenistic steel for SFR applications
HT9 = openmc.Material(4, name="HT9")
HT9.set_density('sum')
HT9.add_nuclide('Fe54', 4.07484175E-03, 'ao')
HT9.add_nuclide('Fe56', 6.39663011E-02, 'ao')
HT9.add_nuclide('Fe57', 1.47726085E-03, 'ao')
HT9.add_nuclide('Fe58', 1.965963E-04, 'ao')
HT9.add_nuclide('Ni58', 2.926221768E-04, 'ao')
HT9.add_nuclide('Ni60', 1.127169432E-04, 'ao')
HT9.add_nuclide('Ni61', 4.89974616E-06, 'ao')
HT9.add_nuclide('Ni62', 1.562296464E-05, 'ao')
HT9.add_nuclide('Ni64', 3.9781692E-06, 'ao')
HT9.add_nuclide('Cr52', 8.68556774E-03, 'ao')
HT9.add_nuclide('Cr53', 9.8487366E-04, 'ao')
HT9.add_nuclide('Cr54', 2.451559E-04, 'ao')
HT9.add_nuclide('Mn55', 4.5921E-04, 'ao')
HT9.add_nuclide('Mo92', 7.120717100E-05, 'ao')
HT9.add_nuclide('Mo94', 4.484140500E-05, 'ao')
HT9.add_nuclide('Mo95', 7.762708800E-05, 'ao')
HT9.add_nuclide('Mo96', 8.169466900E-05, 'ao')
HT9.add_nuclide('Mo97', 4.704672000E-05, 'ao')
HT9.add_nuclide('Mo98', 1.195280730E-04, 'ao')
HT9.temperature = 700.0

# Coolant - Sodium
coolant = openmc.Material(5, "coolant")
coolant.set_density('g/cm3', 0.86)
coolant.add_nuclide('Na23', 1.0, 'ao')
coolant.temperature = 700.0

# Gas
helium = openmc.Material(6, "helium")
helium.set_density('sum')
helium.add_nuclide('He3', 2.4044E-04, 'ao')
helium.temperature = 700.0

# Define the materials file.
mats = openmc.Materials([fuel1, fuel2, blanket, HT9, coolant, helium])

# Create materials.xml file
mats.export_to_xml()

###############################################################################
#                           MATERIAL CARD FINISHED
###############################################################################


###############################################################################
#                            BEGIN GEOMETRY CARD
###############################################################################
# Define surfaces
# Surface for fuel rod
s1 = openmc.ZCylinder(R=0.4742, surface_id=1)
s2 = openmc.ZCylinder(R=0.4893, surface_id=2)
s3 = openmc.ZCylinder(R=0.5419, surface_id=3)
s4 = openmc.ZPlane(z0=0, surface_id=4)          # bottom  of lower gas plenum
s5 = openmc.ZPlane(z0=89.9175, surface_id=5)    # top of lower gas plenum
s6 = openmc.ZPlane(z0=120.0873, surface_id=6)   # top of axial reflector - bottom of active core
s7 = openmc.ZPlane(z0=220.6533, surface_id=7)   # top of active core
s8 = openmc.ZPlane(z0=230.7099, surface_id=8)   # top of upper gas plenum
s9 = openmc.ZPlane(z0=311.1624, surface_id=9)   # top of axial reflector
s10 = openmc.ZPlane(z0=9.9209, surface_id=10)   # FA inner
s11 = openmc.ZPlane(z0=10.3734, surface_id=11)   # FA outer
s12 = openmc.ZPlane(z0=10.61025, surface_id=12)   # FA outer for Sodium

# Define fuel pin cells
fuel1_pin = openmc.Universe(name='fuel1, helium, HT9, coolant', universe_id=1)
c11 = openmc.Cell(cell_id=11, fill=fuel1, region=-s1)
c12 = openmc.Cell(cell_id=12, fill=helium, region=+s1 & -s2)
c13 = openmc.Cell(cell_id=13, fill=HT9, region=+s2 & -s3)
c14 = openmc.Cell(cell_id=14, fill=coolant, region=+s3)
fuel1_pin.add_cells((c11, c12, c13, c14))

fuel2_pin = openmc.Universe(name='fuel2, helium, HT9, coolant', universe_id=2)
c21 = openmc.Cell(cell_id=21, fill=fuel2, region=-s1)
c22 = openmc.Cell(cell_id=22, fill=helium, region=+s1 & -s2)
c23 = openmc.Cell(cell_id=23, fill=HT9, region=+s2 & -s3)
c24 = openmc.Cell(cell_id=24, fill=coolant, region=+s3)
fuel2_pin.add_cells((c21, c22, c23, c24))

blanket_pin = openmc.Universe(name='blanket, helium, HT9, coolant', universe_id=3)
c31 = openmc.Cell(cell_id=31, fill=blanket, region=-s1)
c32 = openmc.Cell(cell_id=32, fill=helium, region=+s1 & -s2)
c33 = openmc.Cell(cell_id=33, fill=HT9, region=+s2 & -s3)
c34 = openmc.Cell(cell_id=34, fill=coolant, region=+s3)
blanket_pin.add_cells((c31, c32, c33, c34))

plenum_pin = openmc.Universe(name='helium, HT9, coolant', universe_id=4)
c41 = openmc.Cell(cell_id=41, fill=helium, region=-s2)
c42 = openmc.Cell(cell_id=42, fill=HT9, region=+s2 & -s3)
c43 = openmc.Cell(cell_id=43, fill=coolant, region=+s3)
plenum_pin.add_cells((c41, c42, c43))

coolantuniverse = openmc.Universe(name='coolantuniverse', universe_id=5)
c50 = openmc.Cell(cell_id=50, fill=coolant)
coolantuniverse.add_cells([c50])

# GEOMETRY - fuel rod definition ------------------------------
# Fuel rod with MOX 20at% with bottom and top axial reflector and two gas plena
fuelrod_20 = openmc.Universe(name='fuelrod_20', universe_id=6)
c101 = openmc.Cell(cell_id=101, fill=coolant, region=-s4)             # Space below rod
c102 = openmc.Cell(cell_id=102, fill=plenum_pin, region=+s4 & -s5)    # Lower gas plenum
c103 = openmc.Cell(cell_id=103, fill=blanket_pin, region=+s5 & -s6)   # Lower axial reflector
c104 = openmc.Cell(cell_id=104, fill=fuel1_pin, region=+s6 & -s7)     # Active core
c105 = openmc.Cell(cell_id=105, fill=plenum_pin, region=+s7 & -s8)    # Upper gas plenum
c106 = openmc.Cell(cell_id=106, fill=blanket_pin, region=+s8 & -s9)   # Upper axial reflector
c107 = openmc.Cell(cell_id=107, fill=coolant, region=+s9)             # Spacer above rod
fuelrod_20.add_cells((c101, c102, c103, c104, c105, c106, c107))

# Fuel rod with 25% MOX with bottom and top axial reflector and two gas plena
fuelrod_25 = openmc.Universe(name='fuelrod_25', universe_id=7)
c201 = openmc.Cell(cell_id=201, fill=coolant, region=-s4)             # Space below rod
c202 = openmc.Cell(cell_id=202, fill=plenum_pin, region=+s4 & -s5)    # Lower gas plenum
c203 = openmc.Cell(cell_id=203, fill=blanket_pin, region=+s5 & -s6)   # Lower axial reflector
c204 = openmc.Cell(cell_id=204, fill=fuel2_pin, region=+s6 & -s7)     # Active core
c205 = openmc.Cell(cell_id=205, fill=plenum_pin, region=+s7 & -s8)    # Upper gas plenum
c206 = openmc.Cell(cell_id=206, fill=blanket_pin, region=+s8 & -s9)   # Upper axial reflector
c207 = openmc.Cell(cell_id=207, fill=coolant, region=+s9)             # Spacer above rod
fuelrod_25.add_cells((c201, c202, c203, c204, c205, c206, c207))

# Empty rod
emptyrod = openmc.Universe(name='emptyrod', universe_id=8)
c301 = openmc.Cell(cell_id=301, fill=coolant, region=-s4)             # Space below rod
c302 = openmc.Cell(cell_id=302, fill=plenum_pin, region=+s4 & -s5)    # Lower gas plenum
c303 = openmc.Cell(cell_id=303, fill=blanket_pin, region=+s5 & -s6)   # Lower axial reflector
c304 = openmc.Cell(cell_id=304, fill=coolant, region=+s6 & -s7)       # Active core
c305 = openmc.Cell(cell_id=305, fill=plenum_pin, region=+s7 & -s8)    # Upper gas plenum
c306 = openmc.Cell(cell_id=306, fill=blanket_pin, region=+s8 & -s9)   # Upper axial reflector
c307 = openmc.Cell(cell_id=307, fill=coolant, region=+s9)             # Spacer above rod
emptyrod.add_cells((c301, c302, c303, c304, c305, c306, c307))

# GEOMETRY - fuel assembly Lattice -------------------------------
#FA Geometry
FA_01_surf = openmc.get_hexagonal_prism(orientation = 'x', edge_length = 9.9209) # Inner "radius" of the assembly = (Sub-assembly duct outer flat-to-flat distance)/2 
FA_02_surf = openmc.get_hexagonal_prism(orientation = 'x', edge_length = 10.3734) # Outer "radius" of the assembly = (Sub-assembly duct outer flat-to-flat distance + Sub-assembly duct wall thickness)/2
FA_03_surf = openmc.get_hexagonal_prism(orientation = 'x', edge_length = 10.61025) # Outer "radius" of the assembly region (Assembly Pitch)/2  ; Where Assembly Pitch = 21.2205 cm
FA_04_surf = openmc.ZPlane(z0=0) # Bottom of the assembly
FA_05_surf = openmc.ZPlane(z0=311.1624) # Top of the assembly

#FA20 MOX
F20 = openmc.HexLattice(name='F20', lattice_id=55)
F20.center = (0, 0)
F20.pitch = ([1.1897])
F20.outer = coolantuniverse
F20.universes = [fuelrod_20]*54, [fuelrod_20]*48, [fuelrod_20]*42, [fuelrod_20]*36, [fuelrod_20]*30, [fuelrod_20]*24, [fuelrod_20]*18, [fuelrod_20]*12, [fuelrod_20]*6, [fuelrod_20]

#FA25 MOX
F25 = openmc.HexLattice(name='FA25')
F25.center = (0, 0)
F25.pitch = ([1.1897])
F25.outer = coolantuniverse
F25.universes = [fuelrod_25]*54, [fuelrod_25]*48, [fuelrod_25]*42, [fuelrod_25]*36, [fuelrod_25]*30, [fuelrod_25]*24, [fuelrod_25]*18, [fuelrod_25]*12, [fuelrod_25]*6, [fuelrod_25]

#Making FA20 as universe
FA20 = openmc.Universe(name='FA20')
c501 = openmc.Cell(fill=F20, region=-FA_01_surf)
c502 = openmc.Cell(fill=HT9, region=+FA_01_surf & -FA_02_surf)
c503 = openmc.Cell(fill=coolant, region=+FA_02_surf)
FA20.add_cells((c501, c502, c503))

#Making FA25 as universe
FA25 = openmc.Universe(name='FA25')
c504 = openmc.Cell(fill=F25, region=-FA_01_surf)
c505 = openmc.Cell(fill=HT9, region=+FA_01_surf & -FA_02_surf)
c506 = openmc.Cell(fill=coolant, region=+FA_02_surf)
FA25.add_cells((c504, c505, c506))

#Making F00
F00 = openmc.Universe(name='F00')
c507 = openmc.Cell(fill=coolant, region=-FA_03_surf)
F00.add_cells([c507])

#Making FSS
FSS = openmc.Universe(name='FSS')
c508 = openmc.Cell(fill=HT9, region=-FA_03_surf)
FSS.add_cells([c508])

# Geometry of the core
core = openmc.HexLattice(name='core', lattice_id=65)
core.center = (0, 0)
core.pitch = ([21.2205])
core.outer = coolantuniverse
core.universes = [FSS]*36, [FSS]*30, [FSS]*24, [FA25]*18, [FA20]*12, [FA20]*6, [FA20]

# Whole model surfaces
C01 = openmc.ZPlane(z0=-50, boundary_type='vacuum')
C02 = openmc.ZPlane(z0=0)
C03 = openmc.ZPlane(z0=311.1624)
C04 = openmc.ZPlane(z0=350, boundary_type='vacuum')
C80 = openmc.ZCylinder(R=100)
C81 = openmc.ZCylinder(R=140, boundary_type='vacuum')

# Making whole model
root = openmc.Universe()
c1 = openmc.Cell(fill=core, region=-C80 & -C03 & +C02)
c2 = openmc.Cell(fill=coolant, region=+C80 & -C81 & -C04 & +C01)
c3 = openmc.Cell(fill=coolant, region=-C80 & +C03)
c4 = openmc.Cell(fill=coolant, region=-C80 & -C02)
FA20.add_cells((c1,c2,c3,c4))

# Create geometry.xml file
geom = openmc.Geometry(root)
geom.export_to_xml()

###############################################################################
#                            GEOMETRY CARD FINISHED
###############################################################################

###############################################################################
#                                  SETTING CARD
###############################################################################

point = openmc.stats.Point((0,0,0))
src = openmc.Source(space=point)
settings = openmc.Settings()
settings.source = src
settings.batches = 500
settings.inactive = 20
settings.particles = 1000
settings.export_to_xml()