import openmc
from openmc.data import atomic_mass, atomic_weight

MB10 = atomic_mass('B10')
MB11 = atomic_mass('B11')
MU234 = atomic_mass('U234')
MU235 = atomic_mass('U235')
MU238 = atomic_mass('U238')

materials = openmc.Materials()
air = openmc.Material(name='Air')
air.temperature = 300
air.set_density('g/cc', 0.000616)
air.add_element('O', 0.2095, 'ao')
air.add_element('N', 0.7809, 'ao')
air.add_element('Ar', 0.00933, 'ao')
air.add_element('C', 0.00027, 'ao')

## Fuels
'''
fuel_16 = openmc.Material(name='UO2 (1.6%)')
fuel_16.set_density('g/cm3', 10.31341)
fuel_16.temperature = 300
fuel_16.add_nuclide('U234', 3.0131e-06)
fuel_16.add_nuclide('U235', 3.7503e-04)
fuel_16.add_nuclide('U238', 2.2625e-02)
fuel_16.add_nuclide('O16', 4.5897e-02)
fuel_16.add_nuclide('O17', 1.7436e-05)
fuel_16.add_nuclide('O18', 9.2032e-05)

fuel_24 = openmc.Material(name='UO2 (2.4%)')
fuel_24.set_density('g/cm3', 10.29748)
fuel_24.temperature = 300
fuel_24.add_nuclide('U234', 4.4842e-06)
fuel_24.add_nuclide('U235', 5.5814e-04)
fuel_24.add_nuclide('U238', 2.2407e-02)
fuel_24.add_nuclide('O16', 4.5830e-02)
fuel_24.add_nuclide('O17', 1.7411e-05)
fuel_24.add_nuclide('O18', 9.1898e-05)

fuel_31 = openmc.Material(name='UO2 (3.1%)')
fuel_31.set_density('g/cm3', 10.30166)
fuel_31.temperature = 300
fuel_31.add_nuclide('U234', 5.7987e-06)
fuel_31.add_nuclide('U235', 7.2175e-04)
fuel_31.add_nuclide('U238', 2.2253e-02)
fuel_31.add_nuclide('O16', 4.5853e-02)
fuel_31.add_nuclide('O17', 1.7420e-05)
fuel_31.add_nuclide('O18', 9.1942e-05)

fuel_32 = openmc.Material(name='UO2 (3.2%)')
fuel_32.set_density('g/cm3', 10.34115)
fuel_32.temperature = 300
fuel_32.add_nuclide('U234', 5.9959e-06)
fuel_32.add_nuclide('U235', 7.4630e-04)
fuel_32.add_nuclide('U238', 2.2317e-02)
fuel_32.add_nuclide('O16', 4.6029e-02)
fuel_32.add_nuclide('O17', 1.7487e-05)
fuel_32.add_nuclide('O18', 9.2296e-05)

fuel_34 = openmc.Material(name='UO2 (3.4%)')
fuel_34.set_density('g/cm3', 10.35917)
fuel_34.temperature = 300
fuel_34.add_nuclide('U234', 6.4018e-06)
fuel_34.add_nuclide('U235', 7.9681e-04)
fuel_34.add_nuclide('U238', 2.2307e-02)
fuel_34.add_nuclide('O16', 4.6110e-02)
fuel_34.add_nuclide('O17', 1.7517e-05)
fuel_34.add_nuclide('O18', 9.2459e-05)
'''

## Specify enrichments to be calculated
enrs = [0.0161006, 0.0239993, 0.0310221, 0.0319547, 0.0340585]
dens = [10.31341, 10.29748, 10.30166, 10.34115, 10.35917]

