"""Merge the per-chunk run outputs of one experiment into one file.

    python merge_chunks.py <chunk_dir> <output_file>

The output has the same layout as a non-chunked compile_run_output file
(e.g. round3/output/*/lmr_mira_hist_reg11.nc), so round3's
preprocess_mira.py works on it unchanged. merge() is also used by
merge_ens.py in-process: starting a new python per member is slow on the
compute nodes (the imports come over NFS).
"""
import glob
import os
import sys

import xarray as xr

def merge(chunk_dir, out_file, verbose=True):

    files = sorted(glob.glob(os.path.join(chunk_dir, '*_chunk*.nc')))
    if not files:
        raise RuntimeError(f'No chunk files in {chunk_dir}')

    # Load each chunk into memory first (they are small): concatenating and
    # sorting lazily-opened files is very slow
    dss = []
    for f in files:
        with xr.open_dataset(f) as d:
            dss.append(d.load())
    # Only the per-glacier variables are concatenated: the others (time,
    # calendar_year, ...) are the same in every chunk
    ds = xr.concat(dss, dim='rgi_id', data_vars='minimal', coords='minimal',
                   compat='override', join='outer')
    ds = ds.sortby('rgi_id')
    assert ds.rgi_id.to_index().is_unique, 'Same glacier in several chunks?'

    n_nan = int(ds.volume.isel(time=0).isnull().sum())
    if verbose:
        print(f'{len(files)} chunks, {ds.sizes["rgi_id"]} glaciers '
          f'({n_nan} without output), years {int(ds.time[0])}-{int(ds.time[-1])}')

    os.makedirs(os.path.dirname(os.path.abspath(out_file)), exist_ok=True)
    tmp_file = out_file + '.tmp'
    ds.to_netcdf(tmp_file)
    os.replace(tmp_file, out_file)
    if verbose:
        print(f'Written: {out_file}')

    # The per-glacier dynamic calibration tables, if any (DYN_CALIB runs)
    csvs = sorted(glob.glob(os.path.join(chunk_dir, 'dyn_calib_*_chunk*.csv')))
    if csvs:
        import pandas as pd
        df = pd.concat([pd.read_csv(f, index_col=0) for f in csvs]).sort_index()
        out_csv = os.path.join(os.path.dirname(os.path.abspath(out_file)),
                               'dyn_calib_reg11.csv')
        df.to_csv(out_csv)
        if verbose:
            print(f'Written: {out_csv} ({len(df)} glaciers, '
                  f'{df.used_spinup_option.value_counts().to_dict()})')


if __name__ == '__main__':
    merge(*sys.argv[1:3])