%load_ext load_style
%load_style talk.css

%matplotlib inline
from matplotlib import pyplot as plt

import os
import numpy as np
import pandas as pd
from mpl_toolkits.basemap import Basemap as bm

import xray

def plot_field(X, lat, lon, vmin, vmax, step, cmap=plt.get_cmap('jet'), ax=False, title=False, grid=False):
    if not ax: 
        f, ax = plt.subplots(figsize=(10, (X.shape[0] / float(X.shape[1])) * 10))
    m.ax = ax
    im = m.contourf(lons, lats, X, np.arange(vmin, vmax+step, step), latlon=True, cmap=cmap, extend='both', ax=ax)
    m.drawcoastlines()
    if grid: 
        m.drawmeridians(np.arange(0, 360, 60), labels=[0,0,0,1])
        m.drawparallels([-40, 0, 40], labels=[1,0,0,0])
    m.colorbar(im)
    if title: 
        ax.set_title(title)

PCs = pd.read_csv('../outputs/EOF_ERSST_PCs.csv', index_col=0, parse_dates=True)

PCs.head()

from sklearn.cluster import KMeans

nclusters = 6

kmeans = KMeans(init='k-means++', n_clusters=nclusters, n_init=10, n_jobs=-1)

kmeans.fit(PCs.values) 

kmeans.labels_

dpath = os.path.join(os.environ['HOME'],'data/SST/ER_SST/V4')
# dpath = os.path.join(os.environ['HOME'],'data/SST/ER_SST')

ncfname = os.path.join(dpath,'ersst.realtime.nc')

dset = xray.open_dataset(ncfname)

dsub = dset.sel(time=slice('1980','2014'),lat=slice(-40,40), lon=slice(120,290))

dsub

lat = dsub['lat'].values
lon = dsub['lon'].values
lons, lats = np.meshgrid(lon, lat)

labels = pd.DataFrame(kmeans.labels_, index=dsub['time'], columns=['cluster'])

labels.head()

pd.unique(labels.cluster)

c = 0

index = labels.query('cluster == {}'.format(c))

cluster = dsub.sel(time=index.index).mean('time')

plt.imshow(cluster['anom'][0,::-1,:])

m = bm(projection='cyl',llcrnrlat=-40,urcrnrlat=40,\
            llcrnrlon=120,urcrnrlon=290,\
            lat_ts=0,resolution='c')

f, axes = plt.subplots(nrows=3,ncols=2, figsize=(14,10))
f.subplots_adjust(hspace=0.1, wspace=0.1)
axes = axes.flatten()
for c in xrange(nclusters):
    index = labels.query('cluster == {}'.format(c))
    cluster = dsub.sel(time=index.index).mean('time')
    ax = axes[c]
    plot_field(cluster['anom'][0,:,:], lats, lons, -2, 2, 0.1, \
               ax=ax, cmap=plt.get_cmap('RdBu_r'), \
               title="Cluster #{}: {} months".format(c+1, len(index)))