# Loop around enrichments
for enr, den in zip(enrs, dens):

    # Calculate molar mass of Uranium
    enr_25 = enr
    enr_24 = 0.008*enr_25
    enr_28 = 1.0 - (enr_24 + enr_25)
    MU = 1.0/(enr_24/MU234 + enr_25/MU235 + enr_28/MU238)

    # Determine molar mass of UO2
    MUO2 = MU + 2.0*atomic_weight('O')

    # Compute weight percent of U in UO2
    wUpUO2 = MU/MUO2

    # Calculate Uranium atom fraction
    a_U = wUpUO2*MUO2/MU

    # Calculate Oxygen atom fraction
    a_O = (1.0 - wUpUO2)*MUO2/atomic_weight('O')

    # Create material
    name = enr*100
    mat_name = 'Fuel {0:1.1f}%'.format(name)
    mat = openmc.Material(name=mat_name)
    mat.temperature = 300
    mat.set_density('g/cc', den)
    mat.add_element('O', a_O, 'ao')
    mat.add_element('U', a_U, 'ao', enrichment=enr*100)
    materials += mat


borosilicate_glass = openmc.Material(name='Borosilicate Glass')
borosilicate_glass.set_density('g/cm3', 2.26)
borosilicate_glass.temperature = 300
borosilicate_glass.add_nuclide('Al27', 1.7352e-03)
borosilicate_glass.add_nuclide('B10', 9.6506e-04)
borosilicate_glass.add_nuclide('B11', 3.9189e-03)
borosilicate_glass.add_nuclide('O16', 4.6514e-02)
borosilicate_glass.add_nuclide('O17', 1.7671e-05)
borosilicate_glass.add_nuclide('O18', 9.3268e-05)
borosilicate_glass.add_nuclide('Si28', 1.6926e-02)
borosilicate_glass.add_nuclide('Si29', 8.5944e-04)
borosilicate_glass.add_nuclide('Si30', 5.6654e-04)

ag_in_cd_rods = openmc.Material(name='Ag-In-Cd Control Rod')
ag_in_cd_rods.set_density('g/cm3', 10.16)
ag_in_cd_rods.temperature = 300
ag_in_cd_rods.add_nuclide('Ag107', 2.3523e-02)
ag_in_cd_rods.add_nuclide('Ag109', 2.1854e-02)
ag_in_cd_rods.add_nuclide('Cd106', 3.3882e-05)
ag_in_cd_rods.add_nuclide('Cd108', 2.4166e-05)
ag_in_cd_rods.add_nuclide('Cd110', 3.3936e-04)
ag_in_cd_rods.add_nuclide('Cd111', 3.4821e-04)
ag_in_cd_rods.add_nuclide('Cd112', 6.5611e-04)
ag_in_cd_rods.add_nuclide('Cd113', 3.3275e-04)
ag_in_cd_rods.add_nuclide('Cd114', 7.8252e-04)
ag_in_cd_rods.add_nuclide('Cd116', 2.0443e-04)
ag_in_cd_rods.add_nuclide('In113', 3.4219e-04)
ag_in_cd_rods.add_nuclide('In115', 7.6511e-03)

b4c = openmc.Material(name='B4C Control Rod')
b4c.set_density('g/cm3', 1.76)
b4c.temperature = 300
b4c.add_nuclide('B10', 1.5206e-02)
b4c.add_nuclide('B11', 6.1514e-02)
b4c.add_nuclide('C12', 1.8972e-02)
b4c.add_nuclide('C13', 2.1252e-04)

helium = openmc.Material(name='Helium')
helium.set_density('g/cm3', 0.0015981)
helium.temperature = 300
helium.add_nuclide('He3', 4.8089e-10)
helium.add_nuclide('He4', 2.4044e-04)

inconel_718 = openmc.Material(name='Inconel 718')
inconel_718.set_density('g/cm3', 8.2)
inconel_718.temperature = 300
inconel_718.add_nuclide('Cr50', 7.8239e-04)
inconel_718.add_nuclide('Cr52', 1.5088e-02)
inconel_718.add_nuclide('Cr53', 1.7108e-03)
inconel_718.add_nuclide('Cr54', 4.2586e-04)
inconel_718.add_nuclide('Fe54', 1.4797e-03)
inconel_718.add_nuclide('Fe56', 2.3229e-02)
inconel_718.add_nuclide('Fe57', 5.3645e-04)
inconel_718.add_nuclide('Fe58', 7.1392e-05)
inconel_718.add_nuclide('Mn55', 7.8201e-04)
inconel_718.add_nuclide('Ni58', 2.9320e-02)
inconel_718.add_nuclide('Ni60', 1.1294e-02)
inconel_718.add_nuclide('Ni61', 4.9094e-04)
inconel_718.add_nuclide('Ni62', 1.5653e-03)
inconel_718.add_nuclide('Ni64', 3.9864e-04)
inconel_718.add_nuclide('Si28', 5.6757e-04)
inconel_718.add_nuclide('Si29', 2.8820e-05)
inconel_718.add_nuclide('Si30', 1.8998e-05)

