"""Hintereisferner: lam8_jjas_dyn against WGMS mass balance and Leclercq length.

Uses the HIST ens183 run of the round 4 ensemble (the calibration member) as
it is: nothing is re-run. The v5 lam8_jjas gdir is only fetched for the
observations (gdir.get_ref_mb_data / get_ref_length_data) and for the fixed
geometry mass balance.

Mass balance, two ways (see the OGGM "plot_mass_balance" tutorial):
- transient: volume change of the run / its area, x ice density. Includes
  the geometry feedbacks.
- fixed geometry: the mass balance model on the inversion flowlines (the
  RGI-date geometry, 2003 for HEF), with its own calibration: the static one
  of the gdir (melt_f, prcp_fac, temp_bias), which matches Hugonnet 2000-2020
  on that geometry. The run has the same prcp_fac and temp_bias, but the
  dynamically calibrated melt_f, which matches Hugonnet with the transient
  geometry. So each of the two is calibrated to the same observation.

Nothing is written to the gdir: both melt_f are passed explicitly.

Length: as change relative to REF_YEAR, since the absolute length of an
elevation band flowline is not comparable to a measured one.

Time: OGGM years are calendar years here (volume at time Y = 1 Jan of Y, so
V[Y+1] - V[Y] is the balance of calendar year Y). The WGMS years are
hydrological (Oct Y-1 to Sep Y). They are compared as they are, as in the
tutorial: 3 months of offset.

Usage: python validate_hef.py   (EXP / GDIR_EXP / MEMBER from the env)
"""
import os

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr

import oggm.cfg as cfg
from oggm import utils, workflow
from oggm.core.massbalance import MultipleFlowlineMassBalance

RGI_ID = 'RGI60-11.00897'
NAME = 'Hintereisferner'
EXP = os.environ.get('EXP', 'lam8_jjas_dyn')
GDIR_EXP = os.environ.get('GDIR_EXP', 'lam8_jjas')
MEMBER = os.environ.get('MEMBER', '183')
# The other HIST members, for the length plot (ens123 left out as elsewhere)
OTHER_MEMBERS = [m for m in ['053', '124', '153', '178', '183'] if m != MEMBER]
# Length changes are relative to this year (the RGI date of HEF)
REF_YEAR = 2003

HERE = os.path.dirname(os.path.abspath(__file__))
R4_DIR = os.path.dirname(HERE)
V5_EXP_DIR = os.path.join(R4_DIR, '..', 'lmr_gdirs_v5_202610', 'exps', GDIR_EXP)
GDIR_URL = ('https://cluster.klima.uni-bremen.de/~fmaussion/runs/lmr_2025/'
            f'lmr_gdirs_v5_202610/exps/{GDIR_EXP}/')
ENS_DIR = os.path.join(R4_DIR, 'ensemble', 'output', EXP)
MELT_F_CSV = os.path.join(R4_DIR, 'sensitivity', 'output', EXP,
                          'dyn_calib_reg11.csv')
FIG_DIR = utils.mkdir(os.path.join(HERE, 'figs'))


def read_run(member):
    fp = os.path.join(ENS_DIR, f'lmr_mira_hist_ens{member}_reg11.nc')
    with xr.open_dataset(fp) as ds:
        ds = ds[['volume', 'area', 'length']].sel(rgi_id=RGI_ID).load()
    df = ds.to_dataframe()[['volume', 'area', 'length']]
    df.index = df.index.astype(int)
    return df


def stats(mod, obs):
    d = mod - obs
    return {'n': len(d), 'mean_obs': obs.mean(), 'mean_mod': mod.mean(),
            'bias': d.mean(), 'rmse': np.sqrt((d ** 2).mean()),
            'std_obs': obs.std(), 'std_mod': mod.std(),
            'r': np.corrcoef(mod, obs)[0, 1]}


cfg.initialize(file=os.path.join(V5_EXP_DIR, 'params.cfg'),
               logging_level='WARNING')
