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

from __future__ import print_function

from datetime import datetime, timedelta
from multiprocessing import Pool
import os
import sys
import time
import netCDF4
import numpy as np

import matplotlib
matplotlib.use('Agg')

import c2sm
import smartcarb as sc
import config


def get_attrs(v):
    return dict((k, v.getncattr(k)) for k in v.ncattrs())


def create_2d_fields(input_filename, constant_filename, output_filename,
        overwrite=True, do_ground=False):
    """\
    Create input data for calculating satellite swaths.

    input_filename: COSMO output for a given hour.
    constant_filename: COSMO output for constant fields

    output_filename: output filename for 2d fields
    """
    sys.stdout.flush()
    attrs = {
            'DESCRIPTION':  'Vertical integrated/averaged tracer fields',
            'DATAORIGIN':   'Tracer fields from COSMO-GHG simulations',
            'CREATOR':      config.user_name,
            'EMAIL':        config.user_email,
            'AFFILIATION':  config.user_affiliation,
            'VERSION':      '0.1',
            'DATE CREATED': time.ctime(time.time()),
            'STUDY':        config.study,
    }

    if not os.path.exists(input_filename):
        print(input_filename, '...missing')
        sys.stdout.flush()
        return 42

    if not overwrite and os.path.exists(output_filename):
        return 42

    print(input_filename, '...working')
    sys.stdout.flush()

    tracers = config.tracers.copy()

    with netCDF4.Dataset(constant_filename) as nc:
        h = nc.variables['HHL'][0]

    with netCDF4.Dataset(input_filename) as nc, netCDF4.Dataset(output_filename, 'w') as out:
        lon = nc.variables['lon'][:]
        lat = nc.variables['lat'][:]
        rlon = nc.variables['rlon'][:]
        rlat = nc.variables['rlat'][:]

        print(output_filename)
        sys.stdout.flush()
        sc.io.init_netcdf(out, rlon, rlat, lon, lat, attrs=attrs)

        T = nc.variables['T'][0]
        p = nc.variables['P'][0]
        ps = nc.variables['PS'][0]

        # specific humidity
        q = nc.variables['QV'][0]

        clct = nc.variables['CLCT'][0]
        clct_attrs = get_attrs(nc.variables['CLCT'])

        # write to netCDF 
        sc.io.append_variable(out, 'CLCT', clct, attrs=clct_attrs)

        # mass of dry air (kg/m2)
        mair = c2sm.conv.calculate_mair(p, ps, h)

        for tracer in tracers:
            xm = nc.variables[tracer][0]
            attrs = get_attrs(nc.variables[tracer])


            if tracer.startswith('CO2_'):
                if do_ground:
                    ground = c2sm.conv.xm_to_xv(xm[-1], 'CO2')
                column = c2sm.conv.calculate_xco2(xm, mair, q)
                column = column.astype('f4')

                if tracer == 'CO2_BG':
                    column2 = c2sm.conv.calculate_xco2(xm, mair, 0.0)
                    column2 = column2.astype('f4')

            elif tracer.startswith('NOX_'):
                if do_ground:
                    ground = c2sm.conv.xm_to_cm(xm[-1], p[-1], T[-1])
                    ground = c2sm.conv.nox_to_no2(ground)
                    ground = c2sm.conv.cm_to_xm(ground, p[-1], T[-1])
                    ground = c2sm.conv.xm_to_xv(ground, 'NO2')

                tracer = tracer.replace('NOX_', 'NO2_')

                column = c2sm.conv.calculate_vcd(xm, p, T, h, 'NOx')
                column = column.astype('f4')

            elif tracer.startswith('CO_'):
                if do_ground:
                    ground = c2sm.conv.xm_to_xv(xm[-1], 'CO')
                column = c2sm.conv.calculate_vcd(xm, p, T, h, 'CO')
                column = column.astype('f4')
            else:
                raise ValueError('Tracers "%s" not handled.' % tracer)


            if tracer.startswith('CO2_'):
                attrs['standard_name'] = attrs['standard_name'].replace(
                        'CO2_mass_fraction', 'CO2_column-averaged_dry-air_mole_fraction')
                attrs['long_name'] = attrs['long_name'].replace(
                        'CO2_mass_fraction', 'CO2_column-averaged_dry-air_mole_fraction')
                attrs['units'] = 'ppm'

                if tracer == 'CO2_BG':
                    attrs2 = attrs.copy()
                    attrs2['long_name'] = attrs2['long_name'].replace('dry-air', 'moist-air')
                    attrs2['standard_name'] = attrs2['standard_name'].replace('dry-air', 'moist-air')


            elif tracer.startswith('NO2_'):
                attrs['standard_name'] = attrs['standard_name'].replace(
                        'NOX_mass_fraction', 'NO2_vertical_column_density')
                attrs['long_name'] = attrs['long_name'].replace(
                        'NOX_mass_fraction', 'NO2_vertical_column_density')
                attrs['units'] = 'molecules cm-2'

            elif tracer.startswith('CO_'):
                attrs['standard_name'] = attrs['standard_name'].replace(
                        'CO_mass_fraction', 'CO_vertical_column_density')
                attrs['long_name'] = attrs['long_name'].replace(
                        'CO_mass_fraction', 'CO_vertical_column_density')
                attrs['units'] = 'molecules cm-2'



            # write ground concentrations
            if do_ground:
                sc.io.append_variable(
                    out, '%s_ground' % tracer,
                    ground, attrs=None
                )

            # write columns
            if tracer.startswith('CO2'):
                sc.io.append_variable(
                    out,
                    'X%s' % tracer,
                    column, attrs=attrs
                )
                if tracer == 'CO2_BG':
                    sc.io.append_variable(
                        out,
                        'Y%s' % tracer,
                        column2, attrs=attrs2
                    )
            else:
                sc.io.append_variable(
                    out,
                    '%s' % tracer,
                    column, attrs=attrs
                )



