import zmq
import IPython
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from mpl_toolkits.axes_grid1 import ImageGrid

print 'IPython Version:', IPython.__version__
print 'Matplotlib Version:', matplotlib.__version__
print 'ZMQ Version:', zmq.__version__


fpath = '/Users/pmarshwx/code/notebooks/nhc_data/hurdat2_atlantic.dat'
lines = open(fpath, 'r').readlines()

sandy_w = [30, 30, 40, 45, 45, 45, 45, 45, 50, 50, 50, 50, 60,
           65, 70, 70, 70, 80, 80, 80, 80, 85, 90, 110, 110, 110,
           105, 105]
sandy_p = [1003, 1003, 999, 998, 998, 998, 998, 997, 993, 993, 993, 993, 989,
           988, 986, 986, 983, 973, 973, 973, 970, 968, 954, 957, 957, 957,
           960, 967]
sandy_w = np.array(sandy_w)
sandy_p = np.array(sandy_p)

wind = []
pres = []
for line in lines:
    parts = line.split(',')
    if len(parts) < 10: continue
    wind.append(int(parts[5]))
    pres.append(int(parts[6]))
wind = np.ma.array(wind)
pres = np.ma.array(pres)

wind[wind < 0] = np.ma.masked
wind[pres < 0] = np.ma.masked
pres[wind < 0] = np.ma.masked
pres[pres < 0] = np.ma.masked

# inds = np.where(pres <= 954)
wvals = wind.compressed()
pvals = pres.compressed()

MPH2KTS = 0.868976242

xbins = range(35, 176, 5)
ybins = range(880, 1031, 5)
norm = mpl.colors.Normalize(vmin=0, vmax=3)
fig = plt.figure(figsize=(12,8))
grid = ImageGrid(fig, 111, nrows_ncols = (1, 1), direction="column",
                 axes_pad = 1, add_all=True, label_mode = "1",
                 share_all = True, cbar_location="bottom", cbar_mode="each",
                 cbar_size="5%", cbar_pad=1, aspect=False)
ax0 = grid[0]; cax0 = grid.cbar_axes[0]
counts, xbins, ybins, cbar0 = ax0.hist2d(wvals, pvals, bins=(xbins, ybins), cmap=plt.cm.cool, norm=norm)
counts = counts.T
c = cbar0.get_array()
c = np.ma.asanyarray(c)
c[c==0] = np.ma.masked
c = np.log10(c)
cbar0.set_array(c)

cax0.colorbar(cbar0)
cax0.set_xlabel('Frequency of Occurrence', size=10)
ticks = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000]
logticks = np.log10(ticks)
# ticks = np.delete(ticks, 1)
cax0.set_xticks(logticks)
clabs = ['%i' % (10**x) for x in cax0.get_xticks()]
cax0.set_xticklabels(ticks)

ax0.plot(sandy_w*MPH2KTS, sandy_p, linestyle='--', linewidth=1, marker='o',
    markersize=10, color='red', zorder=10)

ax0.axvline(64, linewidth=1, linestyle='--', color='k')
ax0.axvline(83, linewidth=1, linestyle='--', color='k')
ax0.axvline(96, linewidth=1, linestyle='--', color='k')
ax0.axvline(113, linewidth=1, linestyle='--', color='k')
ax0.axvline(137, linewidth=1, linestyle='--', color='k')

ax0.grid()
ax0.set_xlim(xbins[0], xbins[-1]-1)
ax0.set_xticks(xbins)
ax0.set_xticklabels(ax0.get_xticks().astype(int), size=8)
ax0.set_xlabel('\nMaximum Wind Speed (kts)', size=12)

ax0.set_ylim(ybins[0], ybins[-1]-1)
ax0.set_yticks(ybins)
ax0.set_yticklabels(ax0.get_yticks().astype(int), size=8)
ax0.set_ylabel('Minimum Pressure (hPa)\n', size=12)

ax0.tick_params(axis='both', direction='out', which='major',
                length=10, width=1, pad=10, top=False, right=False)


