from __future__ import division

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

from salishsea_tools import (
    nc_tools,
    viz_tools,
    stormtools,
    tidetools,
)
import pytz, datetime
import glob
import os
import urllib2
import csv
import cStringIO
import requests
from xml.etree import cElementTree as ElementTree
import pandas as pd
import arrow
from matplotlib import dates


%matplotlib inline




grid = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')
NEMOlat=grid.variables['nav_lat'][:]
NEMOlon=grid.variables['nav_lon'][:]
bathy=grid.variables['Bathymetry'][:]

GEM = nc.Dataset('/ocean/dlatorne/MEOPAR/GEM2.5/NEMO-atmos/res_y2014m09d10.nc')
GEMlon = GEM.variables['nav_lon'][:]-360
GEMlat = GEM.variables['nav_lat'][:]

nc_tools.show_variable_attrs(GEM)

OP = nc.Dataset('/ocean/sallen/allen/research/Meopar/Operational/oper_allvar_y2014m09d12.nc')
OPlon =OP.variables['nav_lon'][:]-360
OPlat = OP.variables['nav_lat'][:]
nc_tools.show_variables(OP)

nc_tools.show_variable_attrs(OP)

wind_speed = {}; wind_dir = {}; time = {}; lat={}; lon={}; press={}; temp={}; therm={}; solar={}; precip={}; qair={}


start = '01-Sep-2014'; end = '30-Sep-2014';
start2 = '01-Oct-2014'; end2 = '31-Oct-2014';

stations = ['Sandheads','EntranceIsland', 'PamRocks', 'YVR', 'SistersIsland', 'Esquimalt']

for key in stations:
    [wind_speed[key],wind_dir[key],temp[key],time[key], 
     lat[key], lon[key]] = stormtools.get_EC_observations(key,start,end)
    [ws,wd,T,t, la, lo] = stormtools.get_EC_observations(key,start2,end2)
    wind_speed[key]=np.append(wind_speed[key],ws)
    wind_dir[key]=np.append(wind_dir[key],wd)
    time[key]=np.append(time[key],t)
    temp[key]=np.append(temp[key],T)

def find_model_point(lon,lat,X,Y):
        # Tolerance for searching for grid points
    # (approx. distances between adjacent grid points)
    tol1 = 0.015 # lon
    tol2 = 0.015# lat

    # Search for a grid point with lon/lat within tolerance of
    # measured location
    x1, y1 = np.where(
        np.logical_and(
            (np.logical_and(X > lon-tol1, X < lon+tol1)),
            (np.logical_and(Y > lat-tol2, Y < lat+tol2))))
    return x1[0], y1[0]

filesGEM = glob.glob('/ocean/dlatorne/MEOPAR/GEM2.5/NEMO-atmos/res_y2014m*.nc')
filesGEM.sort(key=os.path.basename)


filesOP = glob.glob('/ocean/sallen/allen/research/Meopar/Operational/oper_allvar_y2014m*.nc')
filesOP.sort(key=os.path.basename)

def compile_GEM(j,i):
    wind=[]; direc=[]; t=[]; pr=[]; sol=[]; the=[]; pre=[]; tem=[]; qr=[];
    for f in filesGEM:
        G = nc.Dataset(f)
        u = G.variables['u_wind'][:,j,i]; v=G.variables['v_wind'][:,j,i];
        pr.append(G.variables['atmpres'][:,j,i]); sol.append(G.variables['solar'][:,j,i]); 
        qr.append(G.variables['qair'][:,j,i]); the.append(G.variables['therm_rad'][:,j,i]); 
        pre.append(G.variables['precip'][:,j,i]);
        tem.append(G.variables['tair'][:,j,i])
        speed = np.sqrt(u**2 + v**2)
        wind.append(speed)
        
        d = np.arctan2(v, u)
        d = np.rad2deg(d + (d<0)*2*np.pi);
        direc.append(d)
        
        ts=G.variables['time_counter']
        torig = nc_tools.time_origin(G)
        for ind in np.arange(ts.shape[0]):
            t.append((torig + datetime.timedelta(seconds=ts[ind])).datetime)
        
    wind = np.array(wind).reshape(len(filesGEM)*24,)
    direc = np.array(direc,'double').reshape(len(filesGEM)*24,)
    t = np.array(t).reshape(len(filesGEM)*24,)
    pr = np.array(pr).reshape(len(filesGEM)*24,)
    tem = np.array(tem).reshape(len(filesGEM)*24,)
    sol = np.array(sol).reshape(len(filesGEM)*24,)
    the = np.array(the).reshape(len(filesGEM)*24,)
    qr = np.array(qr).reshape(len(filesGEM)*24,)
    pre = np.array(pre).reshape(len(filesGEM)*24,)
    
    return wind, direc, t, pr, tem, sol, the, qr, pre

