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

from __future__ import print_function
from datetime import datetime, timedelta
from multiprocessing import Pool
from netCDF4 import Dataset

import glob
import os
import sys

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

import smartcarb as sc
import config

zoom_ins = {
        'hamburg': {'xlim': [-1.0, 1.5], 'ylim': [5.6, 7.7]},
        'berlin': {'xlim': [1.5, 3.0], 'ylim': [4.8, 6.05]}
}


def level2_maps(filename):
    print(os.path.basename(filename))

    # load datasets
    with Dataset(filename) as nc:

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

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

        # XCO2
        xco2 = (
                nc.variables['XCO2_BG'][:] + nc.variables['XCO2_TOT'][:] +
                nc.variables['XCO2_RA'][:] - nc.variables['XCO2_GPP'][:]
        )
        # NO2 columns
        no2 = nc.variables['NO2_BG_column'][:] + nc.variables['NO2_TOT_column'][:]
        no2 *= 1e-15

        # CO column
        co = nc.variables['CO_BG_column'][:] + nc.variables['CO_TOT_column'][:]
        co *= 1e-18


    # split filename
    parts = os.path.basename(filename).split('_')
    instrument = '_'.join(parts[:-4])
    extra = '_'.join(parts[-4:-1])


    # make output dirs
    for symbol in ['CO2', 'CO', 'NO2']:
        path = os.path.join(config.path_figures, instrument, symbol)
        os.makedirs(path, exist_ok=True)

    for symbol, tracer, vmin, vmax, label in [
                ('CO2', xco2, 399.0, 403.0, 'XCO$_2$ (ppm)'),
                ('NO2',  no2,   0.0,  16.0, 'NO$_2$ columns (10$^{15}$ cm$^{-2}$)'),
                ('CO',    co,   1.0,   2.0, 'CO columns (10$^{18}$ cm$^{-2}$)'),
            ]:

        for clct in [None, clct]:

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

            if clct is None:
                filename = '%s_%s_%s.png' % (instrument, symbol.lower(), extra)
            else:
                filename = '%s_%s_clct_%s.png' % (instrument, symbol.lower(), extra)

            filename = os.path.join(config.path_figures, instrument, symbol, filename)
            fig.savefig(filename)

            plt.close()





def main(filename):
    print(filename)

    # create folders if necessary
    os.makedirs(config.path_figures, exist_ok=True)
    for name in ['CO2', 'NO2', 'CO']:
        path = os.path.join(config.path_figures, name)
        os.makedirs(path, exist_ok=True)

    border = 0.05

    date = os.path.splitext(os.path.basename(filename))[0].split('_')[-1]
    date_string = '%s-%s-%s %s:00 UTC' % (date[:4], date[4:6], date[6:8], date[8:10])


    if os.path.exists(filename):

        with Dataset(filename) as nc:

            rlon = nc.variables['rlon'][:]
            rlat = nc.variables['rlat'][:]

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

            # XCO2
            xco2 = (
                nc.variables['XCO2_BG'][0] + nc.variables['XCO2_TOT'][0] +
                nc.variables['XCO2_RA'][0] - nc.variables['XCO2_GPP'][0] -
                nc.variables['XCO2_BG'][0].mean()
            )
            # NO2 columns
            no2 = nc.variables['NO2_BG'][0] + nc.variables['NO2_TOT'][0]
            no2 *= 1e-15

            # CO column
            co = nc.variables['CO_BG'][0] + nc.variables['CO_TOT'][0]
            co *= 1e-18

    else:
        shape = (sc.config.je, sc.config.ie)
        rlon = sc.config.rlon
        rlat = sc.config.rlat

        xco2 = np.full(shape, np.nan)
        clct = np.full(shape, np.nan)
        no2 = np.full(shape, np.nan)
        co = np.full(shape, np.nan)

    for symbol, tracer, vmin, vmax, label in [
                ('CO2', xco2, -2.0,  2.0, 'XCO$_2$ signal (ppm)'),
                ('NO2',  no2,  0.0, 16.0, 'NO$_2$ columns (10$^{15}$ cm$^{-2}$)'),
                ('CO',    co,  1.0,  2.0, 'CO columns (10$^{18}$ cm$^{-2}$)'),
            ]:

        for clct in [None, clct]:

            # figure
            fig = sc.map.make_cosmo_map(rlon, rlat, tracer, vmin=vmin, vmax=vmax,
                    border=border, label=label, clct=clct)

            # add time label
            ax = fig.get_axes()[0]
            ax.text(0, 1, date_string, transform=ax.transAxes, va='top', ha='left',
                    fontsize='large', backgroundcolor='white')

            # savefig
            if clct is None:
                filename = 'cosmo_%s_%s.png' % (symbol.lower(), date)
            else:
                filename = 'cosmo_%s_clct_%s.png' % (symbol.lower(), date)

            filename = os.path.join(config.path_figures, symbol, filename)
            fig.savefig(filename)

            # zoom-in
            ax.set_xlim(0.4,4.5)
            ax.set_ylim(3.5,6.9)

            # savefig
            if clct is None:
                filename = 'cosmo_%s_zoom_%s.png' % (symbol.lower(), date)
            else:
                filename = 'cosmo_%s_clct_zoom_%s.png' % (symbol.lower(), date)

            filename = os.path.join(config.path_figures, symbol, filename)
            fig.savefig(filename)

            # close figure
            plt.close()


if __name__ == '__main__':

    pool = Pool(processes=24)
    pool.map(main, sys.argv[1:])


