# SPDX-FileCopyrightText: 2024 European Centre for Medium-Range Weather Forecasts (ECMWF)
# SPDX-License-Identifier: Apache-2.0

# script containing all relevant code for producing ecdsAPI requests.
from acacia_s2s_toolkit import argument_output, merge_lag_ensemble
import numpy as np
import os
import eccodes as ec
import xarray as xr
from datetime import datetime, timedelta
import subprocess
import cdsapi
import glob
from pathlib import Path

def get_ecds_client():
    os.environ["CDSAPI_RC"] = os.path.expanduser("~/.cdsapirc.ecds")
    return cdsapi.Client()

def cleanup_patterns(*patterns):
    for pattern in patterns:
        for path_str in glob.glob(pattern):
            path = Path(path_str)
            if path.is_file():
                path.unlink()

def create_initial_ecdsAPI_request(fcdate,grid,area,origin,webapi_param,leadtimes):
    request_dict = {
            "class": "s2",
            "date": f"{fcdate}",
            "expver": "prod",
            "grid": f"{grid}",
            "levtype": "sfc",
            "origin": f"{origin}",
            "param": f"{webapi_param}",
            "step": f"{leadtimes}",
            "time": "00:00:00",
            "stream": "enfo",
            "type": "cf"
            }

    # add area component
    area_formatted = '/'.join(str(x) for x in area)
    request_dict['area'] = f"{area_formatted}"

    return request_dict

def request_forecast(fcdate,origin,grid,variable,area,data_format,webapi_param,leadtime_hour,leveltype,filename,plevs,fc_enslags):
    # import registration detailis
    client = get_ecds_client()
    dataset = "s2s-forecasts"
    # to enable lagged ensemble, loop through requested ensembles
    for lag in np.atleast_1d(fc_enslags):
        lag = int(lag)
        leadtimes, convert_fcdate = argument_output.output_formatted_leadtimes(leadtime_hour,fcdate,variable,origin,lag=lag,fc_enslags=fc_enslags) # you need to know new fcdate, but also appropriate leadtimes so all are forecasting same period
        # create initial control request
        request_dict = create_initial_ecdsAPI_request(convert_fcdate,grid,area,origin,webapi_param,leadtimes)

        # change components of request based on level type, and grid
        # if grid doesn't equal '1.5/1.5', add 'repres' dictionary item which sets the requested representation, in this case, 'll'=latitude/longitude.
        if grid != '1.5/1.5':
            # add repres
            request_dict['repres'] = 'll'

        # if a pressure level type is selected, just need to change levtype and add list of pressure levels.
        if leveltype == 'pressure':
            request_dict['levtype'] = 'pl'
            # convert plevs
            plevels = '/'.join(str(x) for x in plevs)
            request_dict['levelist'] = f"{plevels}"

        # specific change needed for pv
        if variable == 'pv':
            request_dict['levtype'] = 'pt'
            request_dict['levelist'] = '320'

        # retrieve the control forecast
        target=f'{filename}_control_{lag}'
        client.retrieve(dataset,request_dict,target)

        # then download perturbed. change type of forecast, add number of ensemble members, and change target filename
        request_dict['type'] = 'pf'
        # add model number (will not be needed for ECDSapi)
        #num_pert_fcs = argument_output.get_single_parameter(origin,fcdate,'fcNumEns')
        #pert_fcs = '/'.join(str(x) for x in np.arange(1,num_pert_fcs+1))
        #request_dict['number'] = f"{pert_fcs}"

        target = f'{filename}_perturbed_{lag}' 
        client.retrieve(dataset,request_dict,target)

        # once requesting control and perturbed forecast, combine the two.
        # set forecast type in control to pf (perturbed forecast).
        set_cf_to_pf(f'{filename}_control_{lag}',f'{filename}_control2_{lag}')
       
        # merge both control and perturbed forecast
        cmd = ["cdo","-O","merge",f"{filename}_control2_{lag}",f"{filename}_perturbed_{lag}",f"{filename}_allens_{lag}",]
        result = subprocess.run(cmd,stdout=subprocess.DEVNULL,stderr=subprocess.PIPE,text=True)
        if result.returncode != 0:
            print(result.stderr)
            raise subprocess.CalledProcessError(result.returncode, cmd, stderr=result.stderr)

    # create new 'member' dimension based on same date. For instance, 5 members per date and three initialisations used
    # smae process following even with one forecast initialisation date to ensure same structure for all output. 
    combined_forecast = merge_lag_ensemble.merge_all_ens_members(f'{filename}',leveltype)
    combined_forecast.to_netcdf(f'{filename}.nc')

    # remove previous files  
    # remove previous files
    cleanup_patterns(f"{filename}_control*",f"{filename}_perturbed*",f"{filename}_allens*",)  

