# -*- 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

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']['delayed']
DELAY_DATE = DATES['delay']['start_date']

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, "archive")


def generate_median_file(selected_date, only_interp=False, overwrite=False):
    """ Generates the median of stored data models """

    from utils import get_currdate_tstep

    filepaths = []
    cdo_date = datetime.strptime(selected_date, "%Y%m%d").strftime("%Y-%m-%d")
    current_time_before = (DELAY_DATE and datetime.strptime(selected_date,
        "%Y%m%d") < datetime.strptime(DELAY_DATE, "%Y%m%d"))
    start_hour = ((not current_time_before and DELAY) or (current_time_before
        and not DELAY)) and '12' or '00'

    if DEBUG:
        print("CDO DATE", cdo_date)

    for model in MODELS:
        if DEBUG:
            print("MODEL", model)

        model_start = MODELS[model]['start']
        model_start_before = ('start_before' in MODELS[model]) and MODELS[model]['start_before'] or False
        print(f"""
        *******************************
        START: {model_start}, BEFORE: {model_start_before}, CURR_BEFORE: {current_time_before}
        """)
        model_date, _, tsteps = get_currdate_tstep(model_start, model_start_before, current_time_before, DELAY, selected_date)
        print(f"""
        *******************************
        {model_date}, {start_hour}, {tsteps}
        """)

        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 os.path.exists(interp_filepath) and overwrite is False:
            continue

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

        if DEBUG:
            print('MODEL', model, filepath, filepath_interp, model_date, 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)
        if only_interp:
            continue
        time.sleep(1)
        if DEBUG:
            print('CDO SPLIT SCRIPT', split_script)
        subprocess.call(split_script, shell=True)

        filepaths.append(interp_filepath)

    if not only_interp:

        # median_start_hour = "{:02d}".format(current_time_before and MODELS['median']['start_before'] or MODELS['median']['start'])
        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__":
    from os import path
    sys.path.append(os.path.dirname(os.path.abspath('.')))
    if len(sys.argv) > 1:
        generate_median_file(sys.argv[1])
    else:
        generate_median_file(datetime.date.now().strftime('%Y%m%d'))
