# this script uses cadquery to make a nice twisting CAD geometry
# the geometry is saved as a STEP file
# The step file is read into Gmesh and volume meshed (unstructured mesh)
# The mesh is saved as a .msh file
# The mesh file is converted to a h5 file (unstructured moab mesh format)
# The h5 file is read into an OpenMC unstructured mesh filter
# The simulation is run and a vtk file is produced for the mesh tally

from math import cos, floor, pi, sin

import cadquery as cq
from cadquery import exporters

import gmsh

# makes the CAD
def hypocycloid(t, r1, r2):
    return ((r1-r2)*cos(t)+r2*cos(r1/r2*t-t), (r1-r2)*sin(t)+r2*sin(-(r1/r2*t-t)))

def epicycloid(t, r1, r2):
    return ((r1+r2)*cos(t)-r2*cos(r1/r2*t+t), (r1+r2)*sin(t)-r2*sin(r1/r2*t+t))

def gear(t, r1=4, r2=1):
    if (-1)**(1+floor(t/2/pi*(r1/r2))) < 0:
        return epicycloid(t, r1, r2)
    else:
        return hypocycloid(t, r1, r2)

result = (cq.Workplane('XY').parametricCurve(lambda t: gear(t*2*pi, 6, 1))
    .twistExtrude(15, 90).faces('>Z').workplane().circle(2).cutThruAll())

# saves the step file
exporters.export(result, 'twist.step')


# meshes the step file
gmsh.initialize()
gmsh.option.setNumber("General.Terminal", 1)

volumes = gmsh.model.occ.importShapes('twist.step')
gmsh.model.occ.synchronize()

gmsh.model.mesh.generate(3)

gmsh.write("twist.vtk")
gmsh.finalize()

import meshio
# vtk to h5m appears to be one of the few conversions that does not print out errors
mesh = meshio.read("twist.vtk")
mesh.write("twist.h5m")

# neutronics simulation 
import openmc

air = openmc.Material(name='air')
air.set_density('g/cc', 0.001205)
air.add_element('N', 0.784431)

materials = openmc.Materials([air])
materials.export_to_xml()

src_pnt = openmc.stats.Point(xyz=(0.0, 0.0, 0.0))

source = openmc.Source(space=src_pnt)

settings = openmc.Settings()
settings.source = source

settings.run_mode = "fixed source"
settings.batches = 10
settings.particles = 100
settings.export_to_xml()

# we are using a csg cell but here is the code for a dagmc cell
# dagmc_univ = openmc.DAGMCUniverse(filename='twist.h5m').bounded_universe()
# geometry = openmc.Geometry(root=dagmc_univ)
# geometry.export_to_xml()

# UM geometry is 10 cm heigh so a 50cm sphere should be fine
surf1=openmc.Sphere(r=50, boundary_type='vacuum')
region1=-surf1
cell1=openmc.Cell(region=region1, fill=air)
geometry = openmc.Geometry([cell1])
geometry.export_to_xml()

unstructured_mesh = openmc.UnstructuredMesh("twist.h5m", library='moab')
mesh_filter = openmc.MeshFilter(unstructured_mesh)

tally = openmc.Tally()
tally.filters = [mesh_filter]
tally.scores = ['flux']
tally.estimator = 'tracklength'
tally.id = 1

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

openmc.run()

# prints warning
#  WARNING: Output for a MOAB mesh (mesh 1) was requested but will not be written.
#           Please use the Python API to generated the desired VTK tetrahedral
#           mesh.
# therefore attempting to get mesh from statepoint file


with openmc.StatePoint('statepoint.10.h5') as sp:
    tally = sp.tallies[1]
    umesh = sp.meshes[1]

    centroids = umesh.centroids  # why is this needed, why can't we use 
    mesh_vols = umesh.volumes  # why is this needed

    flux_scores = tally.get_values(scores=['flux']).squeeze()

print(flux_scores[0:10])
data_dict = {'Flux': flux_scores}

umesh.write_data_to_vtk(filename="flux.vtk", datasets=data_dict)

# then open paraview to see the vtk file produced