def compile_OP(j,i):
    wind=[]; direc=[]; t=[]; pr=[]; sol=[]; the=[]; pre=[]; tem=[]; qr=[];
    for f in filesOP:
        G = nc.Dataset(f)
        u = G.variables['u_wind'][:,j,i]; v=G.variables['v_wind'][:,j,i];
        pr.append(G.variables['atmpres'][:,j,i]); sol.append(G.variables['solar'][:,j,i]); 
        qr.append(G.variables['qair'][:,j,i]); the.append(G.variables['therm_rad'][:,j,i]); 
        pre.append(G.variables['precip'][:,j,i]); tem.append(G.variables['tair'][:,j,i])
        speed = np.sqrt(u**2 + v**2)
        wind.append(speed)
        
        d = np.arctan2(v, u)
        d = np.rad2deg(d + (d<0)*2*np.pi);
        direc.append(d)
        
        ts=G.variables['time_counter']
        torig = datetime.datetime(1970,1,1) #there is no time_origin attriubte in OP files, so I hard coded this
        for ind in np.arange(ts.shape[0]):
            t.append((torig + datetime.timedelta(seconds=ts[ind])))
        
    wind = np.array(wind).reshape(len(filesOP)*24,)
    direc = np.array(direc,'double').reshape(len(filesOP)*24,)
    t = np.array(t).reshape(len(filesOP)*24,)
    pr= np.array(pr).reshape(len(filesOP)*24,)
    tem = np.array(tem).reshape(len(filesOP)*24,)
    sol = np.array(sol).reshape(len(filesOP)*24,)
    the = np.array(the).reshape(len(filesOP)*24,)
    qr = np.array(qr).reshape(len(filesOP)*24,)
    pre = np.array(pre).reshape(len(filesOP)*24,)
    
    return wind, direc, t, pr, tem, sol, the, qr, pre

stationsGEM =['Sandheads_GEM', 'EntranceIsland_GEM','PamRocks_GEM','YVR_GEM',
                     'SistersIsland_GEM','Esquimalt_GEM']

stationsOP =['Sandheads_OP', 'EntranceIsland_OP','PamRocks_OP','YVR_OP',
                     'SistersIsland_OP','Esquimalt_OP']

for (obs, modGEM, modOP) in zip(stations,stationsGEM, stationsOP):
    [j,i]=find_model_point(lon[obs],lat[obs],GEMlon,GEMlat)
    lon[modGEM] = GEMlon[j,i]
    lat[modGEM]=GEMlat[j,i]

    [wind_speed[modGEM],wind_dir[modGEM],time[modGEM],
     press[modGEM],temp[modGEM],solar[modGEM],
     therm[modGEM],qair[modGEM],precip[modGEM]] = compile_GEM(j,i)
    
    [j,i]=find_model_point(lon[obs],lat[obs],OPlon,OPlat)
    lon[modOP] = OPlon[j,i]
    lat[modOP]=OPlat[j,i]

    [wind_speed[modOP],wind_dir[modOP],time[modOP],
     press[modOP],temp[modOP],solar[modOP],
     therm[modOP],qair[modOP],precip[modOP]] = compile_OP(j,i)

