#!/bin/python3.12
import openmc
import argparse
import sys
import hashlib
import time
import math

############ COMMAND LINE ARGUMENTS ############ 
cli_parser = argparse.ArgumentParser(
        prog="Basic Pin TRISO Homogenization Study",
        description="Program to investigate various homogenization techniques for simple TRISO compacts."
        )

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.5",type=float,help="Radius for RPT or RRPT homogenization [cm], default 0.5 cm. In the case of RRPT, this is the inner radius.")
cli_parser.add_argument("-p","--packing-factor",default=0.29,type=float,help="Packing factor for TRISO packing (0.0, 1.0), default 0.29")
cli_parser.add_argument("-f","--fuel",default="uo2",choices=["uo2","uco5","uco10","uco20","un"],help="What type of fuel to use [uo2/uco10/uco20/un], default uo2.")
cli_parser.add_argument("-e","--execute",default=False,action="store_true",help="Whether or not to automatically run OpenMC, default false.")
cli_args = cli_parser.parse_args()

############ 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 ############
mat_temp_global = 293
mat_temp_fuel = 293
mat_list = []

## 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_graphite_matrix.add_s_alpha_beta("c_Graphite")
mat_list.append(mat_graphite_matrix)

## 200s - TRISO Pins ##
mat_triso_uo2 = openmc.Material(200,name="TRISO Pin UO2 Kernel")
mat_triso_uo2.add_element("U",1.0,enrichment=5.0)
mat_triso_uo2.add_element("O",2.0)
mat_triso_uo2.set_density("g/cc",10.97)
mat_triso_uo2.temperature = mat_temp_fuel
mat_triso_uo2.add_s_alpha_beta("c_U_in_UO2")
mat_triso_uo2.add_s_alpha_beta("c_O_in_UO2")
mat_list.append(mat_triso_uo2)

mat_triso_uco_5wo = openmc.Material(201,name="TRISO Pin 5% UCO Kernel")
mat_triso_uco_5wo.add_nuclide("U235",0.0172417)
mat_triso_uco_5wo.add_nuclide("U238",0.3275923)
mat_triso_uco_5wo.add_nuclide("C12",0.13793)
mat_triso_uco_5wo.add_nuclide("O16",0.51724)
mat_triso_uco_5wo.set_density("g/cc",10.744)
mat_triso_uco_5wo.temperature = mat_temp_fuel
mat_triso_uco_5wo.add_s_alpha_beta("c_U_in_UC_10p")
mat_triso_uco_5wo.add_s_alpha_beta("c_C_in_UC_10p")
mat_list.append(mat_triso_uco_5wo)

mat_triso_uco_10wo = openmc.Material(202,name="TRISO Pin 10% UCO Kernel")
mat_triso_uco_10wo.add_nuclide("U235",0.034483)
mat_triso_uco_10wo.add_nuclide("U238",0.310400)
mat_triso_uco_10wo.add_nuclide("C12",0.13793)
mat_triso_uco_10wo.add_nuclide("O16",0.51724)
mat_triso_uco_10wo.set_density("g/cc",10.744)
mat_triso_uco_10wo.temperature = mat_temp_fuel
mat_triso_uco_10wo.add_s_alpha_beta("c_U_in_UC_10p")
mat_triso_uco_10wo.add_s_alpha_beta("c_C_in_UC_10p")
mat_list.append(mat_triso_uco_10wo)

mat_triso_uco_20wo = openmc.Material(203,name="TRISO Pin 20% UCO Kernel")
mat_triso_uco_20wo.add_nuclide("U235",0.068794)
mat_triso_uco_20wo.add_nuclide("U238",0.27604)
mat_triso_uco_20wo.add_nuclide("C12",0.13793)
mat_triso_uco_20wo.add_nuclide("O16",0.51724)
mat_triso_uco_20wo.set_density("g/cc",10.744)
mat_triso_uco_20wo.temperature = mat_temp_fuel
mat_triso_uco_20wo.add_s_alpha_beta("c_U_in_UC_HALEU")
mat_triso_uco_20wo.add_s_alpha_beta("c_C_in_UC_HALEU")
mat_list.append(mat_triso_uco_20wo)

