month = {1:'jan', 2: 'feb', 3: 'mar', 4: 'apr', 5: 'may', 6: 'jun', 7: 'jul', 
         8: 'aug', 9: 'sep', 10: 'oct', 11: 'nov', 12: 'dec'}
sals, transports = np.array([]), np.array([])
months = 4*np.ones_like(sals)
years = 15*np.ones_like(sals)

south_sals, south_transports = np.array([]), np.array([])
south_months = 4*np.ones_like(sals)
south_years = 15*np.ones_like(sals)
for y in [15, 16, 17, 18]:
    print(y)
    for m in month:
        newsals, newtransports = get_data_backward(month[m], y)
        newmonths = m*np.ones_like(newsals)
        newyears = y*np.ones_like(newsals)
        sals = np.concatenate((sals, newsals))
        transports = np.concatenate((transports, newtransports))
        months = np.concatenate((months, newmonths))
        years = np.concatenate((years, newyears))

        south_newsals, south_newtransports = get_data_backward_south(month[m], y)
        south_newmonths = m*np.ones_like(south_newsals)
        south_newyears = y*np.ones_like(south_newsals)
        south_sals = np.concatenate((south_sals, south_newsals))
        south_transports = np.concatenate((south_transports, south_newtransports))
        south_months = np.concatenate((south_months, south_newmonths))
        south_years = np.concatenate((south_years, south_newyears))

print (transports.sum()/4./365.25/24)
south_transports.sum()/4./365.25/24

salt_bins = np.arange(20, 33, 0.5)
month_bins = np.arange(0.5, 13, 1)
c, xedge, yedge, im = plt.hist2d((months), 
                                (sals),
                                weights=(transports)/24/30.5/4.,
                                bins=[month_bins, salt_bins], cmap=cm.matter, vmax=18000)
plt.colorbar();

fsals, ftransports = np.array([]), np.array([])
fmonths = 4*np.ones_like(fsals)
fyears = 15*np.ones_like(fsals)

south_fsals, south_ftransports = np.array([]), np.array([])
south_fmonths = 4*np.ones_like(south_fsals)
south_fyears = 15*np.ones_like(south_fsals)
for y in [15, 16, 17, 18]:
    print(y)
    for m in month:
        newsals, newtransports = get_data_forward(month[m], y)
        newmonths = m*np.ones_like(newsals)
        newyears = y*np.ones_like(newsals)
        fsals = np.concatenate((fsals, newsals))
        ftransports = np.concatenate((ftransports, newtransports))
        fmonths = np.concatenate((fmonths, newmonths))
        fyears = np.concatenate((fyears, newyears))

        south_newsals, south_newtransports = get_data_forward_south(month[m], y)
        south_newmonths = m*np.ones_like(south_newsals)
        south_newyears = y*np.ones_like(south_newsals)
        south_fsals = np.concatenate((south_fsals, south_newsals))
        south_ftransports = np.concatenate((south_ftransports, south_newtransports))
        south_fmonths = np.concatenate((south_fmonths, south_newmonths))
        south_fyears = np.concatenate((south_fyears, south_newyears))

gsals, gtransports = np.array([]), np.array([])
gmonths = 4*np.ones_like(gsals)
gyears = 15*np.ones_like(gsals)

south_gsals, south_gtransports = np.array([]), np.array([])
south_gmonths = 4*np.ones_like(south_gsals)
south_gyears = 15*np.ones_like(south_gsals)
for y in [15, 16, 17, 18]:
    print(y)
    for m in month:
        newsals, newtransports = get_data_backward_gi(month[m], y)
        newmonths = m*np.ones_like(newsals)
        newyears = y*np.ones_like(newsals)
        gsals = np.concatenate((gsals, newsals))
        gtransports = np.concatenate((gtransports, newtransports))
        gmonths = np.concatenate((gmonths, newmonths))
        gyears = np.concatenate((gyears, newyears))

        south_newsals, south_newtransports = get_data_backward_gi_south(month[m], y)
        south_newmonths = m*np.ones_like(south_newsals)
        south_newyears = y*np.ones_like(south_newsals)
        south_gsals = np.concatenate((south_gsals, south_newsals))
        south_gtransports = np.concatenate((south_gtransports, south_newtransports))
        south_gmonths = np.concatenate((south_gmonths, south_newmonths))
        south_gyears = np.concatenate((south_gyears, south_newyears))