key='Buoy'
lat[key]=47.353 
lon[key]=-124.731

[j,i]=find_model_point(lon[key],lat[key],GEMlon,GEMlat)
key='Buoy_GEM'
lon[key] = GEMlon[j,i]
lat[key]=GEMlat[j,i]

[wind_speed[key],wind_dir[key],time[key],
 press[key],temp[key],solar[key],
therm[key],qair[key],precip[key]]= compile_GEM(j,i)

[j,i]=find_model_point(lon[key],lat[key],OPlon,OPlat)
key='Buoy_OP'
lon[key] = OPlon[j,i]
lat[key]=OPlat[j,i]

[wind_speed[key],wind_dir[key],time[key],
press[key],temp[key],solar[key],
therm[key],qair[key],precip[key]]= compile_OP(j,i)

response = urllib2.urlopen('http://www.ndbc.noaa.gov/data/realtime2/46041.cwind')
html = response.read()
text_file = open("Output.txt", "w")
text_file.write(html)
text_file.close()

response = urllib2.urlopen('http://www.ndbc.noaa.gov/data/derived2/46041.dmv')
html = response.read()
text_file = open("Output_derived.txt", "w")
text_file.write(html)
text_file.close()

#read in csv
key='Buoy'
fil= csv.reader(open('Output.txt','rb'),delimiter=' ')
fil.next()
fil.next()
time[key]=[]; wind_speed[key]=[]; wind_dir[key]=[]
for row in fil:
    aList=[]
    for element in row:
        if element != '':
            aList.append(element)
    time[key].append(datetime.datetime(int(aList[0]),int(aList[1]),int(aList[2]),int(aList[3]),int(aList[4])))
    d = -float(aList[5])+270
    d=d+ 360 * (d<0)
    wind_dir[key].append(d)
    wind_speed[key].append(float(aList[6]))
time[key]=np.array(time[key])
wind_dir[key]=np.array(wind_dir[key])
wind_speed[key]=np.array(wind_speed[key])
# and the 10 m derived winds
key='Buoy10m'
fil= csv.reader(open('Output_derived.txt','rb'),delimiter=' ')
fil.next()
fil.next()
time[key]=[]; wind_speed[key]=[]; wind_dir[key]=[]
for row in fil:
    aList=[]
    for element in row:
        if element != '':
            aList.append(element)
    time[key].append(datetime.datetime(int(aList[0]),int(aList[1]),int(aList[2]),int(aList[3]),
                                       int(aList[4])))
    wind_speed[key].append(float(aList[9]))
time[key]=np.array(time[key])
wind_speed[key]=np.array(wind_speed[key])

def compare_winds(key1,key2,key3,sax,eax):
    #compare wind speed and direction for data indicated by key1 and key2
    #time limits on axis given by sax,eax
    fig,axs = plt.subplots(2,1,figsize=(20,8))

    for key in [key1,key2,key3]:
        ax=axs[0]
        ax.plot(time[key],wind_speed[key],label=key)
        ax.set_title('Wind speed')
        ax.set_xlim([sax,eax])
        ax.legend(loc=0)
        ax.set_ylabel('wind speed (m/s)')
    
        ax=axs[1]
        ax.plot(time[key],wind_dir[key],label=key)
        ax.set_title('Wind direction')
        ax.set_xlim([sax,eax])
        ax.legend(loc=0)
        ax.set_ylabel('wind direction (degrees CCW from East)')
    
    fig,ax = plt.subplots(1,1,figsize=(5,6))
    for key in [key1,key2,key3]:
        ax.plot(lon[key],lat[key],'o',label=key)
        viz_tools.plot_coastline(ax,grid,coords='map')
        ax.legend(loc=0)
        
    return ax

