#!/bin/env python3.12
import openmc
import math
import argparse
import hashlib
import sys
import time
import string
import numpy as np
import warnings

# HP-MR Model in OpenMC
# C.Moore, <connor.moore@psi.ch>

############ HLPER FUNCTIONS ############ 

def rot(base: openmc.Universe, angle: float = 0.0) -> openmc.Universe:
    # Take the cells from the existing universe and create a new rotated universe
    new_cell = openmc.Cell(fill=base)
    new_cell.rotation=(0.0,0.0,angle)
    new_cell.name = f"Control Drum Cell, Rotated {angle} deg"
    return openmc.Universe(name=base.name + f" Rotated {angle} deg", cells=[new_cell])

############ COMMAND LINE ARGUMENTS ############ 
cli_parser = argparse.ArgumentParser(
        prog="HP-RM TRISO Homogenization Study",
        description="Program to investigate various homogenization techniques for TRISO compacts in the HP-MR benchmark."
        )

cli_parser.add_argument("-t","--technique",default="none",choices=["none","vwh","rpt","rrpt"],help="What homogenization technique to use (if any) [none/vwh/rpt/rrpt], default none")
cli_parser.add_argument("-r","--radius",default="0.8",type=float,help="Radius for RPT or RRPT homogenization [cm], default 0.8 cm. In the case of RRPT, this is the inner radius.")
cli_parser.add_argument("-z","--z-height",default="1",type=float,help="Depth of pin/assembly slice in the z-direction [cm]. Should divide nicely into 80, default 1.0")
cli_parser.add_argument("-l","--lclrs",default=False,action="store_true",help="Flag for specifying LCLRS accessible cross-sections, default false.")
cli_parser.add_argument("-g","--geometry-level",default="pin",choices=["pin","ass","core"],help="Geometry level to run the computation at [pin/ass/core], default pin.")
cli_parser.add_argument("-d","--drums-in",default=False,action="store_true",help="Flag for specifying the control drums are rotated in, default false")
cli_parser.add_argument("-p","--packing",default=False,action="store_true",help="Flag for specifying if packing information should be read from file")
cli_parser.add_argument("-e","--execute",default=False,action="store_true",help="Flag for automatically running OpenMC, default false.")
cli_parser.add_argument("--triso-lattice",default=False,action="store_true",help="Flag for enabling TRISO lattice domain decomposition in the compact, default false.")
cli_parser.add_argument("--model-hp",default=False,action="store_true",help="Flag for enabling heat pipe modelling of pool/saturation heights, default false.")
cli_parser.add_argument("--pool-height",default=15.0,type=float,help="Heap pipe pooling height [cm]. Requires z_height=80.0, default 15.0, range (0,160)")
cli_parser.add_argument("--sat-height",default=80.0,type=float,help="Heap pipe wick saturation height [cm]. Requires z_height=80.0, default 80.0, range (0,160)")
cli_args = cli_parser.parse_args()

if cli_args.lclrs:
    openmc.config["cross_sections"] = "/libs/endfb81_official/cross_sections.xml"

# Sanity check
if cli_args.model_hp and cli_args.z_height!=80:
    warnings.warn("Warning: ignoring --model-hp flag as the slice height is not full-core.", RuntimeWarning)


############ INTPUT-OUTPUT MAPPING ############ 
with open(sys.argv[0], "rb") as file:
    contents = file.read()

filehash = hashlib.md5(contents).hexdigest()
inputs=str(cli_args)[10:-1].replace(',' , '\n'+17*' ').replace('=',' = ')

print("="*52)
print(f"Run date CET:     {time.ctime()}")
print(f"Input file MD5:   {filehash}")
print(f"Input arguments:  {inputs}")
print("="*52)

############ MATERIALS ############ 
# All materials are specified in the report on HEAT PIPE MICROREACTOR MODELING WITH BLUECRAB
# <https://www.tandfonline.com/doi/full/10.1080/00295639.2024.2375175>
# Materials marked PNNL-15870 are from the 2nd revision of the report "Compendium of Material
# Composition Data for Radiation Transport Modeling " by US Dept. of Homeland Security and PNNL
# <https://www.pnnl.gov/main/publications/external/technical_reports/PNNL-15870Rev2.pdf>
mat_list = []
mat_temp_global = 700.0
mat_temp_fuel = 1000.0
mat_pipe_porosity = 0.70

## 100s - Matrix Graphite ##
mat_graphite_matrix = openmc.Material(100,name="Graphite Matrix (between pins)")
mat_graphite_matrix.add_nuclide("C12",0.9999997)
mat_graphite_matrix.add_nuclide("B10",3e-7)
mat_graphite_matrix.set_density("g/cc",1.806)
mat_graphite_matrix.temperature = mat_temp_global
mat_list.append(mat_graphite_matrix)

## 200s - TRISO Pins ##
mat_triso_uco = openmc.Material(200,name="TRISO Pin UCO Kernel")
mat_triso_uco.add_nuclide("U235",0.068794)
mat_triso_uco.add_nuclide("U238",0.27604)
mat_triso_uco.add_nuclide("C12",0.13793)
mat_triso_uco.add_nuclide("O16",0.51724)
mat_triso_uco.set_density("g/cc",10.744)
mat_triso_uco.temperature = mat_temp_fuel
mat_list.append(mat_triso_uco)

mat_triso_buffer = openmc.Material(201,name="TRISO Pin Buffer")
mat_triso_buffer.add_nuclide("C12",1.0)
mat_triso_buffer.set_density("g/cc",1.04)
mat_triso_buffer.temperature = mat_temp_fuel
mat_triso_buffer.add_s_alpha_beta("c_Graphite")
mat_list.append(mat_triso_buffer)

mat_triso_pyc1 = openmc.Material(202,name="TRISO Pin PyC1")
mat_triso_pyc1.add_nuclide("C12",1.0)
mat_triso_pyc1.set_density("g/cc",1.882)
mat_triso_pyc1.temperature = mat_temp_fuel
mat_triso_pyc1.add_s_alpha_beta("c_Graphite")
mat_list.append(mat_triso_pyc1)

mat_triso_sic = openmc.Material(203,name="TRISO Pin SiC")
mat_triso_sic.add_nuclide("Si28",0.4611)
mat_triso_sic.add_nuclide("Si29",0.0234)
mat_triso_sic.add_nuclide("Si30",0.0154)
mat_triso_sic.add_nuclide("C12",0.5)
mat_triso_sic.set_density("g/cc",3.171)
mat_triso_sic.temperature = mat_temp_fuel
mat_triso_sic.add_s_alpha_beta("c_C_in_SiC")
mat_triso_sic.add_s_alpha_beta("c_Si_in_SiC")
mat_list.append(mat_triso_sic)

mat_triso_pyc2 = openmc.Material(204,name="TRISO Pin PyC2")
mat_triso_pyc2.add_nuclide("C12",1.0)
mat_triso_pyc2.add_s_alpha_beta("c_Graphite")
mat_triso_pyc2.temperature = mat_temp_fuel
mat_triso_pyc2.set_density("g/cc",1.882)
mat_list.append(mat_triso_pyc2)