gfsals, gftransports = np.array([]), np.array([])
gfmonths = 4*np.ones_like(gfsals)
gfyears = 15*np.ones_like(gfsals)

south_gfsals, south_gftransports = np.array([]), np.array([])
south_gfmonths = 4*np.ones_like(south_gfsals)
south_gfyears = 15*np.ones_like(south_gfsals)
for y in [15, 16, 17, 18]:
    print(y)
    for m in month:
        newsals, newtransports = get_data_forward_gi(month[m], y)
        newmonths = m*np.ones_like(newsals)
        newyears = y*np.ones_like(newsals)
        gfsals = np.concatenate((gfsals, newsals))
        gftransports = np.concatenate((gftransports, newtransports))
        gfmonths = np.concatenate((gfmonths, newmonths))
        gfyears = np.concatenate((gfyears, newyears))

        south_newsals, south_newtransports = get_data_forward_gi_south(month[m], y)
        south_newmonths = m*np.ones_like(south_newsals)
        south_newyears = y*np.ones_like(south_newsals)
        south_gfsals = np.concatenate((south_gfsals, south_newsals))
        south_gftransports = np.concatenate((south_gftransports, south_newtransports))
        south_gfmonths = np.concatenate((south_gfmonths, south_newmonths))
        south_gfyears = np.concatenate((south_gfyears, south_newyears))

print (np.concatenate((sals*transports, gsals*gtransports)).sum()/
      np.concatenate((transports, gtransports)).sum())
print (np.concatenate((fsals*ftransports, gfsals*gftransports)).sum()/
      np.concatenate((ftransports, gftransports)).sum())
print (np.concatenate((south_sals*south_transports, south_gsals*south_gtransports)).sum()/
      np.concatenate((south_transports, south_gtransports)).sum())
print (np.concatenate((south_fsals*south_ftransports, south_gfsals*south_gftransports)).sum()/
      np.concatenate((south_ftransports, south_gftransports)).sum())

print (32.44709649760663 - 30.86775963940009)
print (30.7054201601476 - 27.856508047159505)

fig, axs = plt.subplots(2, 2, figsize=(15, 20))
salt_bins = np.arange(20, 35, 0.5)
month_bins = np.arange(0.5, 13, 1)
cmap = cm.matter
cmap.set_bad('white')
mylabels = ['a)', 'b)', 'c)', 'd)']

c, xedge, yedge = np.histogram2d(np.concatenate((months, gmonths)), 
                            np.concatenate((sals, gsals)),
                            weights=np.concatenate((transports, gtransports))/24/30.5/4./1000.,
                            bins=[month_bins, salt_bins])
axs[0, 0].pcolormesh(xedge.T, yedge.T, c.T, vmin=0.2, vmax=18, cmap=cmap)
print (c.max())
print ('min', c.min())
c1 = c.copy()

c, xedge, yedge = np.histogram2d(np.concatenate((fmonths, gfmonths)), 
                            np.concatenate((fsals, gfsals)),
                            weights=np.concatenate((ftransports, gftransports))/24/30.5/4./1000.,
                            bins=[month_bins, salt_bins])
axs[0, 1].pcolormesh(xedge.T, yedge.T, c.T, vmin=0.2, vmax=18, cmap=cmap)
print (c.max())
print ('min', c.min())
c2 = c.copy()

c, xedge, yedge = np.histogram2d(np.concatenate((south_months, south_gmonths)), 
                            np.concatenate((south_sals, south_gsals)),
                            weights=np.concatenate((south_transports, south_gtransports))/24/30.5/4./1000.,
                            bins=[month_bins, salt_bins])
axs[1, 0].pcolormesh(xedge.T, yedge.T, c.T, vmin=0.2, vmax=18, cmap=cmap)
print (c.max())
print ('min', c.min())
c3 = c.copy()

c, xedge, yedge = np.histogram2d(np.concatenate((fmonths, gfmonths)), 
                            np.concatenate((south_fsals, south_gfsals)),
                            weights=np.concatenate((south_ftransports, south_gftransports))/24/30.5/4./1000.,
                            bins=[month_bins, salt_bins])