def compare_other_fields(key1,key2,sax,eax):
    fig,axs = plt.subplots(3,2,figsize=(20,10))
    axs=axs.flatten()
    colors={key1: 'g',key2:'r','obs':'b'}
    
    for key in [key1,key2]:
        ax=axs[0]
        ax.plot(time[key],press[key],colors[key],label=key)
        ax.set_title('Sea level pressure')
        ax.set_xlim([sax,eax]); ax.set_ylim([99000, 103000])
        ax.legend(loc=0);
        ax.set_ylabel('Sea level pressure (Pa)')
        ax.legend(loc=0);
    
        ax=axs[1]
        ax.plot(time[key],temp[key],colors[key],label=key)
        ax.set_title('Temperature 2m above ground')
        ax.set_xlim([sax,eax]); ax.set_ylim([270, 300])
        ax.set_ylabel('Temperature (K)')
        
        ax=axs[2]
        ax.plot(time[key],solar[key],colors[key],label=key)
        ax.set_title('Short wave radiation')
        ax.set_xlim([sax,eax]); ax.set_ylim([0,800])
        ax.legend(loc=0)
        ax.set_ylabel('Short wave radiaiton(W/m^2)')
    
        ax=axs[3]
        ax.plot(time[key],therm[key],colors[key],label=key)
        ax.set_title('Thermal radiation')
        ax.set_xlim([sax,eax]); ax.set_ylim([200, 400])
        ax.legend(loc=0)
        ax.set_ylabel('Thermal radiaiton (W/m^2)')
        
        ax=axs[4]
        ax.plot(time[key],qair[key],colors[key],label=key)
        ax.set_title('Specific Humidty')
        ax.set_xlim([sax,eax]); ax.set_ylim([0, 0.012])
        ax.legend(loc=0)
        ax.set_ylabel('Specfific Humidity (kg/kg)')
    
        ax=axs[5]
        ax.plot(time[key],precip[key],colors[key],label=key)
        ax.set_title('Precipitaion')
        ax.set_xlim([sax,eax]); ax.set_ylim([0, 0.02])
        ax.legend(loc=0)
        ax.set_ylabel('Precipitation (kg/m^2)')
        
    key=key2[:-3]
    if key != 'Buoy':
        axs[1].plot(time[key],temp[key],colors['obs'],label='observations')
        axs[1].legend(loc=0)
        
    return axs

sax=datetime.datetime(2014,9,10)
eax=datetime.datetime(2014,10,15)

ax =compare_winds('Sandheads','Sandheads_GEM','Sandheads_OP',sax,eax)
ax.set_xlim([-124,-123])
ax.set_ylim([49,50])
ax =compare_other_fields('Sandheads_GEM','Sandheads_OP',sax,eax)


ax=compare_winds('EntranceIsland','EntranceIsland_GEM','EntranceIsland_OP',sax,eax)
ax.set_xlim([-124,-123])
ax.set_ylim([49,50])
ax =compare_other_fields('EntranceIsland_GEM','EntranceIsland_OP',sax,eax)

ax=compare_winds('PamRocks','PamRocks_GEM','PamRocks_OP',sax,eax)
ax.set_xlim([-124,-123])
ax.set_ylim([49,50])
ax =compare_other_fields('PamRocks_GEM','PamRocks_OP',sax,eax)

ax=compare_winds('YVR','YVR_GEM','YVR_OP',sax,eax)
ax.set_xlim([-124,-123])
ax.set_ylim([49,50])
ax =compare_other_fields('YVR_GEM','YVR_OP',sax,eax)

ax=compare_winds('SistersIsland','SistersIsland_GEM','SistersIsland_OP',sax,eax)
ax.set_xlim([-124.5,-123])
ax.set_ylim([49,50])
ax =compare_other_fields('SistersIsland_GEM','SistersIsland_OP',sax,eax)

ax=compare_winds('Esquimalt','Esquimalt_GEM','Esquimalt_OP',sax,eax)
ax.set_xlim([-124.5,-123])
ax.set_ylim([48,50])
ax =compare_other_fields('Esquimalt_GEM','Esquimalt_OP',sax,eax)