## 300 - Moderator Pins ##
mat_mod_helium = openmc.Material(300,name="Moderator Pin Helium")
mat_mod_helium.add_nuclide("He4",1.0)
mat_mod_helium.set_density("g/cc",0.18e-3)
mat_list.append(mat_mod_helium)

# From the offical Serpent input
mat_mod_ss316 = openmc.Material(301,name="Moderator Pin SS316")
mat_mod_ss316.add_nuclide('C12',  0.001901)
mat_mod_ss316.add_nuclide('Si28', 0.0092693)
mat_mod_ss316.add_nuclide('Si29', 4.7251e-4)
mat_mod_ss316.add_nuclide('Si30', 3.1166e-4)
mat_mod_ss316.add_nuclide('P31',  4.1322e-4)
mat_mod_ss316.add_nuclide('S32',  2.471e-4)
mat_mod_ss316.add_nuclide('S33',  1.9511e-6)
mat_mod_ss316.add_nuclide('S34',  1.1056e-5)
mat_mod_ss316.add_nuclide('S36',  3.0016e-8)
mat_mod_ss316.add_nuclide('Cr50', 7.9116e-3)
mat_mod_ss316.add_nuclide('Cr52', 0.15257)
mat_mod_ss316.add_nuclide('Cr53', 0.01730)
mat_mod_ss316.add_nuclide('Cr54', 4.3063e-3)
mat_mod_ss316.add_nuclide('Mn55', 0.01028)
mat_mod_ss316.add_nuclide('Fe54', 0.039029)
mat_mod_ss316.add_nuclide('Fe56', 0.61213)
mat_mod_ss316.add_nuclide('Fe57', 0.014144)
mat_mod_ss316.add_nuclide('Fe58', 1.3343e-3)
mat_mod_ss316.add_nuclide('Ni58', 0.077516)
mat_mod_ss316.add_nuclide('Ni60', 0.029859)
mat_mod_ss316.add_nuclide('Ni61', 1.2981e-3)
mat_mod_ss316.add_nuclide('Ni62', 4.1389e-3)
mat_mod_ss316.add_nuclide('Ni64', 1.0544e-3)
mat_mod_ss316.add_nuclide('Mo92', 2.1259e-3)
mat_mod_ss316.add_nuclide('Mo94', 1.3299e-3)
mat_mod_ss316.add_nuclide('Mo95', 2.303e-3)
mat_mod_ss316.add_nuclide('Mo96', 2.4191e-3)
mat_mod_ss316.add_nuclide('Mo97', 1.3903e-3)
mat_mod_ss316.add_nuclide('Mo98', 3.5249e-3)
mat_mod_ss316.add_nuclide('Mo100',1.4135e-3)
mat_mod_ss316.set_density("g/cc",7.9)
mat_list.append(mat_mod_ss316)

# Official Serpent smeared composition
mat_mod_ss316_he_smear = openmc.Material(302,name="Moderator Pin SS316/He Smear")
mat_mod_ss316_he_smear.add_nuclide('He4',1.983e-5)
mat_mod_ss316_he_smear.add_nuclide('C12',4.349e-5)
mat_mod_ss316_he_smear.add_nuclide('Si28',2.121e-4)
mat_mod_ss316_he_smear.add_nuclide('Si29',1.081e-5)
mat_mod_ss316_he_smear.add_nuclide('Si30',7.13e-6)
mat_mod_ss316_he_smear.add_nuclide('P31',9.454e-6)
mat_mod_ss316_he_smear.add_nuclide('S32',5.653e-6)
mat_mod_ss316_he_smear.add_nuclide('S33',4.464e-8)
mat_mod_ss316_he_smear.add_nuclide('S34',2.53e-7)
mat_mod_ss316_he_smear.add_nuclide('S36',6.867e-10)
mat_mod_ss316_he_smear.add_nuclide('Cr50',1.81e-4)
mat_mod_ss316_he_smear.add_nuclide('Cr52',3.491e-3)
mat_mod_ss316_he_smear.add_nuclide('Cr53',3.958e-4)
mat_mod_ss316_he_smear.add_nuclide('Cr54',9.852e-5)
mat_mod_ss316_he_smear.add_nuclide('Mn55',2.352e-4)
mat_mod_ss316_he_smear.add_nuclide('Fe54',8.929e-4)
mat_mod_ss316_he_smear.add_nuclide('Fe56',0.014)
mat_mod_ss316_he_smear.add_nuclide('Fe57',3.236e-4)
mat_mod_ss316_he_smear.add_nuclide('Fe58',3.053e-5)
mat_mod_ss316_he_smear.add_nuclide('Ni58',1.773e-3)
mat_mod_ss316_he_smear.add_nuclide('Ni60',6.831e-4)
mat_mod_ss316_he_smear.add_nuclide('Ni61',2.97e-5)
mat_mod_ss316_he_smear.add_nuclide('Ni62',9.469e-5)
mat_mod_ss316_he_smear.add_nuclide('Ni64',2.412e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo92',4.864e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo94',3.043e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo95',5.269e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo96',5.535e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo97',3.181e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo98',8.065e-5)
mat_mod_ss316_he_smear.add_nuclide('Mo100',3.234e-5)
mat_mod_ss316_he_smear.set_density("atom/b-cm",0.0229)
mat_mod_ss316_he_smear.temperature = mat_temp_global
mat_list.append(mat_mod_ss316_he_smear)

# YH2 from official Serpent
mat_mod_yh2 = openmc.Material(303,name="Moderater Pin YH2")
mat_mod_yh2.add_nuclide("Y89",0.357142857)
mat_mod_yh2.add_nuclide("H1",0.642857143)
mat_mod_yh2.set_density("g/cc",4.085)
mat_mod_yh2.temperature = mat_temp_global
mat_mod_yh2.add_s_alpha_beta("c_H_in_YH2")
mat_mod_yh2.add_s_alpha_beta("c_Y_in_YH2")
mat_list.append(mat_mod_yh2)


## 400 - Heat Pipes ##
mat_pipe_helium = openmc.Material(400,name="Heat Pipe Helium")
mat_pipe_helium.add_nuclide("He4",1.0)
mat_pipe_helium.set_density("g/cc",0.18e-3)
mat_list.append(mat_pipe_helium)

