
import sys

from netCDF4 import Dataset
import matplotlib.pyplot as plt
import numpy as np

import smartcarb as sc


filename = "/scratch/snx3000/gkuhl/smartcarb/cosmo2D_2015-01/cosmo_2d_2015012411.nc"
em_correction = 0.55 # for jan and jul

with Dataset(filename) as nc:
    rlon = sc.cfg.rlon
    rlat = sc.cfg.rlat


    b0 = nc.variables['XCO2_B0'][0] # industry
    b1 = nc.variables['XCO2_B1'][0] # heating
    b2 = nc.variables['XCO2_B2'][0] # traffic

    no2_b0 = nc.variables['NO2_B0'][0] # industry
    no2_b1 = nc.variables['NO2_B1'][0] # heating
    no2_b2 = nc.variables['NO2_B2'][0] # traffic

    co_b0 = nc.variables['CO_B0'][0] # industry
    co_b1 = nc.variables['CO_B1'][0] # heating
    co_b2 = nc.variables['CO_B2'][0] # traffic


# apply correction factor for Jan
b0 *= em_correction
no2_b0 *= em_correction
co_b0 *= em_correction


xlims = (1.4, 2.7)
ylims = (5.4, 6.5)

names = ['industry', 'heating', 'traffic', 'sum']

for name, xco2 in zip(names, [b0,b1,b2]):
    fig = sc.map.make_cosmo_map(rlon, rlat, xco2, vmin=0, vmax=0.3,
            label='XCO$_2$ signal (ppm)',
            xlim=xlims, ylim=ylims, edgecolor='w')

    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'w')

    fig.savefig('XCO2_2015012411_%s.png' % name)
    plt.close()


sys.exit()


for name, no2 in zip(names, [no2_b0, no2_b1, no2_b2]):
    fig = sc.map.make_cosmo_map(rlon, rlat, no2/1e15, vmin=0, vmax=5,
            label='NO$_2$ columns (10$^{15}$ cm$^{-2}$)',
            xlim=xlims, ylim=ylims, edgecolor='w')

    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'w')

    fig.savefig('NO2_2015012411_%s.png' % name)
    plt.close()


for name, co in zip(names, [co_b0, co_b1, co_b2]):
    fig = sc.map.make_cosmo_map(rlon, rlat, co/1e17, vmin=0, vmax=0.2,
            label='CO columns (10$^{17}$ cm$^{-2}$)',
            xlim=xlims, ylim=ylims, edgecolor='w')

    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'w')

    fig.savefig('CO_2015012411_%s.png' % name)
    plt.close()


with Dataset("/project/d73/input/emissions/berlin/emissions_berlin2.nc") as nc:
    area = nc['Berlin_Area_Fraction'][2:-2,2:-2,0]


plume = ((b0+b1+b2) > 0.1) | (area > 0)


for name, xco2, no2 in zip(names,
        [b0,b1,b2,b0+b1+b2],
        [no2_b0, no2_b1, no2_b2, no2_b0+no2_b1+no2_b2]):

    vmax = 10 if name == 'heating' else 50
    if name == 'sum':
        vmax = 25.0

    r = no2/1e15/xco2
    r = np.ma.array(r, mask=~(plume & (r > 0.0) & np.isfinite(r)))

    fig = sc.map.make_cosmo_map(rlon, rlat, r, vmin=0, vmax=vmax,
            label='NO$_2$ : XCO$_2$ (10$^{15}$ cm$^{-2}$ : ppm)',
            xlim=xlims, ylim=ylims, edgecolor='k')

    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'k')

    fig.savefig('NO2-CO2_2015012411_%s.png' % name)
    plt.close()




for name, xco2, co in zip(names,
        [b0,b1,b2,b0+b1+b2],
        [co_b0, co_b1, co_b2, co_b0+co_b1+co_b2]):

    if name == 'traffic':
        vmax = 4
    else:
        vmax = 1

    r = co/1e17/xco2
    r = np.ma.array(r, mask=~(plume & (r > 0.0) & np.isfinite(r)))

    fig = sc.map.make_cosmo_map(rlon, rlat, r, vmin=0, vmax=vmax,
            label='CO : XCO$_2$ (10$^{17}$ cm$^{-2}$ : ppm)',
            xlim=xlims, ylim=ylims, edgecolor='k')

    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'k')

    fig.savefig('CO-CO2_2015012411_%s.png' % name)
    plt.close()