fig,axs = plt.subplots(2,1,figsize=(20,8))

for key in ['Buoy10m','Buoy_GEM', 'Buoy_OP']:
    ax=axs[0]
    ax.plot(time[key],wind_speed[key],label=key)
    ax.set_title('Wind speed')
    ax.set_xlim(sax,eax)
    ax.legend(loc=0)
    ax.set_ylabel('wind speed (m/s)')
    
for key in ['Buoy','Buoy_GEM','Buoy_OP']:
    ax=axs[1]
    ax.plot(time[key],wind_dir[key],label=key)
    ax.set_title('Wind direction')
    ax.set_xlim(sax,eax)
    ax.legend(loc=0)
    ax.set_ylabel('wind direction (degrees CCW from East)')
    
fig,ax = plt.subplots(1,1,figsize=(5,6))
for key in ['Buoy','Buoy_GEM','Buoy_OP']:
    ax.plot(lon[key],lat[key],'o',label=key)
    viz_tools.plot_coastline(ax,grid,coords='map')
    ax.legend(loc=0)
    
ax =compare_other_fields('Buoy_GEM','Buoy_OP',sax,eax)

filesU = glob.glob('/data//dlatorne/MEOPAR/SalishSea/results/gem-res-1*/SalishSea_1h_*grid_U.nc')
filesU.sort(key=os.path.basename)

filesV = glob.glob('/data//dlatorne/MEOPAR/SalishSea/results/gem-res-1*/SalishSea_1h_*grid_V.nc')
filesV.sort(key=os.path.basename)

def compile_NEMO(j,i):
    wind=[]; direc=[]; t=[]
    for fU, fV in zip(filesU,filesV):
        NU = nc.Dataset(fU); NV=nc.Dataset(fV)
        u = NU.variables['u_wind_stress'][:,j,i]; v=NV.variables['v_wind_stress'][:,j,i];
        speed = np.sqrt(u**2 + v**2)
        wind = np.append(wind,speed)
        
        d = np.arctan2(v, u)
        d = np.rad2deg(d + (d<0)*2*np.pi);
        direc=np.append(direc,d)
        
        ts=NU.variables['time_counter']
        t=np.append(t,nc_tools.timestamp(NU,np.arange(ts.shape[0])))
        
    return wind, direc, t

key='Sandheads'
[j,i]=tidetools.find_closest_model_point(lon[key],lat[key],NEMOlon,NEMOlat,bathy)

key='Sandheads_NEMO'
lon[key] = NEMOlon[j,i]
lat[key]=NEMOlat[j,i]

[wind_speed[key],wind_dir[key],time[key]] = compile_NEMO(j,i)

fig,axs = plt.subplots(1,1,figsize=(15,6))
s='09-Sep-2014'; e='24-Sep-2014';
unaware=datetime.datetime.strptime(s,"%d-%b-%Y")
sdt = unaware.replace(tzinfo=pytz.timezone('utc'))
unaware=datetime.datetime.strptime(e,"%d-%b-%Y")
edt = unaware.replace(tzinfo=pytz.timezone('utc'))


for key in ['Sandheads','Sandheads_GEM', 'Sandheads_NEMO']:
    ax=axs
    ax.plot(time[key],wind_dir[key],label=key)
    ax.set_title('Wind and wind stress direction - Sandheads_NEMO is wind stress direction used by NEMO')
    ax.set_xlim([sdt,edt])
    ax.legend(loc=0)
    ax.set_ylabel('direction (degrees CCW from East)')
    
fig,ax = plt.subplots(1,1,figsize=(5,6))
for key in ['Sandheads','Sandheads_GEM', 'Sandheads_NEMO']:
    ax.plot(lon[key],lat[key],'o',label=key)
    viz_tools.plot_coastline(ax,grid,coords='map')
    ax.legend(loc=0)
ax.set_xlim([-124,-123])
ax.set_ylim([49,50])