# From the offical Serpent input
mat_pipe_ss316 = openmc.Material(401,name="Heat Pipe SS316")
mat_pipe_ss316.add_nuclide('C12', 0.001901)
mat_pipe_ss316.add_nuclide('Si28',0.0092693)
mat_pipe_ss316.add_nuclide('Si29',4.7251e-4)
mat_pipe_ss316.add_nuclide('Si30',3.1166e-4)
mat_pipe_ss316.add_nuclide('P31', 4.1322e-4)
mat_pipe_ss316.add_nuclide('S32', 2.471e-4)
mat_pipe_ss316.add_nuclide('S33', 1.9511e-6)
mat_pipe_ss316.add_nuclide('S34', 1.1056e-5)
mat_pipe_ss316.add_nuclide('S36', 3.0016e-8)
mat_pipe_ss316.add_nuclide('Cr50',7.9116e-3)
mat_pipe_ss316.add_nuclide('Cr52',0.15257)
mat_pipe_ss316.add_nuclide('Cr53',0.01730)
mat_pipe_ss316.add_nuclide('Cr54',4.3063e-3)
mat_pipe_ss316.add_nuclide('Mn55',0.01028)
mat_pipe_ss316.add_nuclide('Fe54',0.039029)
mat_pipe_ss316.add_nuclide('Fe56',0.61213)
mat_pipe_ss316.add_nuclide('Fe57',0.014144)
mat_pipe_ss316.add_nuclide('Fe58',1.3343e-3)
mat_pipe_ss316.add_nuclide('Ni58',0.077516)
mat_pipe_ss316.add_nuclide('Ni60',0.029859)
mat_pipe_ss316.add_nuclide('Ni61',1.2981e-3)
mat_pipe_ss316.add_nuclide('Ni62',4.1389e-3)
mat_pipe_ss316.add_nuclide('Ni64',1.0544e-3)
mat_pipe_ss316.add_nuclide('Mo92',2.1259e-3)
mat_pipe_ss316.add_nuclide('Mo94',1.3299e-3)
mat_pipe_ss316.add_nuclide('Mo95',2.303e-3)
mat_pipe_ss316.add_nuclide('Mo96',2.4191e-3)
mat_pipe_ss316.add_nuclide('Mo97',1.3903e-3)
mat_pipe_ss316.add_nuclide('Mo98',3.5249e-3)
mat_pipe_ss316.add_nuclide('Mo100',1.4135e-3)
mat_pipe_ss316.set_density("g/cc",7.9)
mat_list.append(mat_pipe_ss316)

mat_pipe_k_liquid = openmc.Material(402,name="Heat Pipe Potassium (Liquid)")
mat_pipe_k_liquid.add_nuclide("K39",0.93258)
mat_pipe_k_liquid.add_nuclide("K40",0.00012)
mat_pipe_k_liquid.add_nuclide("K41",0.06730)
mat_pipe_k_liquid.set_density("g/cc",0.705)
mat_pipe_k_liquid.temperature = mat_temp_global
mat_list.append(mat_pipe_k_liquid)

mat_pipe_k_gas = openmc.Material(403,name="Heat Pipe Potassium (Gaseous)")
mat_pipe_k_gas.add_nuclide("K39",0.93258)
mat_pipe_k_gas.add_nuclide("K40",0.00012)
mat_pipe_k_gas.add_nuclide("K41",0.06730)
mat_pipe_k_gas.temperature = mat_temp_global
mat_pipe_k_gas.set_density("g/cc",1.11e-4)
mat_list.append(mat_pipe_k_gas)

mat_pipe_wick_smear = openmc.Material(404,name="Heat Pipe Wick (Smeared)")
mat_pipe_wick_smear.add_nuclide('C12',5.589e-4)
mat_pipe_wick_smear.add_nuclide('Si28',2.725e-3)
mat_pipe_wick_smear.add_nuclide('Si29',1.389e-4)
mat_pipe_wick_smear.add_nuclide('Si30',9.163e-5)
mat_pipe_wick_smear.add_nuclide('P31',1.215e-4)
mat_pipe_wick_smear.add_nuclide('S32',7.265e-5)
mat_pipe_wick_smear.add_nuclide('S33',5.736e-7)
mat_pipe_wick_smear.add_nuclide('S34',3.25e-6)
mat_pipe_wick_smear.add_nuclide('S36',8.825e-9)
mat_pipe_wick_smear.add_nuclide('Cr50',2.326e-3)
mat_pipe_wick_smear.add_nuclide('Cr52',4.485e-2)
mat_pipe_wick_smear.add_nuclide('Cr53',5.086e-3)
mat_pipe_wick_smear.add_nuclide('Cr54',1.266e-3)
mat_pipe_wick_smear.add_nuclide('Mn55',3.022e-3)
mat_pipe_wick_smear.add_nuclide('Fe54',1.147e-2)
mat_pipe_wick_smear.add_nuclide('Fe56',0.18)
mat_pipe_wick_smear.add_nuclide('Fe57',4.158e-3)
mat_pipe_wick_smear.add_nuclide('Fe58',3.923e-4)
mat_pipe_wick_smear.add_nuclide('Ni58',2.279e-2)
mat_pipe_wick_smear.add_nuclide('Ni60',8.778e-3)
mat_pipe_wick_smear.add_nuclide('Ni61',3.816e-4)
mat_pipe_wick_smear.add_nuclide('Ni62',1.217e-3)
mat_pipe_wick_smear.add_nuclide('Ni64',3.10e-4)
mat_pipe_wick_smear.add_nuclide('Mo92',6.25e-4)
mat_pipe_wick_smear.add_nuclide('Mo94',3.91e-4)
mat_pipe_wick_smear.add_nuclide('Mo95',6.771e-4)
mat_pipe_wick_smear.add_nuclide('Mo96',7.112e-4)
mat_pipe_wick_smear.add_nuclide('Mo97',4.087e-4)
mat_pipe_wick_smear.add_nuclide('Mo98',1.036e-3)
mat_pipe_wick_smear.add_nuclide('Mo100',4.156e-4)
mat_pipe_wick_smear.add_nuclide('K39',0.6584)
mat_pipe_wick_smear.add_nuclide('K40',8.472e-5)
mat_pipe_wick_smear.add_nuclide('K41',4.751e-2)
mat_pipe_wick_smear.temperature = mat_temp_global
mat_pipe_wick_smear.set_density("g/cc",2.753)
mat_list.append(mat_pipe_wick_smear)

