#!/usr/bin/env python3

# Plots the mean function betas given an area key and a directory
# name. Note that this is specific to the OCO-2 example, modify as
# needed for other mean function types.

# Usage is something like:
# ./plot_mean_function_betas World ./dir/to/experiment

from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
rc('text', usetex=True)
from pylab import *

from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.basemap import Basemap
from areas import areas, get_latres

areadict = areas[sys.argv[1]]
dirname = sys.argv[2]

min_lat = areadict["latmin"]
min_lon = areadict["lonmin"]
max_lat = areadict["latmax"]
max_lon = areadict["lonmax"]
ar = (max_lon - min_lon)/(max_lat - min_lat)

nrows = 2
ncols = 2

fig, axs = subplots(nrows, ncols, figsize=((ncols*3*ar,nrows*3)))
# No frame in last plot at all, adapt if you would end up with
# non-full rows in the figure
# axs[-1, -1].axis('off')

fnames = ["beta0.txt", "beta1.txt", "beta2.txt", "delta.txt"]
titles = ["(a) $ \\beta_1$", "(b) $\\beta_2$", "(c) $\\beta_3$", "(d)  $\\delta$"]
for i, fname in enumerate(fnames):
    f = dirname + '/' + fname
    a = axs[i%nrows, i//nrows]
    a.set_title(titles[i], fontsize=18)
    print(a)
    b = Basemap(resolution='i', llcrnrlon=min_lon-180, llcrnrlat=min_lat,
                urcrnrlon=max_lon-180, urcrnrlat=max_lat, ax=a)
    b.drawcoastlines()
    data = loadtxt(f)
    l = nanpercentile(data, 5)
    h = nanpercentile(data, 99.5)
    # im = b.imshow(data, cmap='gist_rainbow', interpolation='bicubic', vmin=l, vmax=h)
    im = b.imshow(data, cmap='jet', interpolation='bicubic', vmin=l, vmax=h)
    divider = make_axes_locatable(a)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)

# Form of the OCO-2 mean function as a suptitle:
# fig.suptitle("$m(x,t; \\mathbf \\beta, \\delta) =  f(t)^{T} \\beta = \\beta_{0} \\sin\\left(\\frac{2\\pi t}{\\Delta_{t}} + \\delta\\right) +  \\beta_1 \\cos\\left(\\frac{4\\pi t}{\\Delta_{t}} + \\delta\\right) + \\beta_2 + Ct$", fontsize=20)

fig.tight_layout()
savefig(dirname + "/beta_fields.png", bbox_inches='tight', pad_inches=0.05, dpi=300)