ss304 = openmc.Material(name='Stainless Steel 304')
ss304.set_density('g/cm3', 8.03)
ss304.temperature = 300
ss304.add_nuclide('Cr50', 7.6778e-04)
ss304.add_nuclide('Cr52', 1.4806e-02)
ss304.add_nuclide('Cr53', 1.6789e-03)
ss304.add_nuclide('Cr54', 4.1791e-04)
ss304.add_nuclide('Fe54', 3.4620e-03)
ss304.add_nuclide('Fe56', 5.4345e-02)
ss304.add_nuclide('Fe57', 1.2551e-03)
ss304.add_nuclide('Fe58', 1.6703e-04)
ss304.add_nuclide('Mn55', 1.7604e-03)
ss304.add_nuclide('Ni58', 5.6089e-03)
ss304.add_nuclide('Ni60', 2.1605e-03)
ss304.add_nuclide('Ni61', 9.3917e-05)
ss304.add_nuclide('Ni62', 2.9945e-04)
ss304.add_nuclide('Ni64', 7.6261e-05)
ss304.add_nuclide('Si28', 9.5281e-04)
ss304.add_nuclide('Si29', 4.8381e-05)
ss304.add_nuclide('Si30', 3.1893e-05)

zirc_4 = openmc.Material()
zirc_4.set_density('g/cm3', 6.55)
zirc_4.temperature = 300
zirc_4.add_nuclide('Cr50', 3.2962e-06)
zirc_4.add_nuclide('Cr52', 6.3564e-05)
zirc_4.add_nuclide('Cr53', 7.2076e-06)
zirc_4.add_nuclide('Cr54', 1.7941e-06)
zirc_4.add_nuclide('Fe54', 8.6698e-06)
zirc_4.add_nuclide('Fe56', 1.3610e-04)
zirc_4.add_nuclide('Fe57', 3.1431e-06)
zirc_4.add_nuclide('Fe58', 4.1829e-07)
zirc_4.add_nuclide('O16', 3.0744e-04)
zirc_4.add_nuclide('O17', 1.1680e-07)
zirc_4.add_nuclide('O18', 6.1648e-07)
zirc_4.add_nuclide('Sn112', 4.6735e-06)
zirc_4.add_nuclide('Sn114', 3.1799e-06)
zirc_4.add_nuclide('Sn115', 1.6381e-06)
zirc_4.add_nuclide('Sn116', 7.0055e-05)
zirc_4.add_nuclide('Sn117', 3.7003e-05)
zirc_4.add_nuclide('Sn118', 1.1669e-04)
zirc_4.add_nuclide('Sn119', 4.1387e-05)
zirc_4.add_nuclide('Sn120', 1.5697e-04)
zirc_4.add_nuclide('Sn122', 2.2308e-05)
zirc_4.add_nuclide('Sn124', 2.7897e-05)
zirc_4.add_nuclide('Zr90', 2.1828e-02)
zirc_4.add_nuclide('Zr91', 4.7601e-03)
zirc_4.add_nuclide('Zr92', 7.2759e-03)
zirc_4.add_nuclide('Zr94', 7.3734e-03)
zirc_4.add_nuclide('Zr96', 1.1879e-03)

borated_water = openmc.Material()
borated_water.set_density('g/cm3', 0.740582068)
borated_water.temperature = 300
borated_water.add_nuclide('B10', 7.9714e-06)
borated_water.add_nuclide('B11', 3.2247e-05)
borated_water.add_nuclide('H1', 4.9456e-02)
borated_water.add_nuclide('H2', 7.7035e-06)
borated_water.add_nuclide('O16', 2.4673e-02)
borated_water.add_nuclide('O17', 9.3734e-06)
borated_water.add_nuclide('O18', 4.9474e-05)

