# -*- coding: utf-8 -*-
# Copyright 2016 Earth Sciences Department, BSC-CNS
#

"""calc_median module

This module provide calculation of MEDIAN.

"""

import numpy as np
import xarray as xr
import json
import os
import os.path
import sys
import subprocess
import time
from datetime import datetime
from datetime import timedelta
from operator import itemgetter

np.set_printoptions(precision=2)

DIR_PATH = os.path.abspath(os.path.dirname(__file__))
MODELS = json.load(open(os.path.join(DIR_PATH, '../conf/models.json')))
DATES = json.load(open(os.path.join(DIR_PATH, '../conf/dates.json')))
NETCDF_TEMPLATE = "{}/archive/{}{}.nc"

DEBUG = True
DELAY = DATES['delay']

lonlat = "-27,65,0,65"
main_output_path = "/data/products/median/"
interpolated_path = os.path.join(main_output_path, "interpolated")
# archive_path = os.path.join(main_output_path, "tmp")
archive_path = os.path.join(main_output_path, "archive")
# selected_date = "20220515"


def generate_median_file(selected_date):
    """ Generates the median of stored data models """

    model_conf = {
        True: {
            0: {
                'tsteps': "5/25",
                'model_date': selected_date,
                'start_hour': "12"
                },
            12: {
                'tsteps': "1/21",
                'model_date': selected_date,
                'start_hour': "12"
                },
            },
        False: {
            0: {
                'tsteps': "1/25",
                'model_date': selected_date,
                'start_hour': "00"
                },
            12: {
                'tsteps': "5/29",
                'model_date': (datetime.strptime(selected_date, "%Y%m%d") -
                    timedelta(days=1)).strftime("%Y%m%d"),
                'start_hour': "00"
                },
            }
        }

    filepaths = []
    cdo_date = datetime.strptime(selected_date, "%Y%m%d").strftime("%Y-%m-%d")
    if DEBUG:
        print("CDO DATE", cdo_date)

    for model in MODELS:

        model_start = MODELS[model]['start']
        cur_dict = model_conf[DELAY][model_start]
        tsteps, model_date, start_hour = itemgetter(*cur_dict.keys())(cur_dict)

        filepath_interp = NETCDF_TEMPLATE.format(
            MODELS[model]['path'],
            selected_date,
            MODELS[model]['template']
        )
        
        filepath = NETCDF_TEMPLATE.format(
            MODELS[model]['path'],
            model_date,
            MODELS[model]['template']
        )
        
        interp_filepath = os.path.join(interpolated_path, os.path.basename(filepath_interp))

        if model == 'median' or not os.path.exists(filepath):
            continue

        if DEBUG:
            print('MODEL', model, filepath, model_date, filepath_interp, selected_date)


        remap_script = "cdo -L -r -settaxis," +\
            cdo_date + "," + start_hour + ":00:00,3hours" +\
            " -selvar,OD550_DUST,SCONC_DUST" +\
            " -chname,od550_dust,OD550_DUST" +\
            " -chname,sconc_dust,SCONC_DUST" +\
            " -seltimestep," + tsteps +\
            " -sellonlatbox," + lonlat +\
            " -remapbil,global_0.5 " + filepath +\
            " " + interp_filepath
        split_script = "cdo splitname {} {}".format(interp_filepath, interp_filepath.replace('.nc', '_'))


        if DEBUG:
            print('CDO REMAP SCRIPT', remap_script)
        subprocess.call(remap_script, shell=True)
        time.sleep(1)
        if DEBUG:
            print('CDO SPLIT SCRIPT', split_script)
        subprocess.call(split_script, shell=True)

        filepaths.append(interp_filepath)

    sconc_ds = xr.open_mfdataset(
            "{}/{}*SCONC_DUST.nc".format(interpolated_path, selected_date),
            concat_dim='model', combine='nested')
    od550_ds = xr.open_mfdataset(
            "{}/{}*OD550_DUST.nc".format(interpolated_path, selected_date),
            concat_dim='model', combine='nested')
    sconc_dust = sconc_ds['SCONC_DUST'].median(dim='model', skipna=True, keep_attrs=True)
    od550_dust = od550_ds['OD550_DUST'].median(dim='model', skipna=True, keep_attrs=True)
    newds = xr.Dataset({ 'OD550_DUST': od550_dust, 'SCONC_DUST': sconc_dust }).astype(np.float32)
    print(newds)
    newds = newds.rename_dims(
            {'lon': 'longitude', 'lat': 'latitude'}).rename_vars(
            {'lon': 'longitude', 'lat': 'latitude'})
    newds['longitude'] = newds['longitude'].astype(np.float32)
    newds['latitude'] = newds['latitude'].astype(np.float32)
    newds.to_netcdf('{}/{}_3H_MEDIAN.nc'.format(archive_path, selected_date),
            encoding={ 'time' : { 'units' : 'hours since '+cdo_date+' 00:00:00' } })

    # clean interpolated directory
    for fname in filepaths:
        # os.remove(fname)
        print("to be removed", fname)


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        generate_median_file(sys.argv[1])
    else:
        generate_median_file(datetime.date.now().strftime('%Y%m%d'))
