#!/usr/bin/env python3

"""
      ====================================================
      Description: AP1000 PWR Assembly 17*17
      Case:        Problem 2 (17*17 PWR Assembly)
      Written By:  Ahmed K Madani
      Conversion:  Adam Parler
      Date:        12/03/2023
      ====================================================
"""

import os
import openmc

try:
    os.mkdir(os.path.dirname(os.path.realpath(__file__)) + '/problem')
except OSError: pass

################## MATERIALS ##################

fuel_31 = openmc.Material()
fuel_31.set_density("g/cc", 10.5)
fuel_31.temperature = 600 # K
fuel_31.add_element("U", 1, enrichment=3.1)
fuel_31.add_nuclide("O16",  2, percent_type="ao")

gap = openmc.Material()
gap.set_density("atom/b-cm", 1.00000e-4)
gap.temperature = 600 # K
gap.add_nuclide("He4", 1, percent_type="ao")

clad = openmc.Material()
clad.set_density("g/cc", 6.55)
clad.temperature = 600 # K
clad.add_nuclide("O16", 0.00125, percent_type="wo")
clad.add_element("Nb", 0.01, percent_type="wo")
clad.add_element("Sn", 0.01, percent_type="wo")
clad.add_element("Fe", 0.001, percent_type="wo")
clad.add_element("Zr", 0.97775, percent_type="wo")

water = openmc.Material()
water.set_density("g/cc", 0.73)
water.temperature = 600 # K
water.add_nuclide("H1", 2, percent_type="ao")
water.add_nuclide("O16", 1, percent_type="ao")
water.add_s_alpha_beta(name='c_H_in_H2O')

tube = openmc.Material()
tube.set_density("g/cc", 7.92)
tube.temperature = 600 # K
tube.add_element("Ni", 0.0950095, percent_type="wo")
tube.add_element("Fe", 0.69493745, percent_type="wo")
tube.add_element("Mn", 0.02, percent_type="wo")
tube.add_element("Cr", 0.19, percent_type="wo")

materials_file = openmc.Materials([fuel_31, gap, clad, water, tube])
materials_file.export_to_xml(path="problem/materials.xml")

################## GEOMETRY ##################

surf1 = openmc.XPlane(x0=-10.659, boundary_type="reflective")
surf2 = openmc.XPlane(x0=10.659, boundary_type="reflective")
surf3 = openmc.YPlane(y0=-10.659, boundary_type="reflective")
surf4 = openmc.YPlane(y0=10.659, boundary_type="reflective")

surf5 = openmc.ZCylinder(x0=0, y0=0, r=0.4095)
surf6 = openmc.ZCylinder(x0=0, y0=0, r=0.4177)
surf7 = openmc.ZCylinder(x0=0, y0=0, r=0.4749)
surf8 = openmc.ZCylinder(x0=0, y0=0, r=0.5715)
surf9 = openmc.ZCylinder(x0=0, y0=0, r=0.612)

surf10 = openmc.ZPlane(z0=-182.9, boundary_type="reflective")
surf11 = openmc.ZPlane(z0=182.9, boundary_type="reflective")

univ1 = openmc.Universe()
cell10 = openmc.Cell(cell_id=10, fill=fuel_31)
cell10.region = -surf5

cell20 = openmc.Cell(cell_id=20, fill=gap)
cell20.region = +surf5 & -surf6

cell30 = openmc.Cell(cell_id=30, fill=clad)
cell30.region = +surf6 & -surf7

univ2 = openmc.Universe()
cell40 = openmc.Cell(cell_id=40, fill=water)
cell40.region = +surf7

cell50 = openmc.Cell(cell_id=50, fill=water)
cell50.region = -surf8

cell60 = openmc.Cell(cell_id=60, fill=tube)
cell60.region = +surf8 & -surf9

cell70 = openmc.Cell(cell_id=70, fill=water)
cell70.region = +surf9

univ1.add_cells([cell10, cell20, cell30, cell40])
univ2.add_cells([cell50, cell60, cell70])

lattice = openmc.RectLattice(lattice_id=5)
lattice.pitch = [1.26, 1.26]
lattice.lower_left = [-10.659, -10.659]
lattice.universes = [[univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ2, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ2, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ2, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ2, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ2, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1],
                      [univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1, univ1]
       ]
cell1 = openmc.Cell(fill=lattice)
cell1.region = +surf1 & -surf2 & +surf3 & -surf4 & +surf10 & -surf11

univ3 = openmc.Universe()
univ3.add_cell(cell1)

geometry_file = openmc.Geometry(univ3)
geometry_file.export_to_xml(path="problem/geometry.xml")

################## TALLIES ##################

filter1 = openmc.CellFilter(filter_id=1, bins=1)
filter8 = openmc.DelayedGroupFilter(filter_id=8, bins=[1, 2, 3, 4, 5, 6])

tally1 = openmc.Tally()
tally1.filters = [filter1]
tally1.scores = ["flux"]

tally2 = openmc.Tally()
tally2.filters = [filter1]
tally2.scores = ["nu-fission"]

tally3 = openmc.Tally()
tally3.filters = [filter1]
tally3.scores = ["delayed-nu-fission"]

tally4 = openmc.Tally()
tally4.filters = [filter1]
tally4.scores = ["fission"]

tally5 = openmc.Tally()
tally5.filters = [filter1]
tally5.scores = ["absorption"]

tally6 = openmc.Tally()
tally6.filters = [filter1]
tally6.scores = ["scatter"]

tally7 = openmc.Tally()
tally7.filters = [filter1]
tally7.scores = ["total"]

tally8 = openmc.Tally()
tally8.filters = [filter8]
tally8.scores = ["delayed-nu-fission"]

tallies_file = openmc.Tallies([tally1, tally2, tally3, tally4, tally5, tally6, tally7, tally8])
tallies_file.export_to_xml(path="problem/tallies.xml")

################## PLOTS ##################

plot = openmc.Plot()
plot.pixels = (1000, 1000)
plot.width = (25.50, 25.50)
# plot._colors = (123, 123, 231)

plot_file = openmc.Plots([plot])
plot_file.export_to_xml(path="problem/plots.xml")

################## SETTINGS ##################

settings_file = openmc.Settings()
settings_file.source = openmc.IndependentSource(space=openmc.stats.Box((-10.659, -10.659, -182.9),
                                                            (10.659,  10.659,  182.9)))
settings_file.batches = 600
settings_file.inactive = 50
settings_file.particles = 10000

settings_file.export_to_xml(path="problem/settings.xml")

openmc.run(cwd='problem')