axs[1, 1].pcolormesh(xedge.T, yedge.T, c.T, vmin=0.2, vmax=18, cmap=cmap)
print (c.max())
print ('min', c.min())
c4 = c.copy()


for ii, ax in enumerate(axs.flatten()):
    ax.invert_yaxis()
    ax.text(0.013, 0.993, mylabels[ii], transform=ax.transAxes,
            fontsize=16, verticalalignment='top', fontfamily='serif',
            bbox=dict(facecolor='0.7', edgecolor='none', pad=3.0))
for ax in axs[1,:]:
    ax.set_xlabel("Month");
    ax.set_xticks(np.arange(1, 13))
    ax.set_xticklabels(labels=['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'], rotation=45)
for ax in axs[0,:]:
    ax.set_xticklabels([])
axs[0, 0].set_ylabel('Salinity at Point Roberts (g kg$^{-1}$)')
axs[1, 0].set_ylabel('Salinity at Victoria Sill (g kg$^{-1}$)')
axs[0, 0].set_title('Flux into Strait of Georgia')
axs[1, 1].set_title('Flux westward from Victoria Sill');
axs[1, 0].set_title('Flux eastward from Victoria Sill')
axs[0, 1].set_title('Flux out of Strait of Georgia');
fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
cb1 = fig.colorbar(im, cax=cbar_ax)
cb1.set_label('mSv', labelpad=2)

fig.savefig('/home/sallen/MEOPAR/estuarine_flux_paper/fluxbysalt_masked_v2.pdf')
fig.savefig('/home/sallen/MEOPAR/estuarine_flux_paper/fluxbysalt_masked_v2.png')

fig, axs = plt.subplots(2, 5, figsize=(20, 15))
salt_bins = np.arange(20, 33, 0.5)
month_bins = np.arange(0.5, 13, 1)
cin_mean, xedge, yedge, im = axs[0, 0].hist2d(np.concatenate((months, gmonths)), 
                            np.concatenate((sals, gsals)),
                            weights=np.concatenate((transports, gtransports))/24/30.5/4.,
                            bins=[month_bins, salt_bins], cmap=cm.matter, vmax=18000)

cout_mean, xedge, yedge, im = axs[1, 0].hist2d(np.concatenate((fmonths, gfmonths)), 
                            np.concatenate((fsals, gfsals)),
                            weights=np.concatenate((ftransports, gftransports))/24/30.5/4.,
                            bins=[month_bins, salt_bins], cmap=cm.matter, vmax=18000)
for ax in axs[:, 0]:
    ax.invert_yaxis()
    ax.set_xlabel("Month");
    ax.set_xticks(np.arange(1, 13))
    ax.set_xticklabels(labels=['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'], rotation=45)
    ax.set_ylabel('Salinity (g kg$^{-1}$)')
#axs[0].set_title('Flux into Strait of Georgia')
#axs[1].set_title('Flux out of Strait of Georgia');
fig.subplots_adjust(bottom=0.2)
cbar_ax = fig.add_axes([0.05, 0.05, 0.25, 0.05])
cb1 = fig.colorbar(im, cax=cbar_ax, orientation='horizontal')
cb1.set_label('m s$^{-3}$', labelpad=2)


for ix, year in enumerate([15, 16, 17, 18]):
    c, xedge, yedge = np.histogram2d(np.concatenate((months[years==year], gmonths[gyears==year])), 
                                np.concatenate((sals[years==year], gsals[gyears==year])),
                                weights=np.concatenate((transports[years==year], gtransports[gyears==year]))/24/30.5,
                                bins=[month_bins, salt_bins])#, cmap=cm.curl, vmax=18000)
    print (f'axis={ix+1}')
    axs[0, ix+1].pcolormesh(xedge, yedge, (c-cin_mean).transpose(), cmap=cm.curl, vmax=5000, vmin=-5000)
    print((c-cin_mean).max(), (c-cin_mean).min())
    
    c, xedge, yedge = np.histogram2d(np.concatenate((fmonths[fyears==year], gfmonths[gfyears==year])), 
                                np.concatenate((fsals[fyears==year], gfsals[gfyears==year])),
                                weights=np.concatenate((ftransports[fyears==year], gftransports[gfyears==year]))/24/30.5,
                                bins=[month_bins, salt_bins])#, cmap=cm.curl, vmax=18000)
    axs[1, ix+1].pcolormesh(xedge, yedge, (c-cout_mean).transpose(), cmap=cm.curl, vmax=5000, vmin=-5000)
    print((c-cout_mean).max(), (c-cout_mean).min())
    
