%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_conductivity_1.0E-5/HA_conductivity_1.0E-5"
config.arch_end = "alcf#dtn_mira/projects/alpha-nek/experiments/"
config.frame = 100
config.lower = .1
config.upper = .9

"""
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())

fig = plt.figure()

nx = frame['t_yz'].shape[0]
ny = frame['t_yz'].shape[1]

full = np.concatenate((frame['t_yz'], np.flipud(frame['t_yz'])), axis=0)

plt.imshow(full[:,int(ny*config.lower):int(ny*config.upper)].transpose(),
           origin='lower')
plt.colorbar();

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


fig = plt.figure()

nx = frame['t_yz'].shape[0]
ny = frame['t_yz'].shape[1]

plt.imshow(frame['fz_yz'][:,int(ny*config.lower):int(ny*config.upper)].transpose(),
           vmin = -params['atwood']*params['g']*2,
           vmax =  params['atwood']*params['g']*2,
           origin='lower')
plt.colorbar();
plt.contour(frame['fz_yz'][:,int(ny*config.lower):int(ny*config.upper)].transpose(), 
            colors='k', 
            levels=[-params['atwood']*params['g'], params['atwood']*params['g']])
#axs[1].imshow(frame['fz_yz'].transpose(), origin='lower')

def laplace(grid):
    from numpy.fft import fftn, ifftn, fftfreq
    kx = fftfreq(grid.shape[0])
    ky = fftfreq(grid.shape[1])
    kgrid = fftn(grid)
    for i in range(kgrid.shape[0]):
        for j in range(kgrid.shape[1]):
            if kx[i] != 0 or ky[j] != 0:
                kgrid[i,j] = kgrid[i,j] / (kx[i]**2 + ky[j]**2)
    kgrid[0,0] = 0.
    rgrid = np.array(ifftn(kgrid), dtype=float)
    return rgrid

fig = plt.figure()

nx = frame['t_yz'].shape[0]
ny = frame['t_yz'].shape[1]
full = np.concatenate((frame['vorticity_yz'], -np.flipud(frame['vorticity_yz'])), axis=0)
plt.imshow(full[:,int(ny*config.lower):int(ny*config.upper)].transpose(),
           origin='lower')
#axs[1].imshow(frame['fz_yz'].transpose(), origin='lower')
plt.colorbar();

full = np.concatenate((frame['vorticity_yz'], -np.flipud(frame['vorticity_yz'])), axis=0)
streamfunction = laplace(full[:,:]);

fig = plt.figure()

nx = frame['t_yz'].shape[0]
ny = frame['t_yz'].shape[1]
print(np.max(streamfunction), np.min(streamfunction))
full = np.concatenate((frame['t_yz'], np.flipud(frame['t_yz'])), axis=0)

#background = np.tile(frame['t_yz'][:,int(ny*config.lower):int(ny*config.upper)].transpose(), (1,4))
background = np.tile(streamfunction[:,int(ny*config.lower):int(ny*config.upper)].transpose(), (1,4))
foreground = np.tile(streamfunction[:,int(ny*config.lower):int(ny*config.upper)].transpose(), (1,4))

plt.imshow(background,
           origin='lower', alpha=0.1)
plt.colorbar();
plt.contour(foreground,
           origin='lower', colors='k',
           #levels=[-100,0,100]
           )
#axs[1].imshow(frame['fz_yz'].transpose(), origin='lower')


fig = plt.figure()

nx = frame['t_yz'].shape[0]
ny = frame['t_yz'].shape[1]
full = np.concatenate((frame['total_pressure_yz'], np.flipud(frame['total_pressure_yz'])), axis=0)
plt.imshow(full[:,int(ny*config.lower):int(ny*config.upper)].transpose(),
           origin='lower')
#axs[1].imshow(frame['fz_yz'].transpose(), origin='lower')
plt.colorbar();

%install_ext http://raw.github.com/jrjohansson/version_information/master/version_information.py 
%load_ext version_information 
%version_information numpy, matplotlib, slict, chest, glopen, globussh

"""
Explanation: Plot the bubble height, the 'H' keys, vs. time.
Use a spline to compute the derivative.
End of explanation
"""


