import numpy as np
import xarray as xr
from scipy.io import netcdf_file
import matplotlib.pyplot as plt

# define a function to load data into xarray.DataArray
def load_data(filename, **kwargs):
    # load z and zi
    with netcdf_file(filename, 'r', mmap=False) as ncfile:
        nc_z  = ncfile.variables['z']
        nc_zi = ncfile.variables['zi']
        z = xr.DataArray(
                nc_z[0,:,0,0],
                dims=('z'),
                coords={'z': nc_z[0,:,0,0]},
                attrs={'long_name': nc_z.long_name.decode(), 'units': nc_z.units.decode()}
                )
        zi = xr.DataArray(
                nc_zi[0,:,0,0],
                dims=('zi'),
                coords={'zi': nc_zi[0,:,0,0]},
                attrs={'long_name': nc_zi.long_name.decode(), 'units': nc_zi.units.decode()}
                )
    # load other variables
    out = xr.load_dataset(
            filename,
            drop_variables=['z', 'zi'],
            **kwargs,
            )
    out = out.assign_coords({
        'z': z,
        'zi': zi,
        })
    out = out.assign_coords({
        'z_2d': (('time', 'z'), nc_z[:,:,0,0]),
        'zi_2d': (('time', 'zi'), nc_zi[:,:,0,0]),
        })
    for var in out.data_vars:
        if 'z' in out.data_vars[var].dims:
            out.data_vars[var].assign_coords({'z':z})
        elif 'zi' in out.data_vars[var].dims:
            out.data_vars[var].assign_coords({'zi':zi})
    # return a reorderd view
    return out.transpose('z', 'zi', 'time', 'lon', 'lat')


# read the data
filename1 = 'langmuir.nc'
filename2 = 'langmuir.cvmix.lf17.nc'
ds_gotm1 = load_data(filename1)
ds_gotm2 = load_data(filename2)


# plot surface forcing
fig, axarr = plt.subplots(4,sharex='col')
fig.set_size_inches([7,8])

ax = axarr[0]
hflux = ds_gotm1.data_vars['heat'] + ds_gotm1.data_vars['I_0']
hflux.plot(ax=ax, color='gray', linewidth=1)
hflux.rolling(time=40, center=True).mean().plot(
    ax=ax, color='k', linewidth=1.5)
ax.set_ylim([-300, 600])
ax.axhline(0.0, color='k', linestyle=':', linewidth=0.75)

ax = axarr[1]
rho_w = 1027.0
pme = ds_gotm1.data_vars['precip'] * rho_w * 1e6 # m/s -> mg/m^2/s
pme.plot(ax=ax, color='gray', linewidth=1)
pme.rolling(time=40, center=True).mean().plot(
    ax=ax, color='k', linewidth=1.5)
ax.set_ylim([-100, 150])
ax.axhline(0.0, color='k', linestyle=':', linewidth=0.75)

ax = axarr[2]
ustar = ds_gotm1.data_vars['u_taus']
ustar.plot(ax=ax, color='gray', linewidth=1)
ustar.rolling(time=40, center=True).mean().plot(
    ax=ax, color='k', linewidth=1.5)

ax = axarr[3]
laturb = ds_gotm2.data_vars['La_Turb']
laturb = laturb.where(laturb < 1.e3, drop=True)
laturb.plot(ax=ax, color='gray', linewidth=1)
laturb.rolling(time=40, center=True).mean().plot(
    ax=ax, color='k', linewidth=1.5)
ax.set_ylim([0.1,0.7])
ax.axhline(0.3, color='k', linestyle=':', linewidth=0.75)

ylabels = ['$Q_0$ [W m$^{-2}$]', '$F_0$ [mg m$^{-2}$ s$^{-1}$]', '$u^*$ [m s$^{-1}$]', 'La$_t$']

for i, ax in enumerate(axarr):
    ax.set_title('')
    ax.set_xlabel('')
    ax.set_ylabel(ylabels[i])
    ax.set_xlim([np.datetime64('2012-03-21'), np.datetime64('2013-03-21')])

plt.tight_layout()
fig.savefig('langmuir-forcing.png', dpi=300)


# plot temperature profile
fig, axarr = plt.subplots(3,sharex='col')
fig.set_size_inches([8,6])

ax = axarr[0]
levels = np.linspace(5,17,41)
ds_gotm1.data_vars['temp_obs'].plot(ax=ax, levels=levels, add_colorbar=False)

ax = axarr[1]
im =ds_gotm1.data_vars['temp'].plot(ax=ax, levels=levels, add_colorbar=False)

ax = axarr[2]
ds_gotm2.data_vars['temp'].plot(ax=ax, levels=levels, add_colorbar=False)

labels = ['(a) Observation', '(b) KC04', '(c) LF17']

for i, ax in enumerate(axarr):
    ax.set_title('')
    ax.set_xlabel('')
    ax.set_ylabel('$z$ [m]')
    ax.set_xlim([np.datetime64('2012-03-21'), np.datetime64('2013-03-21')])
    ax.text(0.03, 0.05, labels[i], transform=ax.transAxes, color='w', fontsize=11, va='bottom')

plt.subplots_adjust(top=0.97, bottom=0.1, right=0.86, left=0.1, hspace=0.1, wspace=0.1)
label_str = '$T$ [$^\circ$C]'
cax = plt.axes([0.8, 0.28, 0.1, 0.5])
cax.set_visible(False)
cb = plt.colorbar(im, ax=cax)
cb.formatter.set_powerlimits((-2, 2))
cb.update_ticks()
cb.set_label(label_str)

fig.savefig('langmuir.png', dpi=300)