def create_multi_input(args):
    create_2d_fields(*args, overwrite=True)

    #output_filename = args[-1]
    #create_2d_field_maps.main(output_filename)


def create_args(starttime, hstart, hstop, output_path=None, structure='flat'):
    """\
    startime (datetime)
    hstart (hour of simulation start)
    hstop (hour simulation stop)
    """
    if starttime == datetime(2015,1,1):
        output_path = config.path_2d_fields_01
    elif starttime == datetime(2015,7,1):
        output_path = config.path_2d_fields_07
    else:
        raise ValueError


    args = []

    if isinstance(hstart, datetime):
        start = hstart
    else:
        start = starttime + timedelta(hours=int(hstart))

    if isinstance(hstart, datetime):
        end = hstop
    else:
        end = starttime + timedelta(hours=int(hstop))

    if structure == 'restart':
        input_path = os.path.join(
            config.ROOT, 'mother',
            '%s_%d_%d' % (starttime.strftime('%Y%m%d%H'), hstart, hstop),
            'cosmo', 'output'
            'cosmo_output'
        )
    else:
        input_path = os.path.join(config.ROOT, starttime.strftime('cosmo_%Y-%m'))

    print(start, end)
    for time in sc.misc.iter_times(start, end, timedelta(hours=1), exclude_end=True):
        args.append([
            os.path.join(input_path, time.strftime('lffd%Y%m%d%H.nc')),
            os.path.join(input_path, starttime.strftime('lffd%Y%m%d%Hc.nc')),
            os.path.join(output_path, time.strftime('cosmo_2d_%Y%m%d%H.nc'))
        ])

    return args




if __name__ == '__main__':

    # create arguments using nested COSMO paths from restart runs
    starttime = datetime.strptime(sys.argv[1], '%Y-%m-%d')

    try:
        start = int(sys.argv[2])
    except ValueError:
        start = datetime.strptime(sys.argv[2], '%Y-%m-%d')

    try:
        stop = int(sys.argv[3])
    except ValueError:
        stop = datetime.strptime(sys.argv[3], '%Y-%m-%d')

    args = create_args(starttime, start, stop, structure='flat')

    os.makedirs(config.path_2d_fields_01, exist_ok=True)
    os.makedirs(config.path_2d_fields_07, exist_ok=True)

    #for arg in args:
    #    create_2d_fields(*arg, overwrite=False)

    # create pool for multiprocessing
    pool = Pool(processes=24)
    try:
        pool.map(create_multi_input, args)
    except:
        print('ERROR-IN-SCRIPT')
        raise




