
from datetime import datetime, timedelta
from multiprocessing import Pool

import os
import glob
import sys

import matplotlib
matplotlib.use('Agg')

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


import smartcarb as sc


def load_data(filename, symbols):

    with netCDF4.Dataset(filename) as nc:
        time_string = datetime(2015,1,1) + timedelta(seconds=int(np.mean(nc.variables['time'][:])))
        time_string = time_string.strftime('%Y-%m-%d %H:%M UTC')

        sat_number = {
                0:   1,
                90:  2,
                180: 3,
                270: 4
        }[int(nc.variables['input'].longitude_at_daytime_equator_pass_in_degrees)]


        lon = nc.variables['longitude_corners'][:,:,:]
        lat = nc.variables['latitude_corners'][:,:,:]

        # cloud fraction
        clct = nc.variables['CLCT'][:]

        data = {}

        try:
            # XCO2
            if 'CO2' in symbols:
                data['CO2'] = (
                    nc.variables['XCO2_BG'][:] + nc.variables['XCO2_TOT'][:] +
                    nc.variables['XCO2_RA'][:] - nc.variables['XCO2_GPP'][:] -
                    nc.variables['XCO2_BG'][:].mean()
                )

            if 'uCO2' in symbols:
                data['uCO2'] = 0.667 * nc.variables['uXCO2'][:]  # uncertainty for 0.5 ppm scenario

            # NO2 columns
            if 'NO2' in symbols:
                data['NO2'] = nc.variables['NO2_BG'][:] + nc.variables['NO2_TOT'][:]
                data['NO2'] *= 1e-15

            if 'uNO2_high' in symbols:
                data['uNO2_high'] = 1e-15 * nc.variables['uNO2_high'][:]

            # CO column
            if 'CO' in symbols:
                data['CO'] = nc.variables['CO_BG'][:] + nc.variables['CO_TOT'][:]
                data['CO'] *= 1e-18

            if 'uCO_low' in symbols:
                data['uCO_low'] = 1e-18 * nc.variables['uCO_low'][:]


        except KeyError:
            print('key error in "%s"' % filename)
            raise


        return lon, lat, data, clct, time_string, sat_number


def main_daily(filenames, symbols):

    fig = None
    sats = []

    for filename in filenames:
        # load data
        lon, lat, data, clct, time_string, sat_number = load_data(filename, symbols)
        sats.append(sat_number)

        # split filename
        path = os.path.dirname(filename)

        # make output dirs
        for symbol in symbols:
            os.makedirs(
                os.path.join(path, '%s_daily' % symbol),
                exist_ok=True
            )

        symbol, vmin, vmax, label = ('CO2', -1.5,   1.5, 'XCO$_2$ signal (ppm)')
        tracer = data[symbol]

        fig = sc.map.make_level2_map(lon, lat, tracer, vmin=vmin, vmax=vmax,
            label=label, clct=clct, fig=fig)

    # add date and sat. number string
    ax = fig.get_axes()[0]
    time_string = "%s, sat#%s" % (time_string[:10], ','.join(str(s) for s in sats))
    ax.text(0, 1, time_string, transform=ax.transAxes, va='top', ha='left',
            fontsize='large', backgroundcolor='white')


    image_filename = os.path.basename(filename).replace('_data.nc', '_%s_clct.png' % symbol)
    image_filename = os.path.join(path, '%s_daily' % symbol, image_filename)
    fig.savefig(image_filename)

    plt.close()


def main(filename, symbols):

    # load data
    lon, lat, data, clct, time_string, sat_number = load_data(filename, symbols)

    # split filename
    path = os.path.dirname(filename)

    # make output dirs
    for symbol in symbols:
        os.makedirs(
            os.path.join(path, symbol),
            exist_ok=True
        )

    for symbol, vmin, vmax, label in [
                ('CO2', -1.5,   1.5, 'XCO$_2$ signal (ppm)'),
                ('uCO2', 0,     1.2, 'XCO$_2$ uncertainty (ppm)'),
                ('NO2',  0.0,   16.0, 'NO$_2$ columns (10$^{15}$ cm$^{-2}$)'),
                ('uNO2', 0.0,   2.0, 'NO$_2$ uncertainty (10$^{15}$ cm$^{-2}$)'),
                ('CO',   -1.0,  1.0, 'CO signal (10$^{18}$ cm$^{-2}$)'),
                ('uCO',  0.0,   2.0, 'CO uncertainty (10$^{18}$ cm$^{-2}$)'),
            ]:
        if symbol not in symbols:
            continue

        tracer = data[symbol]

        for clct in [None, clct]:

            fig = sc.map.make_level2_map(lon, lat, tracer, vmin=vmin, vmax=vmax,
                label=label, clct=clct)

            # add date and sat. number string
            ax = fig.get_axes()[0]
            time_string = "%s, sat#%d" % (time_string[:16], sat_number)
            ax.text(0, 1, time_string, transform=ax.transAxes, va='top', ha='left',
                    fontsize='large', backgroundcolor='white')


            if clct is None:
                image_filename = os.path.basename(filename).replace('.nc', '_%s.png' % symbol)
            else:
                image_filename = os.path.basename(filename).replace('.nc', '_%s_clct.png' % symbol)

            image_filename = os.path.join(path, symbol, image_filename)
            fig.savefig(image_filename)

            plt.close()


def do_multi(filename):
    main(filename, ['CO2', 'NO2'])
    #try:
    #    main(filename, ['NO2', 'uNO2_high', 'CO', 'uCO_low'])
    #except:
    #    print('filename "%s" had exception' % filename)


if __name__ == '__main__':
    filenames = sys.argv[1:]

    for fn in filenames:
        print(fn)
        main(fn, ['CO2'])

    #pool = Pool(processes=24)
    #pool.map(do_multi, filenames)





