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'][:]

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

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)

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'])

[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'])