#!/usr/bin/env python
# coding: utf-8

from datetime import datetime, timedelta

import os

import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
import numpy as np

import smartcarb as sc
import config


"""\
As an example of a plume we use 7.7.2015, 11 UTC.
"""

COSMO_FILENAME = '/scratch/snx1600/gkuhl/smartcarb/2d-fields/cosmo_2d_2015070211.nc'


XLIM = [0.75, 3.45]
YLIM = [3.9, 6.2]



def ratios(letter='B', tracer='NO2'):

    name = {'B': 'Berlin', 'J': 'Janschwalde'}[letter]

    if letter == 'J':
        xlim, ylim = None, None
    else:
        xlim = (0.5, 2.35)
        ylim = (5.2,6.8)


    rlon, rlat, co2 = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_%sV' % letter)
    rlon, rlat, other = sc.map.load_cosmo_data(COSMO_FILENAME, '%s_%sV_column' % (tracer, letter))

    no_plume = co2 < 0.1

    if tracer == 'NO2':
        r = np.ma.array(1e-15*other/co2, mask=no_plume)
        label = 'NO$_2$:CO$_2$ (10$^{15}$ cm$^{-2}$ / ppm)'
        vmax = 25.0
    else:
        r = np.ma.array(1e-18*other/co2, mask=no_plume)
        label = 'CO:CO$_2$ (10$^{18}$ cm$^{-2}$ / ppm)'
        vmax = 0.1


    fig = sc.map.make_cosmo_map(rlon, rlat, r,
            vmin=0, vmax=vmax, border=0.05, clct=None, label=label,
            xlim=xlim, ylim=ylim
    )
    sc.map.add_city_labels(fig.get_axes()[0], 'k')

    fig.savefig('%s_%svsCO2.png' % (name, tracer))
    plt.close()


    if letter == 'J':
        return

    for i,extra in enumerate(['industry', 'heating', 'traffic']):
        rlon, rlat, co2 = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_%s%d' % (letter, i))
        rlon, rlat, other = sc.map.load_cosmo_data(COSMO_FILENAME, '%s_%s%d_column' % (tracer, letter, i))

        if tracer == 'NO2':
            r = np.ma.array(1e-15*other/co2, mask=no_plume)
            label = 'NO$_2$:CO$_2$ (10$^{15}$ cm$^{-2}$ / ppm)'
            vmax = 25.0
        else:
            r = np.ma.array(1e-18*other/co2, mask=no_plume)
            label = 'CO:CO$_2$ (10$^{18}$ cm$^{-2}$ / ppm)'
            vmax = 0.1

        fig = sc.map.make_cosmo_map(rlon, rlat, r,
                vmin=0, vmax=vmax, border=0.05, clct=None, label=label,
                xlim=xlim, ylim=ylim
        )
        sc.map.add_city_labels(fig.get_axes()[0], 'k')

        fig.savefig('%s_%svsCO2_%s.png' % (name, tracer, extra))
        plt.close()





def berlin():
    """\
    Berlin plume
    - total
    - etc.
    """

    label = 'XCO$_2$ signal (ppm)'
    rlon, rlat, co2_bg = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_BG')
    print(co2_bg.mean())

    for tracer in ['XCO2', 'XCO2_BG', 'XCO2_BV', 'XCO2_B0', 'XCO2_B1', 'XCO2_B2', 'XCO2_bio',
            'XCO2_outside']:

        if tracer == 'XCO2_bio':
            rlon, rlat, gpp = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_GPP')
            rlon, rlat, ra = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_RA')
            co2 = ra - gpp

        elif tracer == 'XCO2_outside':
            rlon, rlat, a = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_A')
            rlon, rlat, jv = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_JV')
            co2 = a + jv

        else:
            rlon, rlat, co2 = sc.map.load_cosmo_data(COSMO_FILENAME, tracer)

        if tracer in ['XCO2', 'XCO2_BG']:
            co2 -= co2_bg.mean()

        fig = sc.map.make_cosmo_map(rlon, rlat, co2,
                vmin=-0.6, vmax=0.6, border=0.5, clct=None, label=label,
                xlim=(0.5, 2.35), ylim=(5.2,6.8)
        )

        fig.savefig('Berlin_%s.png' % tracer)
        plt.close()




