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

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')))
NETCDF_TEMPLATE = "{}/archive/{}{}.nc"

DEBUG = True

lonlat = "-27,60,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):

    filepaths = []

    for model in MODELS:

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

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

        print('MODEL', model)

        if MODELS[model]['start'] == 0:
            tsteps = "5/25"
        else:
            tsteps = "1/21"

        cdo_date = datetime.strptime(selected_date, "%Y%m%d").strftime("%Y-%m-%d")
        remap_script = "cdo -L -r -settaxis,"+cdo_date+",12: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', '_'))

        subprocess.call(remap_script, shell=True)
        time.sleep(1)
        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'))