for ix in range(1,5):
    for iy in range(2):
        axs[iy, ix].invert_yaxis()

#axs[0].set_ylabel('Salinity (g kg$^{-1}$)')
#axs[0].set_title('Flux into Strait of Georgia')
#axs[1].set_title('Flux out of Strait of Georgia');

fig, axs = plt.subplots(4, 2, figsize=(12, 20))
salt_bins = np.arange(20, 33, 0.5)
month_bins = np.arange(0.5, 13, 1)

for ix, year in enumerate([15, 16, 17, 18]):
    c, xedge, yedge, im = axs[ix, 0].hist2d(np.concatenate((months[years==year], gmonths[gyears==year])), 
                                np.concatenate((sals[years==year], gsals[gyears==year])),
                                weights=np.concatenate((transports[years==year], gtransports[gyears==year]))/24/30.5/1000.,
                                bins=[month_bins, salt_bins], cmap=cm.matter, vmax=21)
    
    c, xedge, yedge, im = axs[ix, 1].hist2d(np.concatenate((fmonths[fyears==year], gfmonths[gfyears==year])), 
                                np.concatenate((fsals[fyears==year], gfsals[gfyears==year])),
                                weights=np.concatenate((ftransports[fyears==year], gftransports[gfyears==year]))/24/30.5/1000.,
                                bins=[month_bins, salt_bins], cmap=cm.matter, vmax=21)
    for jj in range(2):
        axs[ix, jj].text(1.5, 22, f'20{year}')

for ax in axs[3]:
    ax.set_xlabel("Month");
    ax.set_xticks(np.arange(1, 13))
    ax.set_xticklabels(labels=['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'], rotation=45)
for ax in axs[:, 0]:
    ax.set_ylabel('Salinity (g kg$^{-1}$)')
for ax in axs[:, 1]:
    ax.set_yticks([])
    
axs[0, 0].set_title('Flux into Strait of Georgia')
axs[0, 1].set_title('Flux out of Strait of Georgia');
fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
cb1 = fig.colorbar(im, cax=cbar_ax)
cb1.set_label('mSv', labelpad=2)

for iy in range(2):
    for ix in range(4):
        axs[ix, iy].invert_yaxis()
    for ix in range(3):
        axs[ix, iy].set_xticks([])
    
fig.savefig('/home/sallen/MEOPAR/estuarine_flux_paper/fluxbysaltbyyear.pdf')
fig.savefig('/home/sallen/MEOPAR/estuarine_flux_paper/fluxbysaltbyyear.png')

mycolors = ['#e377c2', '#9467bd', '#17becf', '#1f77b4']
custom_lines = [mpl.lines.Line2D([0], [0], color=mycolors[0]),
                mpl.lines.Line2D([0], [0], color=mycolors[1]),
                mpl.lines.Line2D([0], [0], color=mycolors[2]),
                mpl.lines.Line2D([0], [0], color=mycolors[3])]
mpl.rcParams['axes.prop_cycle'] = cycler(color=mycolors)
fig, axs = plt.subplots(1, 2, figsize=(12, 5))
axs[0].plot(range(1, 13), np.sum(c1[:, 13:], axis=1), 'o-', color=mycolors[3])
axs[1].plot(range(1, 13), np.sum(c[:, 13:], axis=1), 'o-', color=mycolors[3])
axs[0].plot(range(1, 13), c1[:, 12], 'o-', color=mycolors[2])
axs[1].plot(range(1, 13), c[:, 12], 'o-', color=mycolors[2])
axs[0].plot(range(1, 13), c1[:, 11]+c1[:, 10]+c1[:, 9], 'o-', color=mycolors[1])
axs[1].plot(range(1, 13), c[:, 11]+c[:, 10]+c[:, 9], 'o-', color=mycolors[1])
axs[0].plot(range(1, 13), np.sum(c1[:, :9], axis=1), 'o-', color=mycolors[0])
axs[1].plot(range(1, 13), np.sum(c[:, :9], axis=1), 'o-', color=mycolors[0])
axs[1].legend(custom_lines, ['< 25 g/kg', '25-29.5 g/kg', '29.5-31 g/kg', '> 31 g/kg'])
for ax in axs:
    ax.set_ylim((0, 50e3))
    ax.set_xlabel('Month in 2015')