pacific=pytz.timezone('US/Pacific')
[j,i]=find_model_point(lon['Sandheads'],lat['Sandheads'],GEMlon,GEMlat)

#October 13
GEM = nc.Dataset('/ocean/dlatorne/MEOPAR/GEM2.5/NEMO-atmos/res_y2014m10d13.nc')
GEM_precip = GEM.variables['precip'][:,j,i]
ts=GEM.variables['time_counter']
torig = nc_tools.time_origin(GEM)
GEM_time=[]
for ind in np.arange(ts.shape[0]):
        GEM_time.append(((torig + datetime.timedelta(seconds=ts[ind])).datetime).astimezone(pacific))
 
#October 14
GEM = nc.Dataset('/ocean/dlatorne/MEOPAR/GEM2.5/NEMO-atmos/res_y2014m10d14.nc')
GEM_precip=np.append(GEM_precip,GEM.variables['precip'][:,j,i])
ts=GEM.variables['time_counter']
for ind in np.arange(ts.shape[0]):
        GEM_time.append(((torig + datetime.timedelta(seconds=ts[ind])).datetime).astimezone(pacific))
GEM_time=np.array(GEM_time)

precip={'GEM': [], 'OP': []}; time={'GEM': [], 'OP': []}
for t, p in zip(GEM_time, GEM_precip):
    if int(t.day)==13:
        precip['GEM'].append(p);
        time['GEM'].append(t)
precip['GEM']=np.array(precip['GEM'])
time['GEM']=np.array(time['GEM'])

fig,ax=plt.subplots(1,1,figsize=(10,5))
ax.plot(time['GEM'],precip['GEM'])
ax.set_ylabel('Precipitation (kg/m^2/s)')
ax.set_xlabel('Time (Local)')

hfmt = dates.DateFormatter('%m/%d %H:%M',tz=pacific)
ax.xaxis.set_major_formatter(hfmt)

print precip['GEM'], time['GEM']

print np.max(precip['GEM']) + precip['GEM'][-1]

[j,i]=find_model_point(lon['Sandheads'],lat['Sandheads'],OPlon,OPlat)

#October 13
OP = nc.Dataset('/ocean/sallen/allen/research/Meopar/Operational/oper_allvar_y2014m10d13.nc')
OP_precip = OP.variables['precip'][:,j,i]
ts=OP.variables['time_counter']
torig = datetime.datetime(1970,1,1)
torig=torig.replace(tzinfo=pacific)#there is no time_origin attriubte in OP files, so I hard coded this
OP_time=[]
for ind in np.arange(ts.shape[0]):
    OP_time.append(((torig + datetime.timedelta(seconds=ts[ind]))).astimezone(pacific))
 
#October 14
OP = nc.Dataset('/ocean/sallen/allen/research/Meopar/Operational/oper_allvar_y2014m10d14.nc')
OP_precip=np.append(OP_precip,OP.variables['precip'][:,j,i])
ts=OP.variables['time_counter']
for ind in np.arange(ts.shape[0]):
    OP_time.append(((torig + datetime.timedelta(seconds=ts[ind]))).astimezone(pacific))
OP_time=np.array(OP_time)

precip['OP']=[]; time['OP']=[]
for t, p in zip(OP_time, OP_precip):
    if int(t.day)==13:
        precip['OP'].append(p);
        time['OP'].append(t)
precip['OP']=np.array(precip['OP'])
time['OP']=np.array(time['OP'])

fig,ax=plt.subplots(1,1,figsize=(10,5))
ax.plot(time['OP'],precip['OP'])
ax.set_ylabel('Precipitation (kg/m^2/s)')
ax.set_xlabel('Time (Local)')

hfmt = dates.DateFormatter('%m/%d %H:%M',tz=pacific)
ax.xaxis.set_major_formatter(hfmt)

dt = 3600;
scale = 1000;
total_rain = np.trapz(precip['OP']/scale, dx=dt)
print total_rain