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

print("Loading main datasets...")
with xr.open_dataset('output/lmr_mira_hist_reg11.nc') as ds:
    ds_mira_hist = ds.load()
with xr.open_dataset('output/lmr_mira_nat_reg11.nc') as ds:
    ds_mira_nat = ds.sel(time=slice(1850, None)).load()
with xr.open_dataset('output/lmr_mira_hist_cl_reg11.nc') as ds:
    ds_mira_hist_cl = ds.load()
with xr.open_dataset('output/lmr_mira_nat_cl_reg11.nc') as ds:
    ds_mira_nat_cl = ds.sel(time=slice(1850, None)).load()

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()

# --- Hist: independent coverage correction for each simulation ---
print("Computing hist coverage corrections...")

valid_ids_hist_eb = ds_mira_hist.rgi_id.data[~ds_mira_hist.isel(time=0).volume.isnull().data]
vol_cov_hist_eb = df.loc[valid_ids_hist_eb].inv_volume_km3.sum() / total_vol
area_cov_hist_eb = df.loc[valid_ids_hist_eb].rgi_area_km2.sum() / total_area
print(f"  hist_eb: {len(valid_ids_hist_eb)} glaciers, vol {vol_cov_hist_eb:.1%}, area {area_cov_hist_eb:.1%}"
      f" -> corrections x{1/vol_cov_hist_eb:.4f} / x{1/area_cov_hist_eb:.4f}")

valid_ids_hist_cl = ds_mira_hist_cl.rgi_id.data[~ds_mira_hist_cl.isel(time=0).volume.isnull().data]
vol_cov_hist_cl = df.loc[valid_ids_hist_cl].inv_volume_km3.sum() / total_vol
area_cov_hist_cl = df.loc[valid_ids_hist_cl].rgi_area_km2.sum() / total_area
print(f"  hist_cl: {len(valid_ids_hist_cl)} glaciers, vol {vol_cov_hist_cl:.1%}, area {area_cov_hist_cl:.1%}"
      f" -> corrections x{1/vol_cov_hist_cl:.4f} / x{1/area_cov_hist_cl:.4f}")

ser_vol_hist_eb = (ds_mira_hist.sel(rgi_id=valid_ids_hist_eb).volume.sum('rgi_id') * 1e-9 / vol_cov_hist_eb).to_series().rename('hist_eb')
ser_area_hist_eb = (ds_mira_hist.sel(rgi_id=valid_ids_hist_eb).area_min_h.sum('rgi_id') * 1e-6 / area_cov_hist_eb).to_series().rename('hist_eb')
ser_vol_hist_cl = (ds_mira_hist_cl.sel(rgi_id=valid_ids_hist_cl).volume.sum('rgi_id') * 1e-9 / vol_cov_hist_cl).to_series().rename('hist_cl')
ser_area_hist_cl = (ds_mira_hist_cl.sel(rgi_id=valid_ids_hist_cl).area_min_h.sum('rgi_id') * 1e-6 / area_cov_hist_cl).to_series().rename('hist_cl')

# Reference hist values at 1850 used to normalize nat and ensemble
ref_vol_eb = float(ser_vol_hist_eb.loc[1850])
ref_area_eb = float(ser_area_hist_eb.loc[1850])
ref_vol_cl = float(ser_vol_hist_cl.loc[1850])
ref_area_cl = float(ser_area_hist_cl.loc[1850])
print(f"  hist_eb at 1850: vol={ref_vol_eb:.2f} km3, area={ref_area_eb:.1f} km2")
print(f"  hist_cl at 1850: vol={ref_vol_cl:.2f} km3, area={ref_area_cl:.1f} km2")

# Save glacier statistics for reference (glaciers valid in both hist simulations)
valid_ids_hist_both = np.intersect1d(valid_ids_hist_eb, valid_ids_hist_cl)
df.loc[valid_ids_hist_both].to_csv('filtered_histnat/glacier_statistics_reg11.csv')
print(f"  Saved glacier statistics for {len(valid_ids_hist_both)} glaciers (hist_eb ∩ hist_cl)")

# --- Nat: coverage check + normalize to hist at 1850 ---
print("Processing nat simulations...")

def _nat_series(ds, name, ref_vol, ref_area):
    ids = ds.rgi_id.data[~ds.isel(time=0).volume.isnull().data]
    vol_cov = df.loc[ids[np.isin(ids, df.index)]].inv_volume_km3.sum() / total_vol
    if vol_cov < 0.85:
        raise ValueError(f"{name}: volume coverage {vol_cov:.1%} < 85%")
    raw_vol = (ds.sel(rgi_id=ids).volume.sum('rgi_id') * 1e-9).to_series()
    raw_area = (ds.sel(rgi_id=ids).area_min_h.sum('rgi_id') * 1e-6).to_series()
    vol_f = ref_vol / float(raw_vol.loc[1850])
    area_f = ref_area / float(raw_area.loc[1850])
    print(f"  {name}: {len(ids)} glaciers, vol {vol_cov:.1%}, norm vol x{vol_f:.4f}, area x{area_f:.4f}")
    return (raw_vol * vol_f).rename(name), (raw_area * area_f).rename(name)

ser_vol_nat_eb, ser_area_nat_eb = _nat_series(ds_mira_nat, 'nat_eb', ref_vol_eb, ref_area_eb)
ser_vol_nat_cl, ser_area_nat_cl = _nat_series(ds_mira_nat_cl, 'nat_cl', ref_vol_cl, ref_area_cl)