mat_pipe_inner_he = openmc.Material(405,name="Heat Pipe Inner He Smear")
mat_pipe_inner_he.add_nuclide('He4',2.68e-5)
mat_pipe_inner_he.add_nuclide('C12',1.62e-6)
mat_pipe_inner_he.add_nuclide('Si28',7.92e-6)
mat_pipe_inner_he.add_nuclide('Si29',4.04e-7)
mat_pipe_inner_he.add_nuclide('Si30',2.66e-7)
mat_pipe_inner_he.add_nuclide('P31',3.53e-7)
mat_pipe_inner_he.add_nuclide('S32',2.11e-7)
mat_pipe_inner_he.add_nuclide('S33',1.67e-9)
mat_pipe_inner_he.add_nuclide('S34',9.45e-9)
mat_pipe_inner_he.add_nuclide('S36',2.57e-11)
mat_pipe_inner_he.add_nuclide('Cr50',6.76e-6)
mat_pipe_inner_he.add_nuclide('Cr52',1.30e-4)
mat_pipe_inner_he.add_nuclide('Cr53',1.48e-5)
mat_pipe_inner_he.add_nuclide('Cr54',3.68e-6)
mat_pipe_inner_he.add_nuclide('Mn55',8.79e-6)
mat_pipe_inner_he.add_nuclide('Fe54',3.34e-5)
mat_pipe_inner_he.add_nuclide('Fe56',5.23e-4)
mat_pipe_inner_he.add_nuclide('Fe57',1.21e-5)
mat_pipe_inner_he.add_nuclide('Fe58',1.14e-6)
mat_pipe_inner_he.add_nuclide('Ni58',6.63e-5)
mat_pipe_inner_he.add_nuclide('Ni60',2.55e-5)
mat_pipe_inner_he.add_nuclide('Ni61',1.11e-6)
mat_pipe_inner_he.add_nuclide('Ni62',3.54e-6)
mat_pipe_inner_he.add_nuclide('Ni64',9.01e-7)
mat_pipe_inner_he.add_nuclide('Mo92',1.82e-6)
mat_pipe_inner_he.add_nuclide('Mo94',1.14e-6)
mat_pipe_inner_he.add_nuclide('Mo95',1.97e-6)
mat_pipe_inner_he.add_nuclide('Mo96',2.07e-6)
mat_pipe_inner_he.add_nuclide('Mo97',1.19e-6)
mat_pipe_inner_he.add_nuclide('Mo98',3.01e-6)
mat_pipe_inner_he.add_nuclide('Mo100',1.21e-6)
mat_pipe_inner_he.temperature = mat_temp_global
mat_pipe_inner_he.set_density("atom/b-cm",0.0008815)
mat_list.append(mat_pipe_inner_he)

mat_pipe_inner_region = openmc.Material(406,name="Heat Piper Inner Region Smear (vapour K + liquid K + wick)")
mat_pipe_inner_region.add_nuclide('C12',3.808e-6)
mat_pipe_inner_region.add_nuclide('Si28',1.856e-5)
mat_pipe_inner_region.add_nuclide('Si29',9.463e-7)
mat_pipe_inner_region.add_nuclide('Si30',6.242e-7)
mat_pipe_inner_region.add_nuclide('P31',8.277e-7)
mat_pipe_inner_region.add_nuclide('S32',4.949e-7)
mat_pipe_inner_region.add_nuclide('S33',3.908e-9)
mat_pipe_inner_region.add_nuclide('S34',2.214e-8)
mat_pipe_inner_region.add_nuclide('S36',6.012e-11)
mat_pipe_inner_region.add_nuclide('Cr50',1.585e-5)
mat_pipe_inner_region.add_nuclide('Cr52',3.055e-4)
mat_pipe_inner_region.add_nuclide('Cr53',3.465e-5)
mat_pipe_inner_region.add_nuclide('Cr54',8.625e-6)
mat_pipe_inner_region.add_nuclide('Mn55',2.059e-5)
mat_pipe_inner_region.add_nuclide('Fe54',7.814e-5)
mat_pipe_inner_region.add_nuclide('Fe56',1.226e-3)
mat_pipe_inner_region.add_nuclide('Fe57',2.833e-5)
mat_pipe_inner_region.add_nuclide('Fe58',2.673e-6)
mat_pipe_inner_region.add_nuclide('Ni58',1.553e-4)
mat_pipe_inner_region.add_nuclide('Ni60',5.98e-5)
mat_pipe_inner_region.add_nuclide('Ni61',2.60e-6)
mat_pipe_inner_region.add_nuclide('Ni62',8.291e-6)
mat_pipe_inner_region.add_nuclide('Ni64',2.112e-6)
mat_pipe_inner_region.add_nuclide('Mo92',4.258e-6)
mat_pipe_inner_region.add_nuclide('Mo94',2.664e-6)
mat_pipe_inner_region.add_nuclide('Mo95',4.613e-6)
mat_pipe_inner_region.add_nuclide('Mo96',4.845e-6)
mat_pipe_inner_region.add_nuclide('Mo97',2.784e-6)
mat_pipe_inner_region.add_nuclide('Mo98',7.058e-6)
mat_pipe_inner_region.add_nuclide('Mo100',2.831e-6)
mat_pipe_inner_region.add_nuclide('K39',5.895e-3)
mat_pipe_inner_region.add_nuclide('K40',7.586e-7)
mat_pipe_inner_region.add_nuclide('K41',4.254e-4)
mat_pipe_inner_region.temperature = mat_temp_global
mat_pipe_inner_region.set_density("atom/b-cm",0.08324)
mat_list.append(mat_pipe_inner_region)

mat_pipe_he_gap_smear = openmc.Material(407,name="Heat Pipe He Gap SS316 Smear")
mat_pipe_he_gap_smear.add_nuclide('He4',5.629e-6)
mat_pipe_he_gap_smear.add_nuclide('C12',1.287e-4)
mat_pipe_he_gap_smear.add_nuclide('Si28',6.276e-4)
mat_pipe_he_gap_smear.add_nuclide('Si29',3.199e-5)
mat_pipe_he_gap_smear.add_nuclide('Si30',2.11e-5)
mat_pipe_he_gap_smear.add_nuclide('P31',2.798e-5)
mat_pipe_he_gap_smear.add_nuclide('S32',1.673e-5)
mat_pipe_he_gap_smear.add_nuclide('S33',1.321e-7)
mat_pipe_he_gap_smear.add_nuclide('S34',7.486e-7)
mat_pipe_he_gap_smear.add_nuclide('S36',2.032e-9)
mat_pipe_he_gap_smear.add_nuclide('Cr50',5.357e-4)
mat_pipe_he_gap_smear.add_nuclide('Cr52',1.033e-2)
mat_pipe_he_gap_smear.add_nuclide('Cr53',1.171e-3)
mat_pipe_he_gap_smear.add_nuclide('Cr54',2.916e-4)
mat_pipe_he_gap_smear.add_nuclide('Mn55',6.96e-4)
mat_pipe_he_gap_smear.add_nuclide('Fe54',2.643e-3)
mat_pipe_he_gap_smear.add_nuclide('Fe56',0.04145)
mat_pipe_he_gap_smear.add_nuclide('Fe57',9.576e-4)
mat_pipe_he_gap_smear.add_nuclide('Fe58',9.034e-5)
mat_pipe_he_gap_smear.add_nuclide('Ni58',5.248e-3)
mat_pipe_he_gap_smear.add_nuclide('Ni60',2.022e-3)
mat_pipe_he_gap_smear.add_nuclide('Ni61',8.789e-5)
mat_pipe_he_gap_smear.add_nuclide('Ni62',2.802e-4)
mat_pipe_he_gap_smear.add_nuclide('Ni64',7.139e-5)
mat_pipe_he_gap_smear.add_nuclide('Mo92',1.439e-4)
mat_pipe_he_gap_smear.add_nuclide('Mo94',9.004e-5)
mat_pipe_he_gap_smear.add_nuclide('Mo95',1.559e-4)
mat_pipe_he_gap_smear.add_nuclide('Mo96',1.638e-4)
mat_pipe_he_gap_smear.add_nuclide('Mo97',9.413e-5)
mat_pipe_he_gap_smear.add_nuclide('Mo98',2.387e-4)
mat_pipe_he_gap_smear.add_nuclide('Mo100',9.57e-5)
mat_pipe_he_gap_smear.temperature = mat_temp_global
mat_pipe_he_gap_smear.set_density("atom/b-cm",0.06771)
mat_list.append(mat_pipe_he_gap_smear)