mat_triso_un = openmc.Material(204,name="TRISO Pin UN Kernel")
mat_triso_un.add_nuclide("U235",2.1125e-4)
mat_triso_un.add_nuclide("U238",2.8758e-2)
mat_triso_un.add_nuclide("U234",1.6795e-6)
mat_triso_un.add_nuclide("N14",2.8971e-4)
mat_triso_un.add_nuclide("N15",2.8681e-2)
mat_triso_un.set_density("g/cc",14.32)
mat_triso_un.temperature = mat_temp_fuel
mat_triso_un.add_s_alpha_beta("c_U_in_UN")
mat_triso_un.add_s_alpha_beta("c_N_in_UN")
mat_list.append(mat_triso_un)

mat_triso_buffer = openmc.Material(205,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(206,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(207,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(208,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)

## Dynamic Fuel Material ##
mat_fuel_map = {
        "uo2":mat_triso_uo2,
        "uco5":mat_triso_uco_5wo,
        "uco10":mat_triso_uco_10wo,
        "uco20":mat_triso_uco_20wo,
        "un":mat_triso_un
        }

mat_fuel_dyn = mat_fuel_map[cli_args.fuel]

############ GEOMETRY ############ 
geo_z_top=openmc.ZPlane(z0=5.0,boundary_type="reflective")
geo_z_bot=openmc.ZPlane(z0=-5.0,boundary_type="reflective")
geo_z_height = geo_z_top.z0 - geo_z_bot.z0
geo_compact_rad = 0.6225
geo_triso_spheres = [openmc.Sphere(r=rad*1e-4, x0=0.0, y0=0.0) for rad in [175, 275, 310, 345, 385]]
geo_triso_cells = [openmc.Cell(name="TRISO Kernel UCO",fill=mat_fuel_dyn,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_compact_rad,x0=0.0,y0=0.0)&+geo_z_bot&-geo_z_top
geo_triso_compact_domain = (10,10,100)

# Generate packing
geo_triso_compact_packing = openmc.model.pack_spheres(
        radius=geo_triso_spheres[4].r,
        region=geo_triso_compact_region,
        pf=cli_args.packing_factor
        )

# 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_compact_rad,-geo_compact_rad,geo_z_bot.z0),
        pitch=tuple((2*p/d for p,d in zip((geo_compact_rad,geo_compact_rad,10),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_compact_surrounding_region = -openmc.model.RectangularPrism(width=2.0, height=2.0, boundary_type="reflective") & +geo_z_bot & -geo_z_top
geo_triso_compact_surrounding_cell = openmc.Cell(name="TRISO Compact Surrounding Cell",fill=mat_graphite_matrix,region=geo_compact_surrounding_region & ~geo_triso_compact_region)
geo_triso_compact_universe = openmc.Universe(name="Root 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_z_bot&-geo_z_top
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_z_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_z_height))**0.5:.4f}")
print("TRISO Stacking calculation complete!\n")

# Get volume fractions for mixing
geo_triso_compact_vwh_v_total = math.pi*geo_compact_rad**2*2*geo_z_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_z_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_fuel_dyn, mat_triso_buffer, mat_triso_pyc1, mat_triso_sic, mat_triso_pyc2],
        fracs=geo_kernel_fracs,
        percent_type="vo",
        name="VWH kernel"
        )
mat_kernel_vwh.add_s_alpha_beta("c_Graphite")
mat_list.append(mat_kernel_vwh)

# Simple VWH
mat_triso_vwh = openmc.Material.mix_materials(
        materials=[mat_fuel_dyn, 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_triso_vwh.add_s_alpha_beta("c_Graphite")
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_fuel_dyn, 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"
            )
    mat_triso_rpt.add_s_alpha_beta("c_Graphite")
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_z_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_z_bot & -geo_z_top

# 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_compact_surrounding_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_compact_surrounding_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_compact_surrounding_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]

############ SETTINGS AND EXPORT ############

materials = openmc.Materials(mat_list)
materials.export_to_xml()

geometry = openmc.Geometry(geo_dyn_universe)
geometry.export_to_xml()


settings = openmc.Settings()
settings.particles=10000
settings.batches=200
settings.inactive=50
settings.source = openmc.IndependentSource(
            space=openmc.stats.Box(
            lower_left = (-1.0, -1.0, -5.0),
            upper_right = (1.0, 1.0, 5.0)
            ))
settings.export_to_xml()

openmc.run() if cli_args.execute else None