#!/usr/bin/env python
# coding: utf-8

# In[1]:


from __future__ import division

import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np

from salishsea_tools import nc_tools

get_ipython().run_line_magic('matplotlib', 'inline')


# In[7]:


restartU = nc.Dataset('/data/jieliu/MEOPAR/river-treatment/15days_nowcast_restart14/1hnowcastrestart06150629gridU.nc')
restartV = nc.Dataset('/data/jieliu/MEOPAR/river-treatment/15days_nowcast_restart14/1hnowcastrestart06150629gridV.nc')
coldU = nc.Dataset('/data/jieliu/MEOPAR/river-treatment/15days_nowcast_coldstart14/1hnowcastallthesame06150629gridU.nc')
coldV = nc.Dataset('/data/jieliu/MEOPAR/river-treatment/15days_nowcast_coldstart14/1hnowcastallthesame06150629gridV.nc')


# In[20]:


rms = np.zeros(360); abse = np.zeros_like(rms); mvel = np.zeros_like(rms); maxerr = np.zeros_like(rms)
for i in range(360):
    rU = restartU.variables['vozocrtx'][i]
    cU = coldU.variables['vozocrtx'][i]
    rV = restartV.variables['vomecrty'][i]
    cV = coldV.variables['vomecrty'][i]
    difference = np.mean((rU-cU)**2 + (rV-cV)**2)
    maxerr[i] = np.max((rU-cU)**2 + (rV-cV)**2)
    ref = np.mean(rU**2 + rV**2 + cU**2 + cV**2)
    rms[i] = np.sqrt(difference/ref)
    abse[i] =np.sqrt(difference)
    mvel[i] = np.sqrt(ref)
    print i


# In[32]:


fig, ax = plt.subplots(2, 2, figsize=(10,10))
ax[0,0].plot(rms)
ax[0,1].plot(mvel)
ax[1,0].plot(abse)
ax[1,1].plot(maxerr)
ax[1,0].set_xlabel('Time (hours)')
ax[1,1].set_xlabel('Time (hours)')
ax[0,0].set_ylabel('Normalized RMS error')
ax[0,1].set_ylabel('Mean Velocities (m/s)')
ax[1,0].set_ylabel('Root Mean Square error (m/s)')


# In[50]:


fig, ax = plt.subplots(1,2,figsize=(12,5))
mesh = ax[0].pcolormesh(rU[0]-cU[0], cmap='bwr', vmax=0.5, vmin=-0.5)
ax[0].set_title('Difference in u')
fig.colorbar(mesh, ax=ax[0])
ax[1].set_title('Difference in v')
mesh = ax[1].pcolormesh(rV[0]-cV[0], cmap='bwr', vmax=0.5, vmin=-0.5)
fig.colorbar(mesh, ax=ax[1])


# In[51]:


fig, ax = plt.subplots(2,2,figsize=(12,10))
mesh = ax[0,0].pcolormesh(rU[0], cmap='bwr', vmax=2, vmin=-2)
fig.colorbar(mesh, ax=ax[0,0])
mesh = ax[0,1].pcolormesh(rV[0], cmap='bwr', vmax=2, vmin=-2)
fig.colorbar(mesh, ax=ax[0,1])
mesh = ax[1,0].pcolormesh(cU[0], cmap='bwr', vmax=2, vmin=-2)
fig.colorbar(mesh, ax=ax[1,0])
mesh = ax[1,1].pcolormesh(cV[0], cmap='bwr', vmax=2, vmin=-2)
fig.colorbar(mesh, ax=ax[1,1])
ax[0,0].set_title('Restart u')
ax[0,1].set_title('Restart v')
ax[1,0].set_title('Cold u')
ax[1,1].set_title('Cold v')


# In[ ]: