#!/usr/bin/env python3
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# ==========================
# USER SETTINGS
# ==========================
SUMMARY_CSV = "./et_compare_with_openet_polymean/et_metrics_vs_conus_summary.csv"
OUTDIR = "./et_compare_with_openet_polymean/summary_figures_vs_conus"
os.makedirs(OUTDIR, exist_ok=True)

# Optional: only keep sites with at least N samples for BOTH comparisons
MIN_SAMPLES_BOTH = 5

# ==========================
# Read + clean
# ==========================
df = pd.read_csv(SUMMARY_CSV).sort_values("site_id")

# Keep only sites where both comparisons have enough samples (optional but usually helpful)
df = df[
    (df["n_global_vs_conus"] >= MIN_SAMPLES_BOTH) &
    (df["n_openet_vs_conus"] >= MIN_SAMPLES_BOTH)
].copy()

if df.empty:
    raise SystemExit("No sites left after filtering by MIN_SAMPLES_BOTH. Lower the threshold or check input CSV.")

sites = df["site_id"].astype(str).tolist()
x = np.arange(len(sites))
width = 0.42  # two bars per group

# ==========================
# Helper: grouped bar chart
# ==========================
def grouped_bar(metric_col_global, metric_col_openet, ylabel, title, fname, zero_line=False):
    y1 = df[metric_col_global].to_numpy()
    y2 = df[metric_col_openet].to_numpy()

    plt.figure(figsize=(max(10, len(sites) * 0.35), 4.8))
    plt.bar(x - width/2, y1, width, label="Global vs CONUS")
    plt.bar(x + width/2, y2, width, label="OpenET vs CONUS")

    if zero_line:
        plt.axhline(0, linewidth=0.8)

    plt.xticks(x, sites, rotation=60, ha="right")
    plt.ylabel(ylabel)
    plt.title(title)
    plt.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(OUTDIR, fname), dpi=300)
    plt.close()

# ==========================
# Helper: paired boxplot
# ==========================
def paired_boxplot(metric_col_global, metric_col_openet, ylabel, title, fname, zero_line=False):
    data = [
        df[metric_col_global].dropna().to_numpy(),
        df[metric_col_openet].dropna().to_numpy()
    ]

    plt.figure(figsize=(5.2, 5.2))
    plt.boxplot(data, labels=["Global vs CONUS", "OpenET vs CONUS"], widths=0.6)

    if zero_line:
        plt.axhline(0, linewidth=0.8)

    plt.ylabel(ylabel)
    plt.title(title)
    plt.tight_layout()
    plt.savefig(os.path.join(OUTDIR, fname), dpi=300)
    plt.close()

# ==========================
# 1) Grouped bar charts by site
# ==========================
grouped_bar(
    "bias_global_vs_conus",
    "bias_openet_vs_conus",
    "Bias (mm day$^{-1}$)  [Model − CONUS]",
    "Bias by Site (Global & OpenET vs CONUS)",
    "bar_bias_by_site_global_openet_vs_conus.png",
    zero_line=True
)

grouped_bar(
    "rmse_global_vs_conus",
    "rmse_openet_vs_conus",
    "RMSE (mm day$^{-1}$)",
    "RMSE by Site (Global & OpenET vs CONUS)",
    "bar_rmse_by_site_global_openet_vs_conus.png"
)

grouped_bar(
    "corr_global_vs_conus",
    "corr_openet_vs_conus",
    "Correlation (r)",
    "Correlation by Site (Global & OpenET vs CONUS)",
    "bar_corr_by_site_global_openet_vs_conus.png"
)

# ==========================
# 2) Paired boxplots across sites
# ==========================
paired_boxplot(
    "bias_global_vs_conus",
    "bias_openet_vs_conus",
    "Bias (mm day$^{-1}$)  [Model − CONUS]",
    "Bias Distribution Across Sites",
    "box_bias_global_openet_vs_conus.png",
    zero_line=True
)

paired_boxplot(
    "rmse_global_vs_conus",
    "rmse_openet_vs_conus",
    "RMSE (mm day$^{-1}$)",
    "RMSE Distribution Across Sites",
    "box_rmse_global_openet_vs_conus.png"
)

paired_boxplot(
    "corr_global_vs_conus",
    "corr_openet_vs_conus",
    "Correlation (r)",
    "Correlation Distribution Across Sites",
    "box_corr_global_openet_vs_conus.png"
)

print(f"✅ Figures saved to: {OUTDIR}")
print(f"✅ Sites plotted: {len(df)}")