import glob
import os
import re
import sys

import numpy as np
import pandas as pd
import xarray as xr

DATA_DIR = 'output'
OUT_DIR = 'filtered'

print("Loading glacier statistics from OGGM cluster...")
df = pd.read_csv(
    'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.6/L3-L5_files/2025.6/elev_bands/W5E5/per_glacier_spinup/RGI62/b_160/L3/summary/glacier_statistics_11.csv',
    index_col=0
)
total_vol = df.inv_volume_km3.sum()
total_area = df.rgi_area_km2.sum()

# By default, process every experiment found under output/. Pass one or more
# experiment names as CLI args to (re-)process only those, e.g. to add a new
# experiment without re-running the (possibly expensive, ensemble-heavy) rest.
if len(sys.argv) > 1:
    experiments = sys.argv[1:]
else:
    experiments = sorted(
        os.path.basename(os.path.dirname(p))
        for p in glob.glob(f'{DATA_DIR}/*/lmr_mira_hist_round2_reg11.nc')
    )
print(f"Processing {len(experiments)} experiment(s): {experiments}")

for exp in experiments:
    print(f"\n=== Experiment {exp} ===")
    exp_dir = f'{DATA_DIR}/{exp}'
    out_dir = f'{OUT_DIR}/{exp}'
    os.makedirs(out_dir, exist_ok=True)

    print("Loading hist dataset...")
    with xr.open_dataset(f'{exp_dir}/lmr_mira_hist_round2_reg11.nc') as ds:
        ds_hist = ds[['volume', 'area_min_h']].load()

    # --- Hist: check coverage ---
    valid_ids_hist = ds_hist.rgi_id.data[~ds_hist.isel(time=0).volume.isnull().data]
    vol_cov_hist = df.loc[valid_ids_hist].inv_volume_km3.sum() / total_vol
    area_cov_hist = df.loc[valid_ids_hist].rgi_area_km2.sum() / total_area
    print(f"  hist: {len(valid_ids_hist)} running glaciers, vol cov {vol_cov_hist:.2%}, area cov {area_cov_hist:.2%}")
    assert area_cov_hist > 0.99, f"{exp}: hist area coverage {area_cov_hist:.2%} is not > 99%"

    # --- Nat ensemble (if any): find the common running subset of ids across this
    # experiment's own members only - coverage is never mixed across experiments ---
    ens_files = sorted(glob.glob(f'{exp_dir}/lmr_mira_nat_round2_ens*_reg11.nc'))
    has_ensemble = len(ens_files) > 0

    if has_ensemble:
        print(f"  Found {len(ens_files)} nat ensemble members. Scanning for common running subset...")
        common_nat_ids = None
        member_area_cov = {}
        for f in ens_files:
            name = re.search(r'(ens\d+)', f).group(1)
            with xr.open_dataset(f) as ds:
                valid = ds.rgi_id.data[~ds.isel(time=0).volume.isnull().data]
            member_area_cov[name] = df.loc[valid[np.isin(valid, df.index)]].rgi_area_km2.sum() / total_area
            common_nat_ids = valid if common_nat_ids is None else np.intersect1d(common_nat_ids, valid)

        cov_ser = pd.Series(member_area_cov).sort_values()
        print(f"  per-member area coverage: min {cov_ser.min():.2%} ({cov_ser.index[0]}), "
              f"max {cov_ser.max():.2%}, median {cov_ser.median():.2%}")
        print(f"  members below 90% area coverage: {(cov_ser < 0.90).sum()}/{len(cov_ser)}")

        vol_cov_nat = df.loc[common_nat_ids].inv_volume_km3.sum() / total_vol
        area_cov_nat = df.loc[common_nat_ids].rgi_area_km2.sum() / total_area
        print(f"  nat common subset (intersection of ALL {exp} members): {len(common_nat_ids)} glaciers, "
              f"vol cov {vol_cov_nat:.2%}, area cov {area_cov_nat:.2%}")
        assert area_cov_nat <= area_cov_hist + 1e-9, \
            f"{exp}: nat common coverage should not exceed hist coverage"

        final_ids = np.intersect1d(valid_ids_hist, common_nat_ids)
    else:
        print("  No nat ensemble members for this experiment.")
        final_ids = valid_ids_hist

    print(f"  Final common subset (hist{' ∩ nat' if has_ensemble else ''}): {len(final_ids)} glaciers")
    df.loc[final_ids].to_csv(f'{out_dir}/glacier_statistics_reg11.csv')

    # --- Build hist volume/area series ---
    df_vol_hist = (ds_hist.sel(rgi_id=final_ids).volume.sum('rgi_id') * 1e-9).to_series().to_frame('hist')
    df_vol_hist.index.name = 'year'
    df_area_hist = (ds_hist.sel(rgi_id=final_ids).area_min_h.sum('rgi_id') * 1e-6).to_series().to_frame('hist')
    df_area_hist.index.name = 'year'
    df_vol_hist.to_csv(f'{out_dir}/volume_km3_hist_reg11.csv')
    df_area_hist.to_csv(f'{out_dir}/area_km2_hist_reg11.csv')

    # --- Build nat ensemble volume/area DataFrames ---
    if has_ensemble:
        print(f"  Building nat ensemble DataFrames from {len(ens_files)} members...")
        vol_nat, area_nat = {}, {}
        for i, f in enumerate(ens_files):
            name = re.search(r'(ens\d+)', f).group(1)
            with xr.open_dataset(f) as ds:
                member = ds[['volume', 'area_min_h']].sel(rgi_id=final_ids).load()
            vol_nat[name] = (member.volume.sum('rgi_id') * 1e-9).to_series()
            area_nat[name] = (member.area_min_h.sum('rgi_id') * 1e-6).to_series()
            if (i + 1) % 20 == 0 or (i + 1) == len(ens_files):
                print(f"    {i + 1}/{len(ens_files)}")

        df_vol_nat = pd.DataFrame(vol_nat)
        df_vol_nat.index.name = 'year'
        df_area_nat = pd.DataFrame(area_nat)
        df_area_nat.index.name = 'year'
        df_vol_nat.to_csv(f'{out_dir}/volume_km3_nat_ens_reg11.csv')
        df_area_nat.to_csv(f'{out_dir}/area_km2_nat_ens_reg11.csv')

    print(f"  Saved to {out_dir}/")

print("\nDone.")