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

"""
Explanation: Comparison of bubble trajectories
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
"""


def load_from_archive(names, arch):
    cs = []
    for name in names:
        cs.append(Chest(path      = "{:s}-results".format(name),
                        open      = partial(glopen,      endpoint=arch),
                        open_many = partial(glopen_many, endpoint=arch)))
    scs = [CachedSlict(c) for c in cs]

    ps = []
    for name in names:
        with glopen(
                    "{:s}.json".format(name), mode='r',
                    endpoint = arch,
                    ) as f:
            ps.append(json.load(f))
    if len(names) == 1:
        return cs[0], scs[0], ps[0]
    return cs, scs, ps

"""
Explanation: Helper routines
End of explanation
"""


config = Foo()
config.names = [
    "Wilk/Wilk_long/Wilk_long",
#    "Wilk/Wilk_per/Wilk_per",
]
config.ref = ["Wilk/Wilk_per/Wilk_per",]
#config.arch_end = "maxhutch#alpha-admin/~/pub/"
#config.arch_end = "alcf#dtn_mira/projects/alpha-nek/experiments/"
config.arch_end = "alcf#dtn_mira/projects/PetaCESAR/maxhutch/"

"""
Explanation: Configuration for this figure.  config.ref points to a simulation with periodic boundaries.
End of explanation
"""


c, sc, p = load_from_archive(config.names, config.arch_end);
rc, rsc, rp = load_from_archive(config.ref, config.arch_end);

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


c.prefetch(sc[:,'H_exp'].full_keys())
c.prefetch(sc[:,'H_exp_cell'].full_keys())
rc.prefetch(rsc[:,'H_exp'].full_keys())
rc.prefetch(rsc[:,'H_exp_cell'].full_keys())

"""
Explanation: We want to plot the overall spike depth, which is the H_exp field in the chest, and the depth for individual spikes, H_exp_cell.
H_exp = max(H_exp_cell)
End of explanation
"""


spls = []
Hs = []
Ts = []

# reference
rT = np.array(rsc[:,'H_exp'].keys())
rH = np.array([x for x in rsc[:,'H_exp'].values()])/4
rspl = UnivariateSpline(rT, rH, k = 5, s = 1.e-10)
rFr = rspl.derivative()

# overall
T = np.array(sc[:,'H_exp'].keys())
H = np.array([x for x in sc[:,'H_exp'].values()])
spls.append(UnivariateSpline(T, H, k = 5, s = 1.e-10))
Hs.append(H)
Ts.append(T)

# bubble-wise
for i in range(10):
    T = np.array(sc[:,'H_exp'].keys())
    H = np.array([x[i] for x in sc[:,'H_exp_cell'].values()])
    spls.append(UnivariateSpline(T,
                                 H,
                                 k = 5,
                                 s = 1.e-10))
    Hs.append(H)
    Ts.append(T)
        
    
Frs = [spl.derivative() for spl in spls]    
T = np.linspace(sc[:,'H_exp'].keys()[0], sc[:,'H_exp'].keys()[-1], 1000)

"""
Explanation: Use a spline to compute the derivative of 'H' vs time: the Froude number.
First, we process the reference periodic simulation.  Then, the overall height
End of explanation
"""


fig, axs = plt.subplots(1,1)
for spl, Fr in zip(spls[1:], Frs[1:]):
    axs.plot(
        spl(T)* p["kmin"] , 
        Fr(T)/ np.sqrt(p["atwood"]*p["g"] / p["kmin"]),
        label="{:3.1f} modes".format(p["kmin"]));

axs.plot(
    spls[0](T)* p["kmin"], 
    Frs[0](T) / np.sqrt(p["atwood"]*p["g"] / p["kmin"]),
    'k-', label="Overall {:3.1f} modes".format(p["kmin"]) );
axs.plot(
    rspl(T)* rp["kmin"], 
    rFr(T) / np.sqrt(rp["atwood"]*rp["g"] / rp["kmin"]),
    'k--', label="Reference {:3.1f} modes".format(rp["kmin"]) );

axs.plot([0,3], [np.sqrt(1/np.pi), np.sqrt(1/np.pi)], 'k--')
axs.axvline(x=1.4, color='k');
axs.set_ylabel(r'Fr')
axs.set_xlabel(r'$h/\lambda$');
axs.legend(loc=0);
#axs.set_xbound(0,3)
#axs.set_ybound(0,2)
plt.savefig('Figure17.eps')

"""
Explanation: Plot the Froude number, non-dimensionalized by the theoretical dependence on Atwood, acceleration, and wave number, vs the spike depth, normalized by wave-length.
The horiztonal dotted line is the theoretical prediction of Goncharaov.  The vertical solid black line is the farthest that Wilkinson and Jacobs were able to get.
The solid black trajectory is the overall wall-bounded result (H_exp), while the dotted black trajectory is the periodic reference case.
End of explanation
"""


fig, axs = plt.subplots(1,1)
for spl, Fr in zip(spls[1:], Frs[1:]):
    axs.plot(
        spl(T)* p["kmin"] , 
        Fr(T)/ np.sqrt(p["atwood"]*p["g"] / p["kmin"]),
        label="{:3.1f} modes".format(p["kmin"]));

axs.plot(
    spls[0](T)* p["kmin"], 
    Frs[0](T) / np.sqrt(p["atwood"]*p["g"] / p["kmin"]),
    'k-', label="Overall {:3.1f} modes".format(p["kmin"]) );
axs.plot(
    rspl(T)* rp["kmin"], 
    rFr(T) / np.sqrt(rp["atwood"]*rp["g"] / rp["kmin"]),
    'k--', label="Reference {:3.1f} modes".format(rp["kmin"]) );

axs.plot([0,3], [np.sqrt(1/np.pi), np.sqrt(1/np.pi)], 'k--')
axs.axvline(x=1.4, color='k');
axs.set_ylabel(r'Fr')
axs.set_xlabel(r'$h/\lambda$');
axs.legend(loc=0);
axs.set_xbound(.5,1)
axs.set_ybound(.5,.6)
plt.savefig('Figure17.eps')

fig, axs = plt.subplots(1,1)

for H, T in zip(Hs[1:], Ts[1:]):
    axs.plot(
        T,
        H, 'x', 
        label="{:3.1f} modes".format(p["kmin"]));
axs.plot(
        Ts[0],
        Hs[0], 'k-', 
        label="Overall {:3.1f} modes".format(p["kmin"]));
axs.plot(
        rT,
        rH, 'k--', 
        label="Reference {:3.1f} modes".format(p["kmin"]));
axs.set_ylabel(r'$h/\lambda$')
axs.set_xlabel(r'T (s)');
axs.legend(loc=0);

%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: Let's zoom into the stagnation stage.
End of explanation
"""