# Smeared heat pipe wick materials
mat_pipe_wick_sat = openmc.Material.mix_materials(
        material_id=408,
        materials=[mat_pipe_ss316, mat_pipe_k_liquid],
        fracs=[1-mat_pipe_porosity, mat_pipe_porosity],
        percent_type="vo",
        name="Heat Pipe Wick (Saturated)",
        temperature=mat_temp_global
        )
mat_list.append(mat_pipe_wick_sat)

mat_pipe_wick_unsat = openmc.Material.mix_materials(
        material_id=409,
        materials=[mat_pipe_ss316, mat_pipe_k_gas],
        fracs=[1-mat_pipe_porosity, mat_pipe_porosity],
        percent_type="vo",
        name="Heat Pipe Wick (Unsaturated)",
        temperature=mat_temp_global
        )
mat_list.append(mat_pipe_wick_unsat)

## 500 - Core-level details ##

mat_core_be_reflector = openmc.Material(500,name="Core Beryllium Reflector")
mat_core_be_reflector.add_nuclide("Be9",1.0)
mat_core_be_reflector.temperature = mat_temp_global
mat_core_be_reflector.set_density("g/cc",1.848)
mat_core_be_reflector.add_s_alpha_beta("c_Be")
mat_list.append(mat_core_be_reflector)

mat_core_be_control_drum = openmc.Material(501,name="Core Beryllium Control Drum")
mat_core_be_control_drum.add_nuclide("Be9",1.0)
mat_core_be_control_drum.temperature = mat_temp_global
mat_core_be_control_drum.set_density("g/cc",1.848)
mat_core_be_control_drum.add_s_alpha_beta("c_Be")
mat_list.append(mat_core_be_control_drum)

mat_core_b4c_control_drum = openmc.Material(502,name="Core B4C Control Drum Absorber")
mat_core_b4c_control_drum.add_nuclide("B10",0.76)
mat_core_b4c_control_drum.add_nuclide("B11",0.04)
mat_core_b4c_control_drum.add_nuclide("C12",0.20)
mat_core_b4c_control_drum.temperature = mat_temp_global
mat_core_b4c_control_drum.set_density("g/cc",2.51)
mat_list.append(mat_core_b4c_control_drum)

mat_core_b4c_central = openmc.Material(503,name="Core B4C Central Absorber")
mat_core_b4c_central.add_nuclide("B10",0.76)
mat_core_b4c_central.add_nuclide("B11",0.04)
mat_core_b4c_central.add_nuclide("C12",0.20)
mat_core_b4c_central.temperature = mat_temp_global
mat_core_b4c_central.set_density("g/cc",1.25)
mat_list.append(mat_core_b4c_central)


############ GEOMETRY ############ 
# All measurements are in cm unless specified
geo_const_pin_pitch = 2.3
geo_const_lattice_pitch = 26.752
geo_const_triso_radius = 1.0
geo_const_triso_packing = 0.40
geo_const_height = cli_args.z_height
geo_const_z_max = openmc.ZPlane(z0=+geo_const_height,boundary_type="transmission" if cli_args.geometry_level=="core" else "reflective")
geo_const_z_min = openmc.ZPlane(z0=-geo_const_height,boundary_type="transmission" if cli_args.geometry_level=="core" else "reflective")

geo_none_universe = openmc.Universe(name="Empty Universe",cells=[openmc.Cell(name="Empty Cell",fill=None)])

### Inidividual Pin Universes ###

## TRISO fuel kernel ##
geo_triso_spheres = [openmc.Sphere(r=rad*1e-2, x0=0.0, y0=0.0) for rad in [2.125, 3.125, 3.525, 3.875, 4.275]]
geo_triso_cells = [openmc.Cell(name="TRISO Kernel UCO",fill=mat_triso_uco,region=-geo_triso_spheres[0]),
                   openmc.Cell(name="TRISO Kernel Buffer",fill=mat_triso_buffer,region=+geo_triso_spheres[0]&-geo_triso_spheres[1]),
                   openmc.Cell(name="TRISO Kernel PyC1",fill=mat_triso_pyc1,region=+geo_triso_spheres[1]&-geo_triso_spheres[2]),
                   openmc.Cell(name="TRISO Kernel SiC",fill=mat_triso_sic,region=+geo_triso_spheres[2]&-geo_triso_spheres[3]),
                   openmc.Cell(name="TRISO Kernel PyC2",fill=mat_triso_pyc2,region=+geo_triso_spheres[3]&-geo_triso_spheres[4])]
geo_triso_universe = openmc.Universe(name="TRISO Kernel Universe",cells=geo_triso_cells)

## TRISO compact ##
geo_triso_compact_region = -openmc.ZCylinder(r=geo_const_triso_radius,x0=0.0,y0=0.0)&+geo_const_z_min&-geo_const_z_max
geo_triso_compact_domain = (4,4,int(4*geo_const_height)) if cli_args.triso_lattice else (1,1,1)

# Generate packing
if cli_args.packing:
    geo_triso_compact_packing = np.loadtxt("positions.out",delimiter=",")
else:
    geo_triso_compact_packing = openmc.model.pack_spheres(
            radius=geo_triso_spheres[-1].r,
            region=-openmc.ZCylinder(r=geo_const_triso_radius-geo_triso_spheres[-1].r,x0=0.0,y0=0.0) & \
                    +openmc.ZPlane(z0=-geo_const_height+geo_triso_spheres[-1].r) & \
                    -openmc.ZPlane(z0=+geo_const_height-geo_triso_spheres[-1].r),
            pf=geo_const_triso_packing
            )

# Export positions to file for backup
np.savetxt("positions.out",geo_triso_compact_packing,delimiter=",")

# Create the TRISO particles
geo_triso_compact_trisos = [openmc.model.TRISO(geo_triso_spheres[4].r, geo_triso_universe, loc) for loc in geo_triso_compact_packing]

# Support with a lattice for domain decomposition
geo_triso_compact_lattice = openmc.model.create_triso_lattice(
        trisos=geo_triso_compact_trisos,
        lower_left=(-geo_const_triso_radius,-geo_const_triso_radius,-geo_const_height),
        pitch=tuple((2*p/d for p,d in zip((geo_const_triso_radius,geo_const_triso_radius,geo_const_height),geo_triso_compact_domain))),
        shape=geo_triso_compact_domain,
        background=mat_graphite_matrix
        )

# Cast into cells and a universe
geo_triso_compact_fuel_cell = openmc.Cell(name="TRISO Compact Cell",fill=geo_triso_compact_lattice,region=geo_triso_compact_region)
geo_triso_compact_surrounding_cell = openmc.Cell(name="TRISO Compact Surrounding Cell",fill=mat_graphite_matrix,region=~geo_triso_compact_region)

