%matplotlib inline
import matplotlib
matplotlib.rcParams['figure.figsize'] = (10.0, 16.0)
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import interp1d, InterpolatedUnivariateSpline
from scipy.optimize import bisect
import json
from functools import partial
class Foo: pass

"""
Explanation: Figure 1
Start by loading some boiler plate: matplotlib, numpy, scipy, json, functools, and a convenience class.
End of explanation
"""


from chest import Chest
from slict import CachedSlict
from glopen import glopen, glopen_many

"""
Explanation: And some more specialized dependencies:
 1. Slict provides a convenient slice-able dictionary interface
 2. Chest is an out-of-core dictionary that we'll hook directly to a globus remote using...
 3. glopen is an open-like context manager for remote globus files
End of explanation
"""


config = Foo()
config.name     = "HighAspect/HA_visc/HA_visc"
#config.arch_end = "alcf#dtn_mira/projects/alpha-nek"
config.arch_end = "maxhutch#alpha-admin/pub/"
config.frame = 1
config.lower = .25
config.upper = .75

"""
Explanation: Configuration for this figure.
End of explanation
"""


c = Chest(path      = "{:s}-results".format(config.name),
          open      = partial(glopen,      endpoint=config.arch_end),
          open_many = partial(glopen_many, endpoint=config.arch_end))
sc = CachedSlict(c)
with glopen(
            "{:s}.json".format(config.name), mode='r',
            endpoint = config.arch_end,
            ) as f:
    params = json.load(f)

"""
Explanation: Open a chest located on a remote globus endpoint and load a remote json configuration file.
End of explanation
"""


T = sc[:,'H'].keys()[config.frame]
frame = sc[T,:]
c.prefetch(frame.full_keys())

import yt
#test = frame['t_yz'] + 1.
test = np.tile(frame['t_yz'].transpose(),(1,1,1)).transpose() + 1.
data = dict(
    density = (test, "g/cm**3")
)
bbox = np.array([[params['root_mesh'][1], params['extent_mesh'][1]],
                 [params['root_mesh'][2], params['extent_mesh'][2]],
                 [0., 1.]])
#bbox = np.array([[params['root_mesh'][1], params['extent_mesh'][1]],
#                 [params['root_mesh'][2], params['extent_mesh'][2]]])
ds = yt.load_uniform_grid(data, test.shape, bbox=bbox, periodicity=(False,True,False), length_unit="m")

slc = yt.SlicePlot(ds, "z", "density",
                   width=(1,16))
slc.set_buff_size((14336,448))
#slc.pan((.25,7))
slc.show()

sl = ds.slice("z", 0).to_frb((1., 'm'), (128,128), height=(32.,'m'))
plt.imshow(sl['density'].d)
plt.show()

plt.figure()
plt.imshow(test[:,7000:7500,0].transpose())
plt.show()

"""
Explanation: We want to grab all the data for the selected frame.
End of explanation
"""


