#!/usr/bin/env python
"""Unstack LMR (year, month) netCDF files into (lon, lat, time) files.

The files we receive store monthly data as a (year, month) grid (with
``year`` documented as "Calendar year", starting at 0) instead of a proper
time axis, use inconsistent variable names, and need small unit fixes
before they are useful. This script converts a whole directory of such
files into CF-ish (lon, lat, time) netCDFs in one go.

Each input variable is handled according to the ``VARIABLES`` /
``STATIC_VARIABLES`` tables below. To support a new file, just add an entry
there -- no other code needs to change. Files named ``<base>_ens###.nc``
are treated as ensemble members and written to an ``ensemble/`` subfolder.

Usage:
    python unstack.py orig/ForFabien_Pilot_nc_Round2 unstacked/ForFabien_Pilot_nc_Round2
"""

import argparse
import glob
import os
import re

import cftime
import numpy as np
import xarray as xr
from oggm import cfg

ENS_RE = re.compile(r'^(?P<base>.+)_ens(?P<num>\d+)\.nc$')


def temp_c_to_k(da):
    da = da + 273.15
    da.attrs['units'] = 'K'
    return da


def precip_m_to_flux(da):
    """P_monthly is given in meters per month -> kg m-2 s-1."""
    da = da * 1000  # m -> mm == kg m-2 per month
    dimo = np.array([cfg.DAYS_IN_MONTH[m - 1] for m in da['time.month']])
    da = da / (dimo * (60 * 60 * 24))
    da.attrs['units'] = 'kg m-2 s-1'
    return da


# input variable name -> (output variable name, output file basename, transform)
VARIABLES = {
    'P_monthly': ('pr', 'p_monthly', precip_m_to_flux),
    'Tp_monthly': ('tas', 'tp_monthly', temp_c_to_k),
    'TpNAT_monthly': ('tas', 'tp_nat_monthly', temp_c_to_k),
}

# input variable name -> (output variable name, output file basename)
STATIC_VARIABLES = {
    'elevation': ('elevation', 'elev'),
}

DIM_RENAMES = {'longitude': 'lon', 'latitude': 'lat'}


def stack_to_time(ds, varname, calendar='noleap'):
    """Turn a (year, month) DataArray into a (lat, lon, time) one.

    Year 0 is dropped if present: it is not a valid calendar year and
    causes errors downstream.
    """
    ds = ds.sel(year=ds.year[ds.year != 0])
    da = ds[varname]

    year_grid, month_grid = xr.broadcast(ds.year, ds.month)
    years = year_grid.values.ravel()
    months = month_grid.values.ravel()
    dates = np.array([
        cftime.DatetimeNoLeap(int(y), int(m), 1) for y, m in zip(years, months)
    ])

    da = da.stack(time=('year', 'month'))
    da = da.drop_vars(['time', 'year', 'month'])
    da['time'] = dates
    return da.transpose('lat', 'lon', 'time').sortby('time')


def convert_timeseries_file(path, out_dir):
    with xr.open_dataset(path) as ds:
        ds = ds.load()

    varname = list(ds.data_vars)[0]
    if varname not in VARIABLES:
        print(f'  ! no mapping for variable {varname!r}, skipping {path}')
        return

    out_var, out_base, transform = VARIABLES[varname]
    da = transform(stack_to_time(ds, varname))
    da.name = out_var
    out_ds = xr.Dataset({out_var: da})

    m = ENS_RE.match(os.path.basename(path))
    if m:
        out_path = os.path.join(out_dir, 'ensemble', f'{out_base}_ens{m.group("num")}.nc')
    else:
        out_path = os.path.join(out_dir, f'{out_base}.nc')

    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    out_ds.to_netcdf(out_path)
    print(f'  wrote {out_path}')


def convert_static_file(path, out_dir):
    with xr.open_dataset(path) as ds:
        ds = ds.load()

    varname = list(ds.data_vars)[0]
    if varname not in STATIC_VARIABLES:
        print(f'  ! no mapping for variable {varname!r}, skipping {path}')
        return

    out_var, out_base = STATIC_VARIABLES[varname]
    rename = {k: v for k, v in DIM_RENAMES.items() if k in ds.coords}
    if rename:
        ds = ds.rename(rename)
    if varname != out_var:
        ds = ds.rename({varname: out_var})

    out_path = os.path.join(out_dir, f'{out_base}.nc')
    os.makedirs(out_dir, exist_ok=True)
    ds.to_netcdf(out_path)
    print(f'  wrote {out_path}')


def process_dir(in_dir, out_dir):
    files = sorted(glob.glob(os.path.join(in_dir, '*.nc')))
    print(f'Found {len(files)} .nc files in {in_dir}')
    for path in files:
        with xr.open_dataset(path) as ds:
            is_timeseries = 'year' in ds.dims and 'month' in ds.dims
        if is_timeseries:
            convert_timeseries_file(path, out_dir)
        else:
            convert_static_file(path, out_dir)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__,
                                      formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('in_dir', help='directory with the original .nc files')
    parser.add_argument('out_dir', help='directory to write unstacked .nc files to')
    args = parser.parse_args()
    process_dir(args.in_dir, args.out_dir)