geo_triso_compact_universe = openmc.Universe(name="TRISO Compact Universe",cells=[geo_triso_compact_fuel_cell,geo_triso_compact_surrounding_cell])

## Homogenized TRISO Compact ##
geo_vwh_fracs = []
geo_rpt_fracs = []
geo_triso_rpt_region = -openmc.ZCylinder(r=cli_args.radius,x0=0.0,y0=0.0)&+geo_const_z_min&-geo_const_z_max
geo_n_triso = len(geo_triso_compact_packing)
print("\nTRISO Stacking Information from Packing Calculation")
print(f" -> Total TRISO pebbles packed:    {geo_n_triso}")
print(f" -> Total TRISO volume in compact: {geo_n_triso*4/3*math.pi*geo_triso_spheres[4].r**3:.4f}")
print(f" -> Minimum RPT cylinder radius:   {(geo_n_triso*4/3*geo_triso_spheres[4].r**3/(2*geo_const_height))**0.5:.4f}")
print(f" -> Maximum RRPT ring radius:      {(1.15**2 - geo_n_triso*4/3*geo_triso_spheres[4].r**3/(2*geo_const_height))**0.5:.4f}")
print("TRISO Stacking calculation complete!\n")

# Get volume fractions for mixing
geo_triso_compact_vwh_v_total = math.pi*geo_const_triso_radius**2*2*geo_const_height
geo_triso_compact_vwh_v_matrix = geo_triso_compact_vwh_v_total - geo_n_triso * 4/3*math.pi*geo_triso_spheres[4].r**3
geo_triso_compact_rpt_v_total = math.pi*cli_args.radius**2*2*geo_const_height
geo_triso_compact_rpt_v_matrix = geo_triso_compact_rpt_v_total - geo_n_triso * 4/3*math.pi*geo_triso_spheres[4].r**3

geo_v_layers = []
geo_temp_prev_volume = 0.0

for layer in geo_triso_spheres:
    geo_temp_cumulative_volume = 4/3*math.pi*layer.r**3
    geo_temp_shell_volume = geo_temp_cumulative_volume-geo_temp_prev_volume
    geo_v_layers.append(geo_temp_shell_volume*geo_n_triso)
    geo_temp_prev_volume=geo_temp_cumulative_volume

geo_kernel_fracs = [v / sum(geo_v_layers) for v in geo_v_layers]
geo_vwh_fracs = [v / geo_triso_compact_vwh_v_total for v in (geo_v_layers + [geo_triso_compact_vwh_v_matrix])]
geo_rpt_fracs = [v / geo_triso_compact_rpt_v_total for v in (geo_v_layers + [geo_triso_compact_rpt_v_matrix])]

# Strip TSL temporarily for mixing
mat_tsl_saved = {}
for mat in mat_list:
    if hasattr(mat, "_sab") and mat._sab:
        mat_tsl_saved[mat]=list(mat._sab)
        mat._sab=[]

# TRISO only VWH
mat_kernel_vwh = openmc.Material.mix_materials(
        materials=[mat_triso_uco, mat_triso_buffer, mat_triso_pyc1, mat_triso_sic, mat_triso_pyc2],
        fracs=geo_kernel_fracs,
        percent_type="vo",
        name="VWH kernel"
        )
mat_list.append(mat_kernel_vwh)

# Simple VWH
mat_triso_vwh = openmc.Material.mix_materials(
        materials=[mat_triso_uco, mat_triso_buffer, mat_triso_pyc1, mat_triso_sic, mat_triso_pyc2, mat_graphite_matrix],
        fracs=geo_vwh_fracs,
        percent_type="vo",
        name="VWH Homogenized TRISO Compact"
        )
mat_list.append(mat_triso_vwh)

# RPT Homogenization
if cli_args.technique=="rpt": # There is a minimum viable radius for RPT, so only calculate this if needed to aovid errors
    mat_triso_rpt = openmc.Material.mix_materials(
            materials=[mat_triso_uco, mat_triso_buffer, mat_triso_pyc1, mat_triso_sic, mat_triso_pyc2, mat_graphite_matrix],
            fracs=geo_rpt_fracs,
            percent_type="vo",
            name="RPT Homogenized TRISO Compact"
            )
else:
    mat_triso_rpt = openmc.Material(name="PLACEHOLDER//DO NOT USE")
    mat_triso_rpt.add_nuclide("U235",1.0)
    mat_triso_rpt.set_density("g/cc",100.0)
mat_list.append(mat_triso_rpt)

# Restore TSL for other materials
for mat, tsl_list in mat_tsl_saved.items():
    mat._sab = tsl_list

# RRPT transformation
geo_rrpt_r_out = math.sqrt(4*geo_n_triso*geo_triso_spheres[4].r**3/(6*geo_const_height) + cli_args.radius**2)
geo_rrpt_region = +openmc.ZCylinder(r=cli_args.radius,x0=0.0,y0=0.0) & -openmc.ZCylinder(r=geo_rrpt_r_out,x0=0.0,y0=0.0) \
                & +geo_const_z_min & -geo_const_z_max

# Cast into cells and a universe
geo_triso_vwh_fuel_cell = openmc.Cell(name="TRISO VWH Fuel Cell",fill=mat_triso_vwh,region=geo_triso_compact_region)
geo_triso_vwh_surrounding_cell = openmc.Cell(name="TRISO VWH Surrounding Cell",fill=mat_graphite_matrix,region=~geo_triso_compact_region)
geo_triso_vwh_universe = openmc.Universe(name="TRISO VWH Universe",cells=[geo_triso_vwh_fuel_cell,geo_triso_vwh_surrounding_cell])

geo_triso_rpt_fuel_cell = openmc.Cell(name="TRISO RPT Fuel Cell",fill=mat_triso_rpt,region=geo_triso_rpt_region)
geo_triso_rpt_surrounding_cell = openmc.Cell(name="TRISO RPT Surrounding Cell",fill=mat_graphite_matrix,region=~geo_triso_rpt_region)
geo_triso_rpt_universe = openmc.Universe(name="TRISO RPT Universe",cells=[geo_triso_rpt_fuel_cell,geo_triso_rpt_surrounding_cell])

geo_triso_rrpt_fuel_cell = openmc.Cell(name="TRISO RRPT Fuel Cell",fill=mat_kernel_vwh,region=geo_rrpt_region)
geo_triso_rrpt_surrouding_cell = openmc.Cell(name="TRISO RRPT Surrounding Cell",fill=mat_graphite_matrix,region=~geo_rrpt_region)
geo_triso_rrpt_universe = openmc.Universe(name="TRISO RRPT Universe",cells=[geo_triso_rrpt_fuel_cell,geo_triso_rrpt_surrouding_cell])

# Dynamically pick the correct one
geo_homo_map = {
        "none":geo_triso_compact_universe,
        "vwh":geo_triso_vwh_universe,
        "rpt":geo_triso_rpt_universe,
        "rrpt":geo_triso_rrpt_universe
        }

geo_dyn_universe = geo_homo_map[cli_args.technique]

