"""One chunk of the round 6 committed-loss runs (HIST strands).

Reads the L5 of the state tree (full gdirs incl. model_geometry_hist_ens###:
each member's HIST run from its 1850 state, yearly geometry). For each
member and seed: OGGM's random mass balance model (run_random_climate) on
the member's own HIST climate (climate_historical_hist_ens###), years
R6_C_Y0 ± R6_C_HALFSIZE, each year drawn once per window
(unique_samples=True), starting from the HIST geometry of R6_C_START, for
R6_C_NYEARS years. One seed = the same year sequence for all glaciers. Do
not run this directly - see submit_commit.sh.

R6_STATES_URL      the state tree (level 5 is read)
R6_PARAMS_FILE, R6_BORDER, R6_CHUNK_IDX, R6_CHUNK_SIZE
R6_STATES          the members, e.g. "123 023 117 054 178 010 041"
R6_C_START         start year (the HIST geometry of that year), e.g. 2025
R6_C_Y0, R6_C_HALFSIZE, R6_C_NYEARS, R6_C_SEEDS
R6_TEST_IDS        optional
OGGM_WORKDIR, OGGM_OUTDIR

Output: $OGGM_OUTDIR/<run_id>/<run_id>_reg11_chunk###.nc with
run_id = commit<START>_state###_seed#.
"""
import os
import logging

import geopandas as gpd

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

E = os.environ
BORDER = int(E['R6_BORDER'])
CHUNK_IDX = int(E['R6_CHUNK_IDX'])
STATES = [int(s) for s in E['R6_STATES'].split()]
START = int(E['R6_C_START'])
Y0 = int(E['R6_C_Y0'])
HALFSIZE = int(E['R6_C_HALFSIZE'])
NYEARS = int(E['R6_C_NYEARS'])
SEEDS = [int(s) for s in E['R6_C_SEEDS'].split()]
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
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__)

if int(E.get('R6_BATCH', 0)) > 0:
    raise SystemExit(0)   # one batch only (the array has one task per chunk)

rgidf = gpd.read_file(utils.get_rgi_region_file(RGI_REG, version='62'))
rgidf = workflow.get_rgi_chunk(rgidf, CHUNK_IDX, chunk_size=int(E['R6_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, states {STATES}, seeds {SEEDS}: '
             f'random climate {Y0 - HALFSIZE}-{Y0 + HALFSIZE} (unique samples) '
             f'from {START} for {NYEARS} yr')

gdirs = workflow.init_glacier_directories(rgidf, from_prepro_level=5,
                                          prepro_base_url=E['R6_STATES_URL'],
                                          prepro_rgi_version='62',
                                          prepro_border=BORDER)

for i in STATES:
    m = f'ens{i:03d}'
    for seed in SEEDS:
        run_id = f'commit{START}_state{i:03d}_seed{seed}'
        workflow.execute_entity_task(tasks.run_random_climate, gdirs,
                                     nyears=NYEARS, y0=Y0, halfsize=HALFSIZE,
                                     seed=seed, unique_samples=True,
                                     climate_input_filesuffix=f'_hist_{m}',
                                     init_model_filesuffix=f'_hist_{m}',
                                     init_model_yr=START,
                                     output_filesuffix=run_id,
                                     store_model_geometry=False,
                                     store_fl_diagnostics=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)
    log.workflow(f'Chunk {CHUNK_IDX}: {m} done ({len(SEEDS)} seeds)')