"""One chunk x one batch of ensemble members of a round 7 ensemble run.

Round 5's run_ens.py for the AdaPBS gdirs: the L4 tree on v8
(lmr_gdirs_v8_202610/adapbs/<tree>), whose settings.yml already holds the
transient AdaPBS calibration (melt_f, prcp_fac, temp_bias of the top-weighted
particle, calibrated on HIST ens010 with LAPrec precipitation). Nothing is put
back at run time. For each member: write its climate into the gdir, run from
R7E_YS, compile the chunk output. Only the forcing changes between members.
Do not run this directly - see run_ens_slurm.sh and submit_ens.sh.

R7E_GDIR_URL       base url of the L4 tree
R7E_PARAMS_FILE    params.cfg of the v8 experiment under the tree
R7E_BORDER
R7E_CHUNK_IDX, R7E_CHUNK_SIZE
R7E_YS             start year of the runs
R7E_MEMBERS        file with one member per line: "hist 010" or "nat 001"
R7E_BATCH, R7E_BATCH_SIZE   which members of the file this task runs
R7E_HIST_URL       server of the HIST forcing: ensemble/tp_monthly_ens### (T) and
                   ensemble/p_monthly_ens### (LAPrec P), elev.nc
R7E_NAT_URL        server of the NAT precip: ensemble/p_monthly_ens###, elev.nc
R7E_NAT_T_URL      round with the NAT temperatures (ensemble/tp_nat_monthly_ens###)
OGGM_WORKDIR, OGGM_OUTDIR

Output: $OGGM_OUTDIR/<run_id>/<run_id>_reg11_chunk###.nc, with run_id
lmr_mira_hist_ens### or lmr_mira_nat_ens### (round3's names).
"""
import os
import logging

import geopandas as gpd

import oggm.cfg as cfg
from oggm import utils, workflow, tasks

import lmr_forcing

E = os.environ
GDIR_URL = E['R7E_GDIR_URL']
BORDER = int(E['R7E_BORDER'])
CHUNK_IDX = int(E['R7E_CHUNK_IDX'])
CHUNK_SIZE = int(E['R7E_CHUNK_SIZE'])
YS = int(E['R7E_YS'])
BATCH = int(E['R7E_BATCH'])
BATCH_SIZE = int(E['R7E_BATCH_SIZE'])
RGI_REG = '11'

cfg.initialize(file=E['R7E_PARAMS_FILE'], logging_level='WORKFLOW')
cfg.PARAMS['border'] = BORDER
cfg.PATHS['working_dir'] = utils.mkdir(E['OGGM_WORKDIR'])
OUTPUT_DIR = utils.mkdir(E['OGGM_OUTDIR'])
cfg.PARAMS['use_multiprocessing'] = True
cfg.PARAMS['continue_on_error'] = True
# round5's variables
cfg.PARAMS['store_diagnostic_variables'] = ['volume', 'volume_bsl', 'area',
                                            'area_min_h', 'length']
cfg.PARAMS['min_ice_thick_for_area'] = 2.
cfg.PARAMS['min_ice_thick_for_length'] = 2.
workflow.init_mp_pool(True)
log = logging.getLogger(__name__)

with open(E['R7E_MEMBERS']) as f:
    members = [line.split() for line in f if line.strip()]
members = members[BATCH * BATCH_SIZE:(BATCH + 1) * BATCH_SIZE]
if not members:
    log.workflow(f'Batch {BATCH} is empty, nothing to do.')
    raise SystemExit(0)

rgidf = gpd.read_file(utils.get_rgi_region_file(RGI_REG, version='62'))
rgidf = workflow.get_rgi_chunk(rgidf, CHUNK_IDX, chunk_size=CHUNK_SIZE)
rgidf = rgidf.loc[rgidf.O2Region == '1']
if len(rgidf) == 0:
    log.workflow(f'Chunk {CHUNK_IDX} is empty, nothing to do.')
    raise SystemExit(0)
log.workflow(f'Chunk {CHUNK_IDX}: {len(rgidf)} glaciers, batch {BATCH}: '
             f'{len(members)} members ({members[0]} ... {members[-1]}), '
             f'runs from {YS}, border {BORDER}, gdirs {GDIR_URL}')

# The L4 is a delta: OGGM fetches L1-L3 from the parents it names
gdirs = workflow.init_glacier_directories(rgidf, from_prepro_level=4,
                                          prepro_base_url=GDIR_URL,
                                          prepro_rgi_version='62',
                                          prepro_border=BORDER)
n_ada = sum(1 for gd in gdirs if os.path.exists(gd.get_filepath('settings'))
            and 'adapbs_status' in gd.settings)
log.workflow(f'AdaPBS settings in {n_ada} of {len(gdirs)} gdirs')

for kind, num in members:
    num = int(num)
    run_id = f'lmr_mira_{kind}_ens{num:03d}'
    if kind == 'hist':
        kw = dict(server=E['R7E_HIST_URL'], temp_var='tp_monthly',
                  temp_server=E['R7E_HIST_URL'])
    elif kind == 'nat':
        kw = dict(server=E['R7E_NAT_URL'], temp_var='tp_nat_monthly',
                  temp_server=E['R7E_NAT_T_URL'])
    else:
        raise ValueError(f'Unknown member kind: {kind}')
    log.workflow(f'Running {run_id}, chunk {CHUNK_IDX}')
    suffix = f'_{kind}_ens{num:03d}'
    workflow.execute_entity_task(lmr_forcing.process_mira_data, gdirs,
                                 ens_member=num, output_filesuffix=suffix, **kw)
    workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                                 climate_filename='climate_historical',
                                 climate_input_filesuffix=suffix,
                                 output_filesuffix=run_id,
                                 return_value=False,
                                 ys=YS)
    out_dir = utils.mkdir(os.path.join(OUTPUT_DIR, run_id))
    utils.compile_run_output(gdirs, input_filesuffix=run_id,
                             path=os.path.join(out_dir, f'{run_id}_reg{RGI_REG}'
                                               f'_chunk{CHUNK_IDX:03d}.nc'))
    # Free the local disk: this member's climate and model output files
    for gd in gdirs:
        for fn, fs in [('climate_historical', suffix),
                       ('model_diagnostics', run_id),
                       ('model_geometry', run_id)]:
            fp = gd.get_filepath(fn, filesuffix=fs)
            if os.path.exists(fp):
                os.remove(fp)