def make_cosmo_plot(xlim=None, ylim=None, extra='', add_labels=False):
    # CO2
    rlon, rlat, co2_tot = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_TOT')
    rlon, rlat, co2_bg = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_BG')
    rlon, rlat, co2_gpp = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_GPP')
    rlon, rlat, co2_ra = sc.map.load_cosmo_data(COSMO_FILENAME, 'XCO2_RA')
    rlon, rlat, clct = sc.map.load_cosmo_data(COSMO_FILENAME, 'CLCT')

    co2_ref = 400.0
    label = 'XCO$_2$ signal (ppm)'

    # total field
    fig = sc.map.make_cosmo_map(rlon, rlat, co2_tot+co2_bg+co2_ra-co2_gpp-co2_ref,
            vmin=-4, vmax=4, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('XCO2_sum%s.png' % extra)
    plt.close()

    # biosphere
    fig = sc.map.make_cosmo_map(rlon, rlat, co2_ra-co2_gpp,
            vmin=-4, vmax=4, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('XCO2_biosphere%s.png' % extra)
    plt.close()

    # background
    fig = sc.map.make_cosmo_map(rlon, rlat, co2_bg-co2_ref,
            vmin=-4, vmax=4, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('XCO2_background%s.png' % extra)
    plt.close()

    # anthro
    fig = sc.map.make_cosmo_map(rlon, rlat, co2_tot,
            vmin=-4, vmax=4, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('XCO2_anthro%s.png' % extra)
    plt.close()


    #NO2
    rlon, rlat, no2_tot = sc.map.load_cosmo_data(COSMO_FILENAME, 'NO2_TOT_column')
    rlon, rlat, no2_bg = sc.map.load_cosmo_data(COSMO_FILENAME, 'NO2_BG_column')

    label = 'NO$_2$ (10$^{15}$ molecules cm$^{-2}$)'
    fig = sc.map.make_cosmo_map(rlon, rlat, (no2_tot+no2_bg)/1e15,
            vmin=0, vmax=14, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim, edgecolor='white')
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'white')
    fig.savefig('NO2_sum%s.png' % extra)
    plt.close()

    fig = sc.map.make_cosmo_map(rlon, rlat, (no2_tot)/1e15,
            vmin=0, vmax=14, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim, edgecolor='white')
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'white')
    fig.savefig('NO2_anthro%s.png' % extra)
    plt.close()

    fig = sc.map.make_cosmo_map(rlon, rlat, (no2_bg)/1e15,
            vmin=0, vmax=14, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim, edgecolor='white')
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'white')
    fig.savefig('NO2_background%s.png' % extra)
    plt.close()


    # CO
    rlon, rlat, co_tot = sc.map.load_cosmo_data(COSMO_FILENAME, 'CO_TOT_column')
    rlon, rlat, co_bg = sc.map.load_cosmo_data(COSMO_FILENAME, 'CO_BG_column')

    co_ref = 1.4e18

    label = 'CO signal (10$^{18}$ molecules cm$^{-2}$)'
    fig = sc.map.make_cosmo_map(rlon, rlat, (co_tot+co_bg-co_ref)/1e18,
            vmin=-0.2, vmax=0.2, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('CO_sum%s.png' % extra)
    plt.close()

    fig = sc.map.make_cosmo_map(rlon, rlat, (co_tot)/1e18,
            vmin=-.2, vmax=.2, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('CO_anthro%s.png' % extra)
    plt.close()

    fig = sc.map.make_cosmo_map(rlon, rlat, (co_bg-co_ref)/1e18,
            vmin=-.2, vmax=.2, border=0.05, clct=clct, label=label,
            xlim=xlim, ylim=ylim)
    if add_labels:
        sc.map.add_city_labels(fig.get_axes()[0], 'black')
    fig.savefig('CO_background%s.png' % extra)
    plt.close()




def make_noisy_orbit():
    path = '/scratch/snx1600/gkuhl/smartcarb/test/'

    label = 'XCO$_2$ signal (ppm)'

    # load CO2 data and create noise
    filename = 'Sentinel_7_CO2_2015070211_o0000_l000_data.nc'
    filename = os.path.join(path, filename)
    lon, lat, clct = sc.map.load_orbit_data(filename, 'CLCT')
    lon, lat, co2 = sc.map.load_orbit_data(filename, 'XCO2')

    co2_00 = 0.0
    co2_05 = sc.noise.compute_uncertainty(co2, 0.5, None, is_relative=False)
    co2_07 = sc.noise.compute_uncertainty(co2, 0.7, None, is_relative=False)
    co2_10 = sc.noise.compute_uncertainty(co2, 1.0, None, is_relative=False)

    for key, noise in [('n00', co2_00), ('n05', co2_05), ('n07', co2_07), ('n10', co2_10)]:

        fig = sc.map.make_level2_map(lon, lat, (co2 + noise - 400),
                vmin=-4, vmax=4, clct=clct, label=label,
                xlim=XLIM, ylim=YLIM
        )
        ax, cax = fig.get_axes()
        sc.map.add_city_labels(ax, 'black')

        fig.savefig('Sentinel_7_XCO2_%s.png' % key)
        plt.close()


    # Sentinel 5 NO2/CO
    filename = 'Sentinel_5_2015070209_o0001_l000_data.nc'
    filename = os.path.join(path, filename)

    lon, lat, clct = sc.map.load_orbit_data(filename, 'CLCT')
    lon, lat, no2 = sc.map.load_orbit_data(filename, 'NO2')
    no2_std = sc.noise.compute_uncertainty(no2, 0.20, 1.3e15, is_relative=True)
    lon, lat, co = sc.map.load_orbit_data(filename, 'CO')
    co_std = sc.noise.compute_uncertainty(co, 0.1, 4e17, is_relative=True)

    # NO2 plots
    label = 'NO$_2$ column (10$^{15}$ molecules cm$^{-2}$)'
    fig = sc.map.make_level2_map(lon, lat, 1e-15*no2,
            vmin=0, vmax=20, clct=clct, label=label,
            xlim=XLIM, ylim=YLIM, edgecolor='white'
    )
    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'white')
    fig.savefig('Sentinel_5_NO2_n00.png')
    plt.close()

    fig = sc.map.make_level2_map(lon, lat, 1e-15*(no2+no2_std),
            vmin=0, vmax=20, clct=clct, label=label,
            xlim=XLIM, ylim=YLIM, edgecolor='white'
    )
    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'white')
    fig.savefig('Sentinel_5_NO2_n20.png')
    plt.close()

    # CO plots
    label = 'CO signal (10$^{18}$ molecules cm$^{-2}$)'

    fig = sc.map.make_level2_map(lon, lat, 1e-18*co,
            vmin=1.0, vmax=1.5, clct=clct, label=label,
            xlim=XLIM, ylim=YLIM, edgecolor='black'
    )
    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'black')

    fig.savefig('Sentinel_5_CO_n00.png')
    plt.close()

    fig = sc.map.make_level2_map(lon, lat, 1e-18*(co + co_std),
            vmin=1.0, vmax=1.5, clct=clct, label=label,
            xlim=XLIM, ylim=YLIM, edgecolor='black'
    )
    ax, cax = fig.get_axes()
    sc.map.add_city_labels(ax, 'black')

    fig.savefig('Sentinel_5_CO_n10.png')
    plt.close()


    # Sentinel 7 NO2/CO
    for i, filename in enumerate([
            'Sentinel_7_NO2_1km_2015070211_o0000_l000_data.nc',
            'Sentinel_7_NO2_2km_2015070211_o0000_l000_data.nc',
            'Sentinel_7_NO2_4km_2015070211_o0000_l000_data.nc'
        ]):
        filename = os.path.join(path, filename)

        # NO2
        label = 'NO$_2$ column (10$^{15}$ molecules cm$^{-2}$)'
        lon, lat, clct = sc.map.load_orbit_data(filename, 'CLCT')
        lon, lat, no2 = sc.map.load_orbit_data(filename, 'NO2')

        no2_15 = sc.noise.compute_uncertainty(no2, 0.15, 1e15, is_relative=True)
        no2_20 = sc.noise.compute_uncertainty(no2, 0.20, 2e15, is_relative=True)

        for key, noise in [('n00', 0.0), ('n15', no2_15), ('n20', no2_20)]:

            fig = sc.map.make_level2_map(lon, lat, 1e-15*(no2 + noise),
                    vmin=0, vmax=20, clct=clct, label=label,
                    xlim=XLIM, ylim=YLIM, edgecolor='white'
            )
            ax, cax = fig.get_axes()
            sc.map.add_city_labels(ax, 'white')

            fig.savefig('Sentinel_7_NO2_%dkm_%s.png' % ([1,2,4][i], key))
            plt.close()


        # CO
        label = 'CO signal (10$^{18}$ molecules cm$^{-2}$)'
        lon, lat, clct = sc.map.load_orbit_data(filename, 'CLCT')
        lon, lat, co = sc.map.load_orbit_data(filename, 'CO')

        co_10 = sc.noise.compute_uncertainty(co, 0.10, 0.0, is_relative=True)
        co_20 = sc.noise.compute_uncertainty(co, 0.20, 0.0, is_relative=True)

        for key, noise in [('n00', 0.0), ('n10', co_10), ('n20', co_20)]:

            fig = sc.map.make_level2_map(lon, lat, 1e-18*(co + noise),
                    vmin=1.0, vmax=1.5, clct=clct, label=label,
                    xlim=XLIM, ylim=YLIM, edgecolor='black'
            )
            ax, cax = fig.get_axes()
            sc.map.add_city_labels(ax, 'black')

            fig.savefig('Sentinel_7_CO_%dkm_%s.png' % ([1,2,4][i], key))
            plt.close()