ax0.set_title('\nAtlantic Basin Wind vs. Pressure\nBest Track Data 1851-2011',
    size=18);

plt.savefig('WvsP.png')
# plt.show()

MPH2KTS = 0.868976242

xbins = range(35, 176, 5)
ybins = range(880, 1031, 5)
norm = mpl.colors.Normalize(vmin=0, vmax=3)
fig = plt.figure(figsize=(12,8))
grid = ImageGrid(fig, 111, nrows_ncols = (1, 1), direction="column",
                 axes_pad = 1, add_all=True, label_mode = "1",
                 share_all = True, cbar_location="bottom", cbar_mode="each",
                 cbar_size="5%", cbar_pad=1, aspect=False)
ax0 = grid[0]; cax0 = grid.cbar_axes[0]
counts, xbins, ybins, cbar0 = ax0.hist2d(wvals, pvals, bins=(xbins, ybins), cmap=plt.cm.cool, norm=norm)
counts = counts.T
c = cbar0.get_array()
c = np.ma.asanyarray(c)
c[c==0] = np.ma.masked
c = np.log10(c)
cbar0.set_array(c)

x, y = np.meshgrid(xbins[:-1], ybins[:-1])
for i in range(x.shape[0]-1):
    for j in range(y.shape[1]-1):
        if counts[i,j] <= 0: continue
        xx = (x[i,j] + x[i, j+1]) / 2.
        yy = (y[i,j] + y[i+1, j]) / 2.
        ax0.text(xx, yy, int(counts[i,j]), ha='center', va='center',
            size=8, zorder=12)

cax0.colorbar(cbar0)
cax0.set_xlabel('Frequency of Occurrence', size=10)
ticks = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000]
logticks = np.log10(ticks)
# ticks = np.delete(ticks, 1)
cax0.set_xticks(logticks)
clabs = ['%i' % (10**x) for x in cax0.get_xticks()]
cax0.set_xticklabels(ticks)

ax0.plot(sandy_w*MPH2KTS, sandy_p, linestyle='--', linewidth=1, marker='o',
    markersize=10, color='red', zorder=10)

line0 = Line2D(range(5), range(5), ls='--', marker='.', color='r', ms=10)
label0 = 'Tropical Cyclone Sandy'
ax0.legend((line0,), (label0,), loc=3, shadow=True, fancybox=True,
    numpoints=1, fontsize=10)


ax0.axvline(64, linewidth=1, linestyle='--', color='k')
ax0.axvline(83, linewidth=1, linestyle='--', color='k')
ax0.axvline(96, linewidth=1, linestyle='--', color='k')
ax0.axvline(113, linewidth=1, linestyle='--', color='k')
ax0.axvline(137, linewidth=1, linestyle='--', color='k')

ax0.grid()
ax0.set_xlim(xbins[0], xbins[-1]-1)
ax0.set_xticks(xbins)
ax0.set_xticklabels(ax0.get_xticks().astype(int), size=8)
ax0.set_xlabel('\nMaximum Wind Speed (kts)', size=12)

ax0.set_ylim(ybins[0], ybins[-1]-1)
ax0.set_yticks(ybins)
ax0.set_yticklabels(ax0.get_yticks().astype(int), size=8)
ax0.set_ylabel('Minimum Pressure (hPa)\n', size=12)

ax0.tick_params(axis='both', direction='out', which='major',
                length=10, width=1, pad=10, top=False, right=False)


ax0.set_title('\nAtlantic Basin Wind vs. Pressure\nBest Track Data 1851-2011',
    size=18);

credit = 'Created By:\nPatrick Marsh\nhttp://www.patricktmarsh.com'
ax0.text(0.875, 0.925, credit, ha='center', va='center', fontsize=9,
         bbox=dict(boxstyle='round, pad=0.5, rounding_size=0.25', fc="#FAFAFA",
         ec="k", lw=0.5), transform=ax0.transAxes)

plt.savefig('WvsP.png')
# plt.show()