from itertools import product 
from pyne.mesh import Mesh
from pyne.xs.cache import XSCache
from pyne.xs.data_source import CinderDataSource, SimpleDataSource, NullDataSource
from pyne.xs.channels import sigma_a, sigma_s
from pyne.material import Material, from_atom_frac
import numpy as np
from yt.config import ytcfg; ytcfg["yt","suppressStreamLogging"] = "True"
from yt.mods import *
from itaps import iBase, iMesh
from matplotlib import animation
from JSAnimation import IPython_display
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from IPython.display import HTML

xsc = XSCache([0.026e-6, 0.024e-6], (SimpleDataSource, NullDataSource))

def lpoint(idx, n, coords, shape, m):
    lidx = list(idx)
    lidx[n] += 1 if idx[n] == 0 else -1
    left = m.structured_get_hex(*lidx)
    l = m.mesh.getVtxCoords(left)[n]
    if idx[n] == 0:
        l = 2*coords[n] - l 
    return left, l

def rpoint(idx, n, coords, shape, m):
    ridx = list(idx)
    ridx[n] += -1 if idx[n] == shape[n]-2 else 1
    right = m.structured_get_hex(*ridx)
    r = m.mesh.getVtxCoords(right)[n]
    if idx[n] == shape[n]-2:
        r = 2*coords[n] - r
    return right, r

def laplace(tag, idx, m, shape):
    ent = m.structured_get_hex(*idx)
    coords = m.mesh.getVtxCoords(ent)
    lptag = 0.0
    for n in range(3):
        left, l = lpoint(idx, n, coords, shape, m)
        right, r = rpoint(idx, n, coords, shape, m)
        c = coords[n]
        lptag += (((tag[right] - tag[ent])/(r-c)) - ((tag[ent] - tag[left])/(c-l))) / ((r-l)/2)
    return lptag

def timestep(m, dt):
    nx = len(m.structured_get_divisions("x"))
    ny = len(m.structured_get_divisions("y"))
    nz = len(m.structured_get_divisions("z"))
    shape = (nx, ny, nz)
    D = m.mesh.getTagHandle("D")
    k = m.mesh.getTagHandle("k")
    S = m.mesh.getTagHandle("S")
    Sigma_a = m.mesh.getTagHandle("Sigma_a")
    phi = m.mesh.getTagHandle("phi")
    phi_next = m.mesh.getTagHandle("phi_next")
    for idx in product(*[range(xyz-1) for xyz in shape]):
        ent = m.structured_get_hex(*idx)
        phi_next[ent] = (max(D[ent] * laplace(phi, idx, m, shape) + 
                                    (k[ent] - 1.0) * Sigma_a[ent] * phi[ent], 0.0) + S[ent])*dt*2.2e5 + phi[ent]
    ents = m.mesh.getEntities(iBase.Type.region)
    phi[ents] = phi_next[ents]

def render(m, dt, axis="z", field="phi", frames=100):
    pf = PyneMoabHex8StaticOutput(m)
    s = SlicePlot(pf, axis, field, origin='native')
    fig = s.plots['gas', field].figure
    fig.canvas = FigureCanvasAgg(fig)
    axim = fig.axes[0].images[0]

    def init():
        axim = s.plots['gas', 'phi'].image
        return axim

    def animate(i):
        s = SlicePlot(pf, axis, field, origin='native')
        axim.set_data(s._frb['gas', field])
        timestep(m, dt)
        return axim

    return animation.FuncAnimation(fig, animate, init_func=init, frames=frames, interval=100, blit=False)

def isinrod(ent, rx, radius=0.4):
    """returns whether an entity is in a control rod"""
    coord = rx.mesh.getVtxCoords(ent)
    return (coord[0]**2 + coord[1]**2) <= radius**2

def create_reactor(multfact=1.0, radius=0.4):
    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)
    xpoints = [0.0, 0.075, 0.15, 0.225] + list(np.arange(0.3, 0.7, 0.025)) + list(np.arange(0.7, 1.201, 0.05))
    ypoints = xpoints
    zpoints = np.linspace(0.0, 1.0, 8)
    rx = Mesh(structured_coords=[xpoints, ypoints, zpoints], structured=True)
    D = rx.mesh.createTag("D", 1, float)
    k = rx.mesh.createTag("k", 1, float)
    S = rx.mesh.createTag('S', 1, float)
    Sigma_a = rx.mesh.createTag("Sigma_a", 1, float)
    phi = rx.mesh.createTag("phi", 1, float)
    phi_next = rx.mesh.createTag("phi_next", 1, float)
    for ent in rx.structured_iterate_hex("xyz"):
        if isinrod(ent, rx, radius=radius):
            D[ent] = 1.0 / (3.0 * fuel.density * 1e-24 * sigma_s(fuel, xs_cache=xsc))
            Sigma_a[ent] = fuel.density * 1e-24 * sigma_a(fuel, xs_cache=xsc)
            phi[ent] = 4e14
            k[ent] = multfact
        else:
            D[ent] = 1.0 / (3.0 * cool.density * 1e-24 * sigma_s(cool, xs_cache=xsc))
            Sigma_a[ent] = cool.density * 1e-24 * sigma_a(cool, xs_cache=xsc)
            r2 = (rx.mesh.getVtxCoords(ent)[:2]**2).sum()
            phi[ent] = 4e14 * radius**2 / r2 if r2 < 0.7**2 else 0.0
            k[ent] = 0.0
        S[ent] = 0.0
        phi_next[ent] = 0.0
    return rx

rx = create_reactor()
render(rx, dt=2.5e-31, frames=100)