cfg.PARAMS['border'] = 160
cfg.PATHS['working_dir'] = utils.mkdir(os.path.join(HERE, '_wd'))

gdir = workflow.init_glacier_directories([RGI_ID], from_prepro_level=3,
                                         prepro_base_url=GDIR_URL,
                                         prepro_rgi_version='62',
                                         prepro_border=160)[0]

# The two melt_f: the static one of the gdir (fixed geometry), and the
# dynamically calibrated one the run used. Read only, the gdir is not changed.
melt_f_static = gdir.settings['melt_f']
mf = pd.read_csv(MELT_F_CSV, index_col=0, low_memory=False)
melt_f_dyn = float(mf.loc[RGI_ID, 'melt_f_dynamic_calibration'])
if not np.isfinite(melt_f_dyn):
    melt_f_dyn = melt_f_static
print(f'{NAME} ({RGI_ID}), rgi date {gdir.rgi_date}: melt_f {melt_f_static} '
      f'(static, fixed geometry) / {melt_f_dyn} (dyn, the run), prcp_fac '
      f'{gdir.settings["prcp_fac"]:.2f}, temp_bias '
      f'{gdir.settings["temp_bias"]:.2f}')

run = read_run(MEMBER)

# ------------------------------ mass balance -------------------------------

ref_mb = gdir.get_ref_mb_data()
mb = pd.DataFrame(index=ref_mb.index)
mb['wgms'] = ref_mb.ANNUAL_BALANCE

# Transient: as the tutorial (volume change over the area at the year's end)
smb = (run.volume.values[1:] - run.volume.values[:-1]) / run.area.values[1:]
smb = pd.Series(smb * cfg.PARAMS['ice_density'], index=run.index[:-1])
mb['transient'] = smb.reindex(mb.index)

# Fixed geometry, at the RGI date, with the static calibration
mbmod = MultipleFlowlineMassBalance(gdir, use_inversion_flowlines=True,
                                    melt_f=melt_f_static)
# Check that it is the one matching Hugonnet, and what the run's melt_f gives
obs = gdir.observations['ref_mb']
y0, y1 = [int(d[:4]) for d in obs['period'].split('_')]
hug_yrs = np.arange(y0, y1)
mb_hug_fixed = mbmod.get_specific_mb(year=hug_yrs).mean()
mb_hug_fixed_dyn = MultipleFlowlineMassBalance(
    gdir, use_inversion_flowlines=True,
    melt_f=melt_f_dyn).get_specific_mb(year=hug_yrs).mean()
print(f'Hugonnet {y0}-{y1}: {obs["value"]:.1f} +- {obs["err"]:.1f}; fixed '
      f'geometry {mb_hug_fixed:.1f} (static melt_f), '
      f'{mb_hug_fixed_dyn:.1f} (with the dyn melt_f)')
assert abs(mb_hug_fixed - obs['value']) < 1
mb['fixed_geom'] = mbmod.get_specific_mb(year=mb.index.values)
# All years, for the plot
yrs = np.arange(1850, 2025)
mb_fixed_all = pd.Series(mbmod.get_specific_mb(year=yrs), index=yrs)

mb = mb.dropna()
mb.index.name = 'year'
mb.to_csv(os.path.join(HERE, f'hef_mb_{EXP}.csv'), float_format='%.1f')

rows = {}
for p0, p1 in [(mb.index[0], mb.index[-1]), (1953, 1979), (1980, 1999),
               (2000, 2020)]:
    sel = mb.loc[p0:p1]
    for k in ['transient', 'fixed_geom']:
        rows[(f'{p0}-{p1}', k)] = stats(sel[k], sel.wgms)
mb_stats = pd.DataFrame(rows).T
mb_stats.index.names = ['period', 'model']
mb_stats.to_csv(os.path.join(HERE, f'hef_mb_stats_{EXP}.csv'),
                float_format='%.3f')
