import os
from IPython.display import Image
import openmc
import openmc.lib

assert(openmc.lib._dagmc_enabled())

"""
Explanation: Unstructured Mesh: Tallies with CAD and Point Cloud Visualization
In the first notebook on this topic, we looked at how to set up a tally using an unstructured mesh in OpenMC.
In this notebook, we will explore using unstructured mesh in conjunction with CAD-based geometry to perform detailed geometry analysis on complex geomerty.
NOTE: This notebook will not run successfully if OpenMC has not been built with DAGMC support enabled.
End of explanation
"""


from IPython.display import display, clear_output
import urllib.request

manifold_geom_url = 'https://tinyurl.com/rp7grox' # 99 MB
manifold_mesh_url = 'https://tinyurl.com/wojemuh' # 5.4 MB

    
def download(url, filename):
    """
    Helper function for retrieving dagmc models
    """
    def progress_hook(count, block_size, total_size):
        prog_percent = 100 * count * block_size / total_size
        prog_percent = min(100., prog_percent)
        clear_output(wait=True)
        display('Downloading {}: {:.1f}%'.format(filename, prog_percent))
                    
    urllib.request.urlretrieve(url, filename, progress_hook)

"""
Explanation: We'll need to download our DAGMC geometry and unstructured mesh files. We'll be retrieving those using the function and URLs below.
End of explanation
"""


Image("./images/manifold-cad.png", width=800)

"""
Explanation: The model we'll be looking at in this example is a steel piping manifold:
End of explanation
"""


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

steel = openmc.Material(name='steel')
steel.set_density('g/cc', 8.0)
steel.add_element('Si', 0.010048)
steel.add_element('S', 0.00023)
steel.add_element('Fe', 0.669)
steel.add_element('Ni', 0.12)
steel.add_element('Mo', 0.025)
steel.add_nuclide('P31',0.00023)
steel.add_nuclide('Mn55',0.011014)

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

"""
Explanation: This is a nice example of a model which would be extremely difficult to model using CSG. To get started, we'll need two files: 
  1. the DAGMC gometry file on which we'll track particles and 
  2. a tetrahedral mesh of the piping structure on which we'll score tallies
To start, let's create the materials we'll need for this problem. The pipes are steel and we'll model the surrounding area as air.
End of explanation
"""


# get the manifold DAGMC geometry file
download(manifold_geom_url, 'dagmc.h5m') 
# get the manifold tet mesh
download(manifold_mesh_url, 'manifold.h5m')

"""
Explanation: Now let's download the geometry and mesh files.
(This may take some time.)
End of explanation
"""


src_pnt = openmc.stats.Point(xyz=(0.0, 0.0, 0.0))
src_energy = openmc.stats.Discrete(x=[5.e+06], p=[1.0])

source = openmc.Source(space=src_pnt, energy=src_energy)

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

settings.run_mode = "fixed source"
settings.batches = 10
settings.particles = 100

"""
Explanation: Next we'll create a 5 MeV neutron point source at the entrance the single pipe on the low side of the model with
End of explanation
"""


settings.dagmc = True

settings.export_to_xml()

"""
Explanation: And we'll indicate that we're using a CAD-based geometry.
End of explanation
"""


openmc.run()

"""
Explanation: We'll run a few particles through this geometry to make sure everything is working properly.
End of explanation
"""


unstructured_mesh = openmc.UnstructuredMesh("manifold.h5m")

mesh_filter = openmc.MeshFilter(unstructured_mesh)

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

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

settings.batches = 200
settings.particles = 5000
settings.export_to_xml()

openmc.run(output=False)

"""
Explanation: Now let's setup the unstructured mesh tally. We'll do this the same way we did in the previous notebook.
End of explanation
"""


!ls *.vtk

Image("./images/manifold_flux.png", width="800")

"""
Explanation: Again we should see that tally_1.200.vtk file which we can use to visualize our results in VisIt, ParaView, or another tool of your choice that supports VTK files.
End of explanation
"""


# energy filter with bins from 0 to 1 MeV and 1 MeV to 5 MeV
energy_filter = openmc.EnergyFilter((0.0, 1.e+06, 5.e+06))

tally.filters = [mesh_filter, energy_filter]
print(tally)
print(energy_filter)
tallies.export_to_xml()

!cat tallies.xml

openmc.run(output=False)

"""
Explanation: For the purpose of this example, we haven't run enough particles to score in all of the tet elements, but we indeed see larger flux values near the source location at the bottom of the model.
Visualization with statepoint data
It was mentioned in the previous unstructured mesh example that the centroids and volumes of elements are written to the state point file. Here, we'll explore how to use that information to produce point cloud information for visualization of this data.
This is particularly important when combining an unstructured mesh tally with other filters as a .vtk file will not automatically be written with the statepoint file in that scenario. To demonstrate this, let's setup a tally similar to the one above, but add an energy filter and re-run the model.
End of explanation
"""


with openmc.StatePoint("statepoint.200.h5") as sp:
    tally = sp.tallies[1]
    
    umesh = sp.meshes[1]
    centroids = umesh.centroids
    mesh_vols = umesh.volumes
    
    thermal_flux = tally.get_values(scores=['flux'], 
                                    filters=[openmc.EnergyFilter],
                                    filter_bins=[((0.0, 1.e+06),)])
    fast_flux = tally.get_values(scores=['flux'],
                                 filters=[openmc.EnergyFilter],
                                 filter_bins=[((1.e+06, 5.e+06),)])

data_dict = {'Flux 0 - 1 MeV' : thermal_flux,
             'Flux 1 - 5 MeV' : fast_flux,
             'Total Flux' : thermal_flux + fast_flux}

umesh.write_data_to_vtk("manifold", data_dict)

"""
Explanation: Noice the warning at the end of the output above indicating that the .vtk file we used before isn't written in this case.
Let's open up this statepoint file and get the information we need to create the point cloud data instead.
NOTE: You will need the Python vtk module installed to run this part of the notebook.
End of explanation
"""


!ls *.vtk

Image("./images/manifold_pnt_cld.png", width=800)

"""
Explanation: We should now see our new flux file in the directory. It can be used to visualize the results in the same way as our other .vtk files.
End of explanation
"""