## Heat Pipe ##
geo_pipe_cylinders = [openmc.ZCylinder(r=rad, x0=0.0, y0=0.0) for rad in [0.80, 0.90, 0.97, 1.05, 1.07]]

# Dynamic for if HP is being studied or not
if cli_args.model_hp:
    geo_pipe_z_pool = openmc.ZPlane(z0=cli_args.pool_height-80)
    geo_pipe_z_wick = openmc.ZPlane(z0=cli_args.sat_height-80)
    geo_pipe_cells = [openmc.Cell(name="Heat Pipe inner K Gas",fill=mat_pipe_k_gas,region=-geo_pipe_cylinders[0]&+geo_pipe_z_pool),
                      openmc.Cell(name="Heat Pipe inner K Liquid",fill=mat_pipe_k_liquid,region=-geo_pipe_cylinders[0]&-geo_pipe_z_pool),
                      openmc.Cell(name="Heat Pipe Saturated Wick",fill=mat_pipe_wick_sat,region=+geo_pipe_cylinders[0]&-geo_pipe_cylinders[1]&-geo_pipe_z_wick),
                      openmc.Cell(name="Heat Pipe Unsaturated Wick",fill=mat_pipe_wick_unsat,region=+geo_pipe_cylinders[0]&-geo_pipe_cylinders[1]&+geo_pipe_z_wick),
                      openmc.Cell(name="Heat Pipe Outer K Gas",fill=mat_pipe_k_gas,region=+geo_pipe_cylinders[1]&-geo_pipe_cylinders[2]&+geo_pipe_z_pool),
                      openmc.Cell(name="Heat Pipe Outer K Liquid",fill=mat_pipe_k_liquid,region=+geo_pipe_cylinders[1]&-geo_pipe_cylinders[2]&-geo_pipe_z_pool),
                      openmc.Cell(name="Heat Pipe SS316 Envelope",fill=mat_pipe_ss316,region=+geo_pipe_cylinders[2]&-geo_pipe_cylinders[3]),
                      openmc.Cell(name="Heat Pipe Helium Gap",fill=mat_pipe_helium,region=+geo_pipe_cylinders[3]&-geo_pipe_cylinders[4]),
                      openmc.Cell(name="Heat Pipe Graphite Matrix",fill=mat_graphite_matrix,region=+geo_pipe_cylinders[4])]
else:
    geo_pipe_cells = [openmc.Cell(name="Heat Pipe K Gas",fill=mat_pipe_k_gas,region=-geo_pipe_cylinders[0]),
                      openmc.Cell(name="Heat Pipe Wick",fill=mat_pipe_wick_smear,region=+geo_pipe_cylinders[0]&-geo_pipe_cylinders[1]),
                      openmc.Cell(name="Heat Pipe K Layer",fill=mat_pipe_k_liquid,region=+geo_pipe_cylinders[1]&-geo_pipe_cylinders[2]),
                      openmc.Cell(name="Heat Pipe SS316 Envelope",fill=mat_pipe_ss316,region=+geo_pipe_cylinders[2]&-geo_pipe_cylinders[3]),
                      openmc.Cell(name="Heat Pipe Helium Gap",fill=mat_pipe_helium,region=+geo_pipe_cylinders[3]&-geo_pipe_cylinders[4]),
                      openmc.Cell(name="Heat Pipe Graphite Matrix",fill=mat_graphite_matrix,region=+geo_pipe_cylinders[4])]

# Assign to a universe
geo_pipe_universe = openmc.Universe(name="Heat Pipe Universe",cells=geo_pipe_cells)

## Moderator Pin ##
geo_mod_cylinders = [openmc.ZCylinder(r=rad, x0=0.0, y0=0.0) for rad in [0.825, 0.875, 0.900, 0.920]]
geo_mod_cells = [openmc.Cell(name="Moderator Pin YH2",fill=mat_mod_yh2,region=-geo_mod_cylinders[0]),
                 openmc.Cell(name="Moderator Pin Inner Helium Gap",fill=mat_mod_ss316_he_smear,region=+geo_mod_cylinders[0]&-geo_mod_cylinders[1]),
                 openmc.Cell(name="Moderator Pin SS316 Envelope",fill=mat_mod_ss316,region=+geo_mod_cylinders[1]&-geo_mod_cylinders[2]),
                 openmc.Cell(name="Moderator Pin Outer Helium Gap",fill=mat_mod_ss316_he_smear,region=+geo_mod_cylinders[2]&-geo_mod_cylinders[3]),
                 openmc.Cell(name="Moderator Pin Graphite Matrix",fill=mat_graphite_matrix,region=+geo_mod_cylinders[3])]
geo_mod_universe = openmc.Universe(name="Moderator Pin Universe",cells=geo_mod_cells)

### Assembly-level Geometry ###

geo_assembly_lattice = openmc.HexLattice()
geo_assembly_lattice.center=(0,0)
geo_assembly_lattice.pitch=(geo_const_pin_pitch,)
geo_assembly_lattice.outer=openmc.Universe(cells=[openmc.Cell(fill=mat_graphite_matrix)])

# Create the lattice structure
geo_assembly_lattice_map = []
for i in range(0,7):
    if i%2 == 0:
        ring = [geo_pipe_universe]*(36-6*i)
        ring[1::2] = [geo_dyn_universe]*len(ring[1::2])
    else:
        ring = [geo_mod_universe]*(36-6*i)
        ring[::2] = [geo_dyn_universe]*len(ring[::2])
    geo_assembly_lattice_map.append(ring)

geo_assembly_lattice_map[-1] = [geo_pipe_universe if cli_args.geometry_level!="pin" else geo_dyn_universe]

geo_assembly_lattice.universes=geo_assembly_lattice_map

# Bound it to a unit assembly
geo_assembly_region = openmc.model.HexagonalPrism(
        edge_length=geo_const_lattice_pitch/3**0.5 if cli_args.geometry_level!="pin" else geo_const_pin_pitch/3**0.5,
        orientation="y" if cli_args.geometry_level !="pin" else "x",
        origin=(0.0,0.0),
        boundary_type="transmission" if cli_args.geometry_level=="core" else "reflective"
        )

# Clip the lattice to the space and make a universe
geo_assembly_cell = openmc.Cell(name="Unit Assembly Cell",fill=geo_assembly_lattice,region=-geo_assembly_region&+geo_const_z_min&-geo_const_z_max)
geo_assembly_universe = openmc.Universe(name="Unit Lattice Assembly Universe",cells=[geo_assembly_cell])

## Core-level Geometry ##
geo_const_core_z_max = openmc.ZPlane(z0=80)
geo_const_core_z_min = openmc.ZPlane(z0=-80)
geo_const_refl_z_max = openmc.ZPlane(z0=100, boundary_type="vacuum")
geo_const_refl_z_min = openmc.ZPlane(z0=-100, boundary_type="vacuum")
geo_const_core_cyl = openmc.ZCylinder(r=115, x0=0.0, y0=0.0, boundary_type="vacuum")