def request_hindcast(fcdate,origin,grid,variable,area,data_format,webapi_param,leadtime_hour,leveltype,filename,plevs,rf_enslags,rf_years,fc_time=True):
    # to enable lagged ensemble, loop through requested ensembles
    # import registration details
    client = get_ecds_client()
    dataset = "s2s-reforecasts" 
  
    print (rf_enslags)
    for lag in np.atleast_1d(rf_enslags):
        lag = int(lag)
        # convert fcdate
        lagged_fcdate = datetime.strptime(fcdate, '%Y%m%d')+timedelta(days=lag) # work out what the lagged fcdate is.
        convert_fcdate = lagged_fcdate.strftime('%Y%m%d') # convert that date to YYYYMMDD format

        rf_model_date, rfyears = argument_output.check_and_output_all_hc_arguments(variable,origin,convert_fcdate,rf_years) # get the reforecast model date version plus the number of reforecast years

        leadtimes, lag_fcdate_dashed_format = argument_output.output_formatted_leadtimes(leadtime_hour,convert_fcdate,variable,origin) # get the appropriate leadtimes plus the convert fcdate.
        print (leadtimes)

        # create initial control request
        request_dict = create_initial_ecdsAPI_request(lag_fcdate_dashed_format,grid,area,origin,webapi_param,leadtimes)

        # use correct reforecast model date
        request_dict['date'] = f"{rf_model_date}"

        # download reforecast, so change stream
        request_dict['stream'] = f"enfh"

        # create list of hdates
        #hdates = argument_output.create_reforecast_dates(rfyears,convert_fcdate)
        #request_dict['hdate']=f"{hdates}"
        #print (request_dict)

        request_dict['hyear'] = [str(year) for year in rfyears]
        request_dict['hmonth'] = convert_fcdate[4:6]
        request_dict['hday']   = convert_fcdate[6:]

        # change components of request based on level type, and grid
        # if grid doesn't equal '1.5/1.5', add 'repres' dictionary item which sets the requested representation, in this case, 'll'=latitude/longitude.
        if grid != '1.5/1.5':
            # add repres
            request_dict['repres'] = 'll'

        # if a pressure level type is selected, just need to change levtype and add list of pressure levels.
        if leveltype == 'pressure':
            request_dict['levtype'] = 'pl'
            # convert plevs
            plevels = '/'.join(str(x) for x in plevs)
            request_dict['levelist'] = f"{plevels}"

        # specific change needed for pv
        if variable == 'pv':
            request_dict['levtype'] = 'pt'
            request_dict['levelist'] = '320'

        # retrieve the control forecast
        target=f'{filename}_control_{lag}'
        client.retrieve(dataset,request_dict,target)

        # then download perturbed. change type of forecast, add number of ensemble members, and change target filename
        request_dict['type'] = 'pf'
        # add model number (will not be needed for ECDSapi)
        #num_pert_hcs = argument_output.get_single_parameter(origin,convert_fcdate,'rfNumEns')
        #pert_hcs = '/'.join(str(x) for x in np.arange(1,num_pert_hcs+1))
        #request_dict['number'] = f"{pert_hcs}"

        target=f"{filename}_perturbed_{lag}"
        client.retrieve(dataset,request_dict,target)

        # once requesting control and perturbed forecast, combine the two.
        # set forecast type in control to pf (perturbed forecast).
        set_cf_to_pf(f'{filename}_control_{lag}',f'{filename}_control2_{lag}')

        # Merge control and perturbed forecast members into one file.
        # We suppress routine CDO warning noise, but still raise on real failures.
        cmd = ["cdo","-O","merge",f"{filename}_control2_{lag}",f"{filename}_perturbed_{lag}",f"{filename}_allens_{lag}",]
        result = subprocess.run(cmd,stdout=subprocess.DEVNULL,stderr=subprocess.PIPE,text=True)
        if result.returncode != 0:
            print(result.stderr)
            raise subprocess.CalledProcessError(result.returncode, cmd, stderr=result.stderr)

        # shift the time so all reforecasts have the same time values
        if fc_time:
            shift_day_value = lag*-1
        else:
            shift_day_value = 0 
        # only shift the time, if you want a 'forecast-based' time.
        rf_shifttime(f'{filename}_allens_{lag}',f'{filename}_timeshifted_allens_{lag}',shift_days=shift_day_value)

    # create new 'member' dimension based on same date. For instance, 5 members per date and three initialisations used
    # same process following even with one forecast initialisation date to ensure same structure for all output. 
    combined_forecast = merge_lag_ensemble.merge_all_ens_hindcasts(f'{filename}_timeshifted',leveltype)
    combined_forecast.to_netcdf(f'{filename}.nc')

    # remove previous files  
    cleanup_patterns(f"{filename}_control*",f"{filename}_perturbed*",f"{filename}_allens*",)

def rf_shifttime(fn,output_fn,shift_days=0):
    orig_hc = xr.open_dataset(fn,engine='cfgrib')
    if np.size(orig_hc['step'].values) > 1: # different option depending on size of step variable.
        data = [] # loop through each time point (forecast initialisation). only save the valid forecast time.
        for time_pt in orig_hc['time']:
            select = orig_hc.sel(time=time_pt)
            select = select.assign_coords(step=select['valid_time'])
            select = select.drop_vars(["valid_time","time"])
            select = select.rename({'step':'valid_time'})
            data.append(select)
        orig_hc_new = xr.concat(data, dim='valid_time')
        #orig_hc_new = xr.merge(data)
    else:
        orig_hc_new = orig_hc.swap_dims({'time':'valid_time'})
        orig_hc_new = orig_hc_new.drop_vars("time")

    orig_hc_new = orig_hc_new.assign_coords(
            valid_time=orig_hc_new.valid_time+np.timedelta64(shift_days,'D')
            )
    orig_hc_new = orig_hc_new.rename({'valid_time':'time'})
    orig_hc_new.to_netcdf(output_fn)

def set_cf_to_pf(input_file, output_file):
    # Open input file
    with open(input_file, 'rb') as fin,open(output_file,'wb') as fout:
        while True:
            gid = ec.codes_grib_new_from_file(fin)
            if gid is None:
                break
            # Set type from 'cf' to 'pf'
            ec.codes_set(gid, 'type', 'pf')
            ec.codes_write(gid,fout)
            ec.codes_release(gid)