plate_borated_water = openmc.Material()
plate_borated_water.set_density('g/cm3', 0.981002532)
plate_borated_water.add_nuclide('B10', 1.0559e-05)
plate_borated_water.add_nuclide('B11', 4.2716e-05)
plate_borated_water.add_nuclide('H1', 6.5512e-02)
plate_borated_water.add_nuclide('H2', 1.0204e-05)
plate_borated_water.add_nuclide('O16', 3.2683e-02)
plate_borated_water.add_nuclide('O17', 1.2416e-05)
plate_borated_water.add_nuclide('O18', 6.5535e-05)

plate_stainless = openmc.Material()
plate_stainless.set_density('g/cm3', 3.68384807)
plate_stainless.temperature = 300
plate_stainless.add_nuclide('Cr50', 3.5223e-04)
plate_stainless.add_nuclide('Cr52', 6.7924e-03)
plate_stainless.add_nuclide('Cr53', 7.7020e-04)
plate_stainless.add_nuclide('Cr54', 1.9172e-04)
plate_stainless.add_nuclide('Fe54', 1.5882e-03)
plate_stainless.add_nuclide('Fe56', 2.4931e-02)
plate_stainless.add_nuclide('Fe57', 5.7578e-04)
plate_stainless.add_nuclide('Fe58', 7.6625e-05)
plate_stainless.add_nuclide('Mn55', 8.0762e-04)
plate_stainless.add_nuclide('Ni58', 2.5731e-03)
plate_stainless.add_nuclide('Ni60', 9.9117e-04)
plate_stainless.add_nuclide('Ni61', 4.3085e-05)
plate_stainless.add_nuclide('Ni62', 1.3738e-04)
plate_stainless.add_nuclide('Ni64', 3.4985e-05)
plate_stainless.add_nuclide('Si28', 4.3711e-04)
plate_stainless.add_nuclide('Si29', 2.2195e-05)
plate_stainless.add_nuclide('Si30', 1.4631e-05)