print('\nSpecific mass balance against WGMS (mm w.e. yr-1):')
print(mb_stats.round(2).to_string())
print('\nModel area (km2): ' + ', '.join(
    f'{y}: {run.area.loc[y] * 1e-6:.2f}' for y in [1953, 1980, 2003, 2020]))

f, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
ax1.plot(mb.index, mb.wgms, color='k', lw=2, label='WGMS')
ax1.plot(smb.loc[1850:], color='C0', label='OGGM transient')
ax1.plot(mb_fixed_all, color='C1', label=f'OGGM fixed geometry '
                                         f'({gdir.rgi_date})')
ax1.axhline(0, color='grey', lw=0.5)
ax1.set_ylabel('Specific MB (mm w.e. yr$^{-1}$)')
ax1.set_title(f'{NAME}, {EXP} HIST ens{MEMBER}: annual mass balance')
ax1.legend(loc='lower left')
ax2.plot(mb.index, mb.wgms.cumsum() * 1e-3, color='k', lw=2, label='WGMS')
ax2.plot(mb.index, mb.transient.cumsum() * 1e-3, color='C0',
         label='OGGM transient')
ax2.plot(mb.index, mb.fixed_geom.cumsum() * 1e-3, color='C1',
         label='OGGM fixed geometry')
ax2.set_ylabel(f'Cumulative MB since {mb.index[0] - 1} (m w.e.)')
ax2.set_xlim(1850, 2025)
ax2.legend(loc='lower left')
f.tight_layout()
f.savefig(os.path.join(FIG_DIR, f'hef_mb_{EXP}.png'), dpi=150)
plt.close(f)

# Linear trends over the observation period (OLS, per decade)
from scipy.stats import linregress
rows = {}
for k in ['wgms', 'transient', 'fixed_geom']:
    r = linregress(mb.index, mb[k])
    rows[k] = {'trend_per_decade': r.slope * 10, 'stderr': r.stderr * 10,
               'p': r.pvalue}
r = linregress(mb.index, mb.transient - mb.fixed_geom)
rows['transient - fixed_geom'] = {'trend_per_decade': r.slope * 10,
                                  'stderr': r.stderr * 10, 'p': r.pvalue}
mb_trends = pd.DataFrame(rows).T
mb_trends.to_csv(os.path.join(HERE, f'hef_mb_trends_{EXP}.csv'))
print(f'\nLinear trend {mb.index[0]}-{mb.index[-1]} '
      f'(mm w.e. yr-1 per decade):')
print(mb_trends.to_string(float_format=lambda v: f'{v:.3g}'))

