"""One chunk x one batch of NAT members of the round 6 permutation runs.

Reads the L4 state tree (full gdirs: dyn melt_f in settings.yml, the 1850
states as model_flowlines_ens###_1850). For each NAT member j of the batch,
writes its climate once (newnat T, own P, as round5/run_ens.py), then runs
every state i from 1850 to the end of that climate. Do not run this
directly - see run_perm_slurm.sh and submit_perm.sh.

R6_STATES_URL      the L4 state tree
R6_PARAMS_FILE, R6_BORDER   the v7 experiment's
R6_CHUNK_IDX, R6_CHUNK_SIZE
R6_YS              1850
R6_STATES          e.g. "123 023 117 054 178 010 041"
R6_MEMBERS         file with one NAT member number per line
R6_BATCH, R6_BATCH_SIZE
R6_NAT_T_URL, R6_P_URL
R6_TEST_IDS        optional: only these glaciers
OGGM_WORKDIR, OGGM_OUTDIR

Output: $OGGM_OUTDIR/<run_id>/<run_id>_reg11_chunk###.nc with
run_id = perm_state###_nat_ens###.
"""
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
STATES_URL = E['R6_STATES_URL']
BORDER = int(E['R6_BORDER'])
CHUNK_IDX = int(E['R6_CHUNK_IDX'])
CHUNK_SIZE = int(E['R6_CHUNK_SIZE'])
YS = int(E['R6_YS'])
STATES = [int(s) for s in E['R6_STATES'].split()]
BATCH = int(E['R6_BATCH'])
BATCH_SIZE = int(E['R6_BATCH_SIZE'])
TEST_IDS = E.get('R6_TEST_IDS', '').split()
RGI_REG = '11'

cfg.initialize(file=E['R6_PARAMS_FILE'], logging_level='WORKFLOW')
cfg.PARAMS['border'] = BORDER
cfg.PARAMS['use_multiprocessing'] = True
cfg.PARAMS['continue_on_error'] = True
# As the round5 ensemble (run_ens.py)
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.
cfg.PATHS['working_dir'] = E['OGGM_WORKDIR']
OUTPUT_DIR = E['OGGM_OUTDIR']
utils.mkdir(OUTPUT_DIR)
workflow.init_mp_pool(True)
log = logging.getLogger(__name__)

with open(E['R6_MEMBERS']) as f:
    members = [int(line) 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 TEST_IDS:
    rgidf = rgidf.loc[rgidf.RGIId.isin(TEST_IDS)]
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}: NAT '
             f'{members[0]:03d}-{members[-1]:03d} ({len(members)}) x states {STATES}')

gdirs = workflow.init_glacier_directories(rgidf, from_prepro_level=4,
                                          prepro_base_url=STATES_URL,
                                          prepro_rgi_version='62',
                                          prepro_border=BORDER)

for j in members:
    csfx = f'_nat_ens{j:03d}'
    workflow.execute_entity_task(lmr_forcing.process_mira_data, gdirs,
                                 ens_member=j, server=E['R6_P_URL'],
                                 temp_var='tp_nat_monthly',
                                 temp_server=E['R6_NAT_T_URL'],
                                 output_filesuffix=csfx)
    for i in STATES:
        run_id = f'perm_state{i:03d}_nat_ens{j:03d}'
        workflow.execute_entity_task(tasks.run_from_climate_data, gdirs, ys=YS,
                                     climate_input_filesuffix=csfx,
                                     model_flowlines_filesuffix=f'_ens{i:03d}_{YS}',
                                     output_filesuffix=run_id,
                                     store_model_geometry=False,
                                     store_fl_diagnostics=False,
                                     return_value=False)
        out_dir = os.path.join(OUTPUT_DIR, run_id)
        utils.mkdir(out_dir)
        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'))
        for gd in gdirs:
            fp = gd.get_filepath('model_diagnostics', filesuffix=run_id)
            if os.path.exists(fp):
                os.remove(fp)
    for gd in gdirs:
        fp = gd.get_filepath('climate_historical', filesuffix=csfx)
        if os.path.exists(fp):
            os.remove(fp)
    log.workflow(f'Chunk {CHUNK_IDX}: NAT ens{j:03d} done ({len(STATES)} states)')