def co2_orbit(date=datetime(2015,7,2,11)):
    path = '/scratch/snx1600/gkuhl/smartcarb/test/'

    # Sentinel-7_CO2
    filename = date.strftime('Sentinel_7_CO2_%Y%m%d%H_o0000_l000_data.nc')
    filename = os.path.join(path, filename)

    lon, lat, clct = sc.map.load_orbit_data(filename, 'CLCT')
    lon, lat, co2 = sc.map.load_orbit_data(filename, 'XCO2')

    co2_ref = co2.mean()
    label = 'XCO$_2$ signal (ppm)'

    # Sentinel_7 XCO2 
    fig = sc.map.make_level2_map(lon, lat, co2 - co2_ref,
            vmin=-2, vmax=2, clct=clct, label=label,
            xlim=XLIM, ylim=YLIM)
    ax, cax = fig.get_axes()

    sc.map.add_city_labels(ax, 'black')
    fig.savefig(date.strftime('Sentinel_7_CO2_%Y%m%d%H_ALL.png'))
    plt.close()



def no2_co_orbits():

    co_ref = 1.4e18
    path = '/scratch/snx1600/gkuhl/smartcarb/test/'

    for filename in [
        'Sentinel_7_NO2_1km_2015070211_o0000_l000_data.nc',
        'Sentinel_7_NO2_2km_2015070211_o0000_l000_data.nc',
        'Sentinel_7_NO2_4km_2015070211_o0000_l000_data.nc',
        'Sentinel_5_2015070211_o0001_l000_data.nc'
            ]:

        if filename.startswith('Sentinel_5'):
            instrument = 'Sentinel_5'
        else:
            instrument = '_'.join(filename.split('_')[:4])

        filename = os.path.join(path, filename)
        lon, lat, clct = sc.map.load_orbit_data(filename, 'CLCT')

        # CO
        lon, lat, co_tot = sc.map.load_orbit_data(filename, 'CO_TOT_column')
        lon, lat, co_bg = sc.map.load_orbit_data(filename, 'CO_BG_column')

        fig = sc.map.make_level2_map(lon, lat, 1e-18*(co_tot + co_bg - co_ref), vmin=1, vmax=1.5,
            clct=clct, label='CO signal (10$^{18}$ molecules cm$^{-2}$)', xlim=XLIM, ylim=YLIM,
            edgecolor='white')
        ax, cax = fig.get_axes()
        sc.map.add_city_labels(ax, 'black')


        fig.savefig('%s_2015070211_CO_sum.png' % instrument)
        plt.close()

        # NO2
        lon, lat, no2_tot = sc.map.load_orbit_data(filename, 'NO2_TOT_column')
        lon, lat, no2_bg = sc.map.load_orbit_data(filename, 'NO2_BG_column')

        fig = sc.map.make_level2_map(lon, lat, 1e-15*(no2_tot + no2_bg), vmin=0, vmax=20,
            clct=clct, label='NO$_2$ column (10$^{15}$ molecules cm$^{-2}$)',
            xlim=XLIM, ylim=YLIM, edgecolor='white')
        ax, cax = fig.get_axes()
        sc.map.add_city_labels(ax, 'white')

        fig.savefig('%s_2015070211_NO2_sum.png' % instrument)
        plt.close()



def main():
    pass


if __name__ == '__main__':
    for day in [ 1,  2,  3,  4,  5,  6,  7,  9, 11, 20, 23, 26]:
        co2_orbit(datetime(2015,7,day,11))
    for day in [ 18, 24]:
        co2_orbit(datetime(2015,1,day,11))


    #make_noisy_orbit()
    #berlin()
    #ratios('J', 'CO')
    #ratios('B', 'CO')
    #make_cosmo_plot(XLIM, YLIM, '_zoom', True)
    #co2_orbit()
    #no2_co_orbits()