# The observation period only, with the means of the complete decades on top
dec = mb.groupby(mb.index // 10 * 10).agg(['mean', 'count'])
dec = dec.loc[dec[('wgms', 'count')] == 10].xs('mean', axis=1, level=1)
dec.index.name = 'decade'
dec.to_csv(os.path.join(HERE, f'hef_mb_decadal_{EXP}.csv'),
           float_format='%.1f')
print('\nDecadal mean specific mass balance (mm w.e. yr-1):')
print(dec.round(0).to_string())

f, ax = plt.subplots(figsize=(10, 5))
for k, c, lab in [('wgms', 'k', 'WGMS'),
                  ('transient', 'C0', 'OGGM transient'),
                  ('fixed_geom', 'C1',
                   f'OGGM fixed geometry ({gdir.rgi_date})')]:
    ax.plot(mb.index, mb[k], color=c, lw=0.8, alpha=0.6)
    ax.hlines(dec[k], dec.index - 0.5, dec.index + 9.5, color=c, lw=3,
              label=lab)
for d in list(dec.index) + [dec.index[-1] + 10]:
    ax.axvline(d - 0.5, color='grey', lw=0.5, ls=':')
ax.axhline(0, color='grey', lw=0.5)
ax.set_xlim(mb.index[0] - 0.5, mb.index[-1] + 0.5)
ax.set_ylabel('Specific MB (mm w.e. yr$^{-1}$)')
ax.set_title(f'{NAME}, {EXP} HIST ens{MEMBER}: annual mass balance and '
             f'decadal means')
ax.legend(loc='lower left')
f.tight_layout()
f.savefig(os.path.join(FIG_DIR, f'hef_mb_decadal_{EXP}.png'), dpi=150)
plt.close(f)

# --------------------------------- length ----------------------------------

ref_l = gdir.get_ref_length_data()
obs_dl = ref_l.dL.astype(float)
obs_dl = obs_dl - obs_dl.loc[REF_YEAR]
dl = pd.DataFrame({'leclercq': obs_dl})
dl[f'ens{MEMBER}'] = run.length - run.length.loc[REF_YEAR]
others = {}
for m in OTHER_MEMBERS:
    r = read_run(m)
    others[m] = r.length - r.length.loc[REF_YEAR]
    dl[f'ens{m}'] = others[m]
# The centerlines run of run_hef_centerlines.py, if it is there (another
# model setup: default lambda, static calibration)
cl_fp = os.path.join(HERE, 'hef_centerlines_run.csv')
cl = None
if os.path.exists(cl_fp):
    cl = pd.read_csv(cl_fp, index_col=0)
    cl_dl = cl.length - cl.length.loc[REF_YEAR]
    dl['centerlines'] = cl_dl
dl.index.name = 'year'
dl.to_csv(os.path.join(HERE, f'hef_length_{EXP}.csv'), float_format='%.0f')

mod_dl = run.length - run.length.loc[REF_YEAR]
print(f'\nLength change relative to {REF_YEAR} (m), Leclercq / ens{MEMBER} '
      f'[other HIST members] / centerlines:')
for y in ref_l.index:
    if y in (1847, 1855, 1871) or (y % 25 == 0) or y == ref_l.index[-1]:
        oth = ', '.join(f'{others[m].loc[y]:.0f}' for m in OTHER_MEMBERS)
        c = '' if cl is None else f' / {cl_dl.loc[y]:6.0f}'
        print(f'  {y}: {obs_dl.loc[y]:6.0f} / {mod_dl.loc[y]:6.0f}  [{oth}]'
              + c)
# The runs start from the RGI-date geometry in year 1: how far from it are
# they at the RGI date?
print(f'Model length, year 1 (the RGI geometry) -> {REF_YEAR}: '
      f'{run.length.loc[1]:.0f} -> {run.length.loc[REF_YEAR]:.0f} m' +
      ('' if cl is None else f'; centerlines {cl.length.loc[1]:.0f} -> '
                             f'{cl.length.loc[REF_YEAR]:.0f} m'))
print(f'Model length in {REF_YEAR}: {run.length.loc[REF_YEAR]:.0f} m; model '
      f'max since 1700: {mod_dl.loc[1700:].max():.0f} m in '
      f'{mod_dl.loc[1700:].idxmax()}')

f, ax = plt.subplots(figsize=(10, 5))
for i, m in enumerate(OTHER_MEMBERS):
    ax.plot(others[m].loc[1700:], color='C0', lw=0.7, alpha=0.4,
            label='OGGM elev. bands, other HIST members' if i == 0 else None)
ax.plot(mod_dl.loc[1700:], color='C0', lw=2, label=f'OGGM elev. bands ({EXP}), HIST ens{MEMBER}')
if cl is not None:
    ax.plot(cl_dl.loc[1700:], color='C3', lw=2,
            label=f'OGGM centerlines (static calib.), HIST ens{MEMBER}')
ax.plot(obs_dl, 'o-', color='k', ms=3, label='Leclercq et al. (2014)')
ax.axhline(0, color='grey', lw=0.5)
ax.set_ylabel(f'Length change relative to {REF_YEAR} (m)')
ax.set_title(f'{NAME}, {EXP}: length change')
ax.legend()
f.tight_layout()
f.savefig(os.path.join(FIG_DIR, f'hef_length_{EXP}.png'), dpi=150)
plt.close(f)