# Reflector region cell and universe
geo_refl_cell = openmc.Cell(name="Core Reflector Cell",fill=mat_core_be_reflector)
geo_refl_universe = openmc.Universe(name="Core Reflector Universe",cells=[geo_refl_cell])

# Control Drum cell and universe
geo_cd_cyl = [openmc.ZCylinder(r=rad, x0=0.0, y0=0.0) for rad in [12.250, 13.250]]
geo_cd_pl = openmc.XPlane(x0=88.92-3*geo_const_lattice_pitch)
geo_cd_abs_cell = openmc.Cell(name="Control Drum Absorber Cell",fill=mat_core_b4c_control_drum,region=+geo_cd_cyl[0]&-geo_cd_cyl[1]&+geo_cd_pl)
geo_cd_drum_cell = openmc.Cell(name="Core Control Drum Cell",fill=mat_core_be_control_drum,region=~geo_cd_abs_cell.region&-geo_cd_cyl[1])
geo_cd_outer_cell = openmc.Cell(name="Core Control Drum Outer Cell",fill=None,region=+geo_cd_cyl[1])
geo_cd_universe = openmc.Universe(name="Core Control Drum Universe",cells=[geo_cd_abs_cell,geo_cd_drum_cell,geo_cd_outer_cell])

# Core 2D radial lattice
geo_core_rad_lattice=openmc.HexLattice()
geo_core_rad_lattice.orientation="x"
geo_core_rad_lattice.center=(0,0)
geo_core_rad_lattice.pitch=(geo_const_lattice_pitch,)
geo_core_rad_lattice.outer=geo_none_universe

# Assemble lattice map
geo_dyn_cd_rot = 180 if cli_args.drums_in else 0
geo_core_rad_lattice_map=[]
geo_core_rad_lattice_map.append([rot(geo_cd_universe,geo_dyn_cd_rot-30-(i-2)*15) if i%4==2 else geo_none_universe if i%4==0 else geo_refl_universe for i in range(24)])
geo_core_rad_lattice_map.append([rot(geo_cd_universe,geo_dyn_cd_rot-i*20) if i%3==0 else geo_assembly_universe for i in range(18)])
geo_core_rad_lattice_map.append([geo_assembly_universe]*12)
geo_core_rad_lattice_map.append([geo_assembly_universe]*6)
geo_core_rad_lattice_map.append([geo_none_universe])
geo_core_rad_lattice.universes=geo_core_rad_lattice_map

# Cell and universe
geo_core_rad_cell=openmc.Cell(name="Core Radial Slice Cell",fill=geo_core_rad_lattice)
geo_core_rad_universe = openmc.Universe(cells=[geo_core_rad_cell])

# Axial Reflector lattice
geo_core_rad_refl_lattice=openmc.HexLattice()
geo_core_rad_refl_lattice.center=(0,0)
geo_core_rad_refl_lattice.orientation="x"
geo_core_rad_refl_lattice.pitch=(geo_const_lattice_pitch,)
geo_core_rad_refl_lattice.outer=geo_none_universe

# Axial Reflector lattice map
geo_core_rad_refl_lattice_map = []
geo_core_rad_refl_lattice_map.append([rot(geo_cd_universe,geo_dyn_cd_rot-30-(i-2)*15) if i%4==2 else geo_none_universe if i%4==0 else geo_refl_universe for i in range(24)])
geo_core_rad_refl_lattice_map.append([rot(geo_cd_universe,geo_dyn_cd_rot-i*20) if i%3==0 else geo_refl_universe for i in range(18)])
geo_core_rad_refl_lattice_map.append([geo_refl_universe]*12)
geo_core_rad_refl_lattice_map.append([geo_refl_universe]*6)
geo_core_rad_refl_lattice_map.append([geo_none_universe])
geo_core_rad_refl_lattice.universes = geo_core_rad_refl_lattice_map

geo_core_refl_top_cell = openmc.Cell(name="Top Axial Reflector Cell",fill=geo_core_rad_refl_lattice,
                                     region=-geo_const_core_cyl & +geo_const_core_z_max & -geo_const_refl_z_max)
geo_core_refl_bot_cell = openmc.Cell(name="Bottom Axial Reflector Cell",fill=geo_core_rad_refl_lattice,
                                     region=-geo_const_core_cyl & -geo_const_core_z_min & +geo_const_refl_z_min)

# 3D Stacking Lattice
geo_core_axial_lattice = openmc.RectLattice()
geo_core_axial_lattice.lower_left = (-sys.float_info.max/2, -sys.float_info.max/2, -int(160//geo_const_height)/2*geo_const_height)
geo_core_axial_lattice.pitch=(sys.float_info.max, sys.float_info.max, 2*geo_const_height)
geo_core_axial_lattice.universes = [[[geo_core_rad_universe]] for _ in range(int(160//(2*geo_const_height)))]
geo_core_axial_lattice.outer=geo_none_universe

geo_core_cell = openmc.Cell(name="Core 3D Cell",fill=geo_core_axial_lattice, region=-geo_const_core_cyl & +geo_const_core_z_min & -geo_const_core_z_max)
geo_core_universe = openmc.Universe(cells=[geo_core_cell, geo_core_refl_top_cell, geo_core_refl_bot_cell])
#geo_core_universe = openmc.Universe(cells=[geo_core_refl_bot_cell])

## Final Geometry Universe
geometry = openmc.Geometry(geo_core_universe if cli_args.geometry_level=="core" else geo_assembly_universe)

############ SETTINGS ############ 
settings = openmc.Settings()
settings.particles = 100000
settings.batches = 200
settings.inactive = 50
#settings.verbosity = 10

set_dyn_radial = geo_const_core_cyl.r if cli_args.geometry_level=="core" else geo_const_pin_pitch/3**0.5 if cli_args.geometry_level=="pin" else geo_const_lattice_pitch/3**0.5
set_dyn_axial = 180 if cli_args.geometry_level=="core" else geo_const_height

set_source_pts = openmc.IndependentSource(
        space=openmc.stats.Box(
            lower_left = (-set_dyn_radial, -set_dyn_radial, -set_dyn_axial),
            upper_right = (set_dyn_radial, set_dyn_radial, set_dyn_axial)
            ))
        

settings.source = set_source_pts

# Mesh for Shannon Entropy
set_entropy_mesh = openmc.RegularMesh()
set_entropy_mesh.lower_left=(-set_dyn_radial,-set_dyn_radial,-set_dyn_axial)
set_entropy_mesh.upper_right=(+set_dyn_radial,+set_dyn_radial,+set_dyn_axial)
set_entropy_mesh.dimension=(30,30,30)

settings.entropy_mesh=set_entropy_mesh
settings.temperature = {'method': "interpolation"}

############ TALLIES ############ 
tal_global = openmc.Tally(name="Global Reaction Rate")
tal_global.scores = ["absorption","fission"]

tallies = openmc.Tallies([tal_global])
tallies.export_to_xml()

## Export ##
materials = openmc.Materials(mat_list)
materials.export_to_xml()
geometry.export_to_xml()
settings.export_to_xml()

# Run!
openmc.run() if cli_args.execute else None