# --- Ensemble: coverage check + normalize to hist at 1850 ---
ens_files_cl = sorted(glob.glob('output/ensemble/cl/lmr_mira_nat_cl_ens*_reg11.nc'))
print(f"Loading {len(ens_files_cl)} CL ensemble members...")
ens_vol_cl, ens_area_cl = {}, {}
discarded_cl = 0
for i, f in enumerate(ens_files_cl):
    name = re.search(r'(ens\d+)', f).group(1)
    with xr.open_dataset(f) as ds:
        member = ds.sel(time=slice(1850, None)).load()
    member_ids = member.rgi_id.data[~member.isel(time=0).volume.isnull().data]
    vol_cov = df.loc[member_ids[np.isin(member_ids, df.index)]].inv_volume_km3.sum() / total_vol
    if vol_cov < 0.85:
        print(f"  Discarding {name}: volume coverage {vol_cov:.1%} < 85%")
        discarded_cl += 1
        continue
    raw_vol = (member.sel(rgi_id=member_ids).volume.sum('rgi_id') * 1e-9).to_series()
    raw_area = (member.sel(rgi_id=member_ids).area_min_h.sum('rgi_id') * 1e-6).to_series()
    ens_vol_cl[name] = raw_vol * (ref_vol_cl / float(raw_vol.loc[1850]))
    ens_area_cl[name] = raw_area * (ref_area_cl / float(raw_area.loc[1850]))
    if (i + 1) % 20 == 0 or (i + 1) == len(ens_files_cl):
        print(f"  {i + 1}/{len(ens_files_cl)} (kept {len(ens_vol_cl)}, discarded {discarded_cl})")

ens_files_eb = sorted(glob.glob('output/ensemble/eb/lmr_mira_nat_ens*_reg11.nc'))
print(f"Loading {len(ens_files_eb)} EB ensemble members...")
ens_vol_eb, ens_area_eb = {}, {}
discarded_eb = 0
for i, f in enumerate(ens_files_eb):
    name = re.search(r'(ens\d+)', f).group(1)
    with xr.open_dataset(f) as ds:
        member = ds.sel(time=slice(1850, None)).load()
    member_ids = member.rgi_id.data[~member.isel(time=0).volume.isnull().data]
    vol_cov = df.loc[member_ids[np.isin(member_ids, df.index)]].inv_volume_km3.sum() / total_vol
    if vol_cov < 0.85:
        print(f"  Discarding {name}: volume coverage {vol_cov:.1%} < 85%")
        discarded_eb += 1
        continue
    raw_vol = (member.sel(rgi_id=member_ids).volume.sum('rgi_id') * 1e-9).to_series()
    raw_area = (member.sel(rgi_id=member_ids).area_min_h.sum('rgi_id') * 1e-6).to_series()
    ens_vol_eb[name] = raw_vol * (ref_vol_eb / float(raw_vol.loc[1850]))
    ens_area_eb[name] = raw_area * (ref_area_eb / float(raw_area.loc[1850]))
    if (i + 1) % 20 == 0 or (i + 1) == len(ens_files_eb):
        print(f"  {i + 1}/{len(ens_files_eb)} (kept {len(ens_vol_eb)}, discarded {discarded_eb})")

# --- Build and save DataFrames ---
print("Building DataFrames and saving to filtered_histnat/...")

df_vol = pd.concat([ser_vol_hist_eb, ser_vol_nat_eb, ser_vol_hist_cl, ser_vol_nat_cl], axis=1)
df_vol.index.name = 'year'
df_area = pd.concat([ser_area_hist_eb, ser_area_nat_eb, ser_area_hist_cl, ser_area_nat_cl], axis=1)
df_area.index.name = 'year'

df_vol_ens_cl = pd.DataFrame(ens_vol_cl)
df_vol_ens_cl.index.name = 'year'
df_area_ens_cl = pd.DataFrame(ens_area_cl)
df_area_ens_cl.index.name = 'year'

df_vol_ens_eb = pd.DataFrame(ens_vol_eb)
df_vol_ens_eb.index.name = 'year'
df_area_ens_eb = pd.DataFrame(ens_area_eb)
df_area_ens_eb.index.name = 'year'

df_vol.to_csv('filtered_histnat/volume_km3_reg11.csv')
df_area.to_csv('filtered_histnat/area_km2_reg11.csv')
df_vol_ens_cl.to_csv('filtered_histnat/volume_km3_ens_cl_reg11.csv')
df_area_ens_cl.to_csv('filtered_histnat/area_km2_ens_cl_reg11.csv')
df_vol_ens_eb.to_csv('filtered_histnat/volume_km3_ens_eb_reg11.csv')
df_area_ens_eb.to_csv('filtered_histnat/area_km2_ens_eb_reg11.csv')

print("Done. Files saved to filtered_histnat/:")
for fname in [
    'glacier_statistics_reg11.csv',
    'volume_km3_reg11.csv', 'area_km2_reg11.csv',
    'volume_km3_ens_cl_reg11.csv', 'area_km2_ens_cl_reg11.csv',
    'volume_km3_ens_eb_reg11.csv', 'area_km2_ens_eb_reg11.csv',
]:
    print(f"  filtered_histnat/{fname}")