from pyne.xs.channels import sigma_t
from pyne.material import Material, from_atom_frac
from pyne.mesh import Mesh, IMeshTag, MetadataTag, ComputedTag
from yt.config import ytcfg; ytcfg["yt","suppressStreamLogging"] = "True"
from yt.visualization.plot_window import SlicePlot
from yt.frontends.moab.api import PyneMoabHex8StaticOutput

def cantor(n):
    return [0.] + cant(0., 1., n) + [1.]

def cant(x, y, n):
    if n == 0:
        return []
    new_pts = [2.*x/3. + y/3., x/3. + 2.*y/3.]
    return cant(x, new_pts[0], n-1) + new_pts + cant(new_pts[1], y, n-1)

#c5 = cantor(5)
c5x = [0.] + cant(0., 1500., 5) + [1500.]
c5y = [-850.] + cant(-850., 850., 5) + [850.]
coords = [c5x, c5y, [-850.0, 850.0]]

m = Mesh(structured_coords=coords, structured=True)
fuel = from_atom_frac({'U235': 0.045, 'U238': 0.955, 'O16': 2.0}, density=10.7)
cool = from_atom_frac({'H1': 2.0, 'O16': 1.0}, density=1.0)
for i in range(len(m)):
    m.mats[i] = cool
m.mats[len(m)/2] = fuel
m.mats[len(m)/4] = fuel
for i, c in enumerate(c5[:-1]):
    m.mats[i*len(c5)] = fuel

pf = PyneMoabHex8StaticOutput(m)
s = SlicePlot(pf, 'z', 'density', origin='native')
s.annotate_triangle_facets("bllite30matls_zip.h5m", 1, 0.0, plot_args={"colors": 'w'})
s.display()