carbon_steel = openmc.Material()
carbon_steel.set_density('g/cm3', 7.8)
carbon_steel.temperature = 300
carbon_steel.add_nuclide('Al27', 4.3523e-05)
carbon_steel.add_nuclide('B10', 2.5833e-06)
carbon_steel.add_nuclide('B11', 1.0450e-05)
carbon_steel.add_nuclide('C12', 1.0442e-03)
carbon_steel.add_nuclide('C13', 1.1697e-05)
carbon_steel.add_nuclide('Ca40', 1.7043e-05)
carbon_steel.add_nuclide('Ca42', 1.1375e-07)
carbon_steel.add_nuclide('Ca43', 2.3734e-08)
carbon_steel.add_nuclide('Ca44', 3.6673e-07)
carbon_steel.add_nuclide('Ca46', 7.0322e-10)
carbon_steel.add_nuclide('Ca48', 3.2875e-08)
carbon_steel.add_nuclide('Cr50', 1.3738e-05)
carbon_steel.add_nuclide('Cr52', 2.6493e-04)
carbon_steel.add_nuclide('Cr53', 3.0041e-05)
carbon_steel.add_nuclide('Cr54', 7.4778e-06)
carbon_steel.add_nuclide('Cu63', 1.0223e-04)
carbon_steel.add_nuclide('Cu65', 4.5608e-05)
carbon_steel.add_nuclide('Fe54', 4.7437e-03)
carbon_steel.add_nuclide('Fe56', 7.4465e-02)
carbon_steel.add_nuclide('Fe57', 1.7197e-03)
carbon_steel.add_nuclide('Fe58', 2.2886e-04)
carbon_steel.add_nuclide('Mn55', 6.4126e-04)
carbon_steel.add_nuclide('Mo100', 2.9814e-05)
carbon_steel.add_nuclide('Mo92', 4.4822e-05)
carbon_steel.add_nuclide('Mo94', 2.8110e-05)
carbon_steel.add_nuclide('Mo95', 4.8567e-05)
carbon_steel.add_nuclide('Mo96', 5.1015e-05)
carbon_steel.add_nuclide('Mo97', 2.9319e-05)
carbon_steel.add_nuclide('Mo98', 7.4327e-05)
carbon_steel.add_nuclide('Nb93', 5.0559e-06)
carbon_steel.add_nuclide('Ni58', 4.0862e-04)
carbon_steel.add_nuclide('Ni60', 1.5740e-04)
carbon_steel.add_nuclide('Ni61', 6.8420e-06)
carbon_steel.add_nuclide('Ni62', 2.1815e-05)
carbon_steel.add_nuclide('Ni64', 5.5557e-06)
carbon_steel.add_nuclide('P31', 3.7913e-05)
carbon_steel.add_nuclide('S32', 3.4808e-05)
carbon_steel.add_nuclide('S33', 2.7420e-07)
carbon_steel.add_nuclide('S34', 1.5368e-06)
carbon_steel.add_nuclide('S36', 5.3398e-09)
carbon_steel.add_nuclide('Si28', 6.1702e-04)
carbon_steel.add_nuclide('Si29', 3.1330e-05)
carbon_steel.add_nuclide('Si30', 2.0653e-05)
carbon_steel.add_nuclide('Ti46', 1.2144e-06)
carbon_steel.add_nuclide('Ti47', 1.0952e-06)
carbon_steel.add_nuclide('Ti48', 1.0851e-05)
carbon_steel.add_nuclide('Ti49', 7.9634e-07)
carbon_steel.add_nuclide('Ti50', 7.6249e-07)
carbon_steel.add_nuclide('V50', 1.1526e-07)
carbon_steel.add_nuclide('V51', 4.5989e-05)

"""
materials += [fuel_16, fuel_24, fuel_31, fuel_32, fuel_34]
"""
materials += [air, borosilicate_glass, ag_in_cd_rods, b4c, helium,
              inconel_718, ss304, zirc_4, borated_water, plate_borated_water,
              plate_stainless, carbon_steel]
materials.cross_sections="/opt/xdata/cross_sections.xml"
materials.export_to_xml()

## Fuel Pin
fuel_or = openmc.ZCylinder(x0=0, y0=0, r=0.39218, name='Fuel OR')
clad_ir = openmc.ZCylinder(x0=0, y0=0, r=0.40005, name='Clad IR')
clad_or = openmc.ZCylinder(x0=0, y0=0, r=0.45720, name='Clad OR')

## Upper Fuel Pin Plenum
spring = openmc.ZCylinder(x0=0, y0=0, r=0.06459, name='Spring')
clad_ir = openmc.ZCylinder(x0=0, y0=0, r=0.40005, name='Clad IR')
clad_or = openmc.ZCylinder(x0=0, y0=0, r=0.45720, name='Clad OR')

## Empty Guide Tube above Dashpot
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_or = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide OR')

## Empty Guide Tube at Dashpot
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.50419, name='Guide IR')
guide_or = openmc.ZCylinder(x0=0, y0=0, r=0.54610, name='Guide OR')

## Instrument Tube Pin Geometry
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.43688, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide IR')

## Bare Instrument Thimble Pin
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.43688, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')

## BP above Dashpoint
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.21400, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.23051, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.24130, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.42672, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.43688, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide IR')

## BP Plennum
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.21400, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.23051, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.43688, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.50419, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.54610, name='Guide IR')

## Control Rod Pin Upper Geometry
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.37338, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.38608, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide IR')

## Control Rod Pin Lower Geometry
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.38227, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.38608, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide IR')

## Control Rod Pin Spacer
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.37845, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.38608, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide IR')

## Control Rod Plenum Geometry
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.06459, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.38608, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.48387, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.56134, name='Guide IR')
guide_ir = openmc.ZCylinder(x0=0, y0=0, r=0.60198, name='Guide IR')