# doesn't make sense, check ranges.

print ((np.sum(c1, axis=1) - np.sum(c, axis=1)).mean())
print (np.sum(c1, axis=1).mean())
print (np.sum(c1, axis=1).mean()*20*1020*30.5*86400/1000.)
for year in [15, 16]:
    print (year)
    print ((transports[years==year].sum())/(24*365))
    print ((ftransports[fyears==year].sum())/(24*365))
    print ((gtransports[gyears==year].sum())/(24*365))
    print ((gftransports[gfyears==year].sum())/(24*365))
print ((transports.sum()+gtransports.sum() - ftransports.sum() - gftransports.sum())/(24.*365))

saltflux=np.zeros((12))
fsaltflux = np.zeros_like(saltflux)
gsaltflux = np.zeros_like(saltflux)
gfsaltflux = np.zeros_like(saltflux)

print(saltflux.shape)
for i in range(1, 13):
    saltflux[i-1] = (sals[months==i]*transports[months==i]).sum()*1020/1e6/24/30.5
    fsaltflux[i-1] = (fsals[fmonths==i]*ftransports[fmonths==i]).sum()*1020/1e6/24/30.5
    gsaltflux[i-1] = (gsals[gmonths==i]*gtransports[gmonths==i]).sum()*1020/1e6/24/30.5
    gfsaltflux[i-1] = (gfsals[gfmonths==i]*gftransports[gfmonths==i]).sum()*1020/1e6/24/30.5

fig, ax = plt.subplots(1, 1)
ax.plot(range(1, 13), saltflux+gsaltflux, label='Salt Flux In')
ax.plot(range(1, 13), fsaltflux+gfsaltflux, label='Salt Flux Out')
#plt.plot(range(1,13), fsaltflux)
#plt.plot(range(1, 13), gfsaltflux)
ax.legend(loc='upper left')
ax.set_ylabel('Salt Flux (Mg s$^{-1}$)')
ax.set_xlabel('Month in 2015');

print (saltflux.sum()+gsaltflux.sum(), fsaltflux.sum()+gfsaltflux.sum(), gfsaltflux.sum())

fraserriver = 3500 * 86400 * 365
fraserriver/1e13

section=5
ssals, stransports = get_data_forward('apr', section=section)
smonths = 4*np.ones_like(ssals)
for m in month:
    print(m)
    newsals, newtransports = get_data_forward(month[m], section=section)
    newmonths = m*np.ones_like(newsals)
    stransports = np.concatenate((stransports, newtransports))
    smonths = np.concatenate((smonths, newmonths))

print (stransports.sum()/(24*365))

section=5
ssals, stransports = get_data_backward('apr', section=section)
smonths = 4*np.ones_like(ssals)
for m in month:
    print(m)
    newsals, newtransports = get_data_backward(month[m], section=section)
    newmonths = m*np.ones_like(newsals)
    stransports = np.concatenate((stransports, newtransports))
    smonths = np.concatenate((smonths, newmonths))
print (stransports.sum()/(24*365))

section=5
ssals, stransports = get_data_forward_gi('apr', section=section)
smonths = 4*np.ones_like(ssals)
for m in month:
    print(m)
    newsals, newtransports = get_data_forward_gi(month[m], section=section)
    newmonths = m*np.ones_like(newsals)
    stransports = np.concatenate((stransports, newtransports))
    smonths = np.concatenate((smonths, newmonths))

print (stransports.sum()/(24*365))

section=5
ssals, stransports = get_data_backward_gi('apr', section=section)
smonths = 4*np.ones_like(ssals)
for m in month:
    print(m)
    newsals, newtransports = get_data_backward_gi(month[m], section=section)
    newmonths = m*np.ones_like(newsals)
    stransports = np.concatenate((stransports, newtransports))
    smonths = np.concatenate((smonths, newmonths))
print (stransports.sum()/(24*365))

print(21.1+7.7+4+0.3)
print(25.5+2.6+2.5+0.2)