"""
Ensemble evaluation for RockGlacierNet — averages N SwinUNet checkpoints.

Each model runs 8-way TTA (4 rotations x 2 flips); the per-pixel sigmoid
probabilities are averaged across ALL models AND all TTA transforms, then
scored exactly like evaluate.py (valid-region only, threshold sweep, per-image).

Usage (set the checkpoints, then run):
    CKPT_PATHS="dir1/best.h5,dir2/best.h5,dir3/best.h5" \
    DATA_DIR=dataset/patches_v2 \
    BAND_LIST="band_25,band_18,..." \
    OUT_DIR=bandselection/ensemble_test \
    python3 evaluate_ensemble.py
"""

import os
import glob
import json
import numpy as np
np.Inf = np.inf
import pandas as pd
import tensorflow as tf
from sklearn.metrics import confusion_matrix

from data_generator import GLOBAL_MEAN, GLOBAL_STD
from models.swin_unet import build_model

# ─────────────────────────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────────────────────────
DATA_DIR    = os.environ.get("DATA_DIR", "dataset/patches_v2")
TEST_FEAT   = os.environ.get("TEST_FEAT", f"{DATA_DIR}/test/features")
TEST_MASK   = os.environ.get("TEST_MASK", f"{DATA_DIR}/test/masks")
OUT_DIR     = os.environ.get("OUT_DIR", "bandselection/ensemble_test")

# Comma-separated list of best_model.weights.h5 paths (one per ensemble member)
CKPT_PATHS = [p.strip() for p in os.environ["CKPT_PATHS"].split(",") if p.strip()]

# Bands MUST match training (all members share the same band set here).
_band_names = json.load(open(f"{DATA_DIR}/metadata/norm_stats.json"))["band_names"]
if os.environ.get("BAND_LIST"):
    _bands = [b.strip() for b in os.environ["BAND_LIST"].split(",") if b.strip()]
    BAND_INDICES = [_band_names.index(b) for b in _bands]
else:
    BAND_INDICES = None

_ph, _pw = np.load(sorted(glob.glob(f"{TEST_FEAT}/*.npy"))[0]).shape[:2]
N_BANDS  = len(BAND_INDICES) if BAND_INDICES else len(_band_names)
INPUT_SHAPE = (_ph, _pw, N_BANDS)
BATCH_SIZE  = 8
THRESHOLDS  = [0.30, 0.35, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70]

os.makedirs(OUT_DIR, exist_ok=True)
print(f"[ensemble] {len(CKPT_PATHS)} models | bands={N_BANDS} | patch={_ph}x{_pw}")
for p in CKPT_PATHS:
    print(f"           {p}")


# ─────────────────────────────────────────────────────────────────────────────
# Normalization — use the first checkpoint's saved stats (all identical here)
# ─────────────────────────────────────────────────────────────────────────────
def _load_norm(ckpt_path):
    cand = os.path.join(os.path.dirname(ckpt_path), "norm_stats.json")
    if os.path.exists(cand):
        st = json.load(open(cand))
        m = np.array(st["mean"], np.float32).reshape(1, 1, -1)
        s = np.array(st["std"], np.float32).reshape(1, 1, -1)
        print(f"[ensemble] norm stats: {cand}")
        return m, s
    print("[ensemble] norm stats: built-in fallback")
    return GLOBAL_MEAN, GLOBAL_STD


NORM_MEAN, NORM_STD = _load_norm(CKPT_PATHS[0])


def metrics_from_cm(tp, fp, fn, tn):
    prec = tp / (tp + fp + 1e-9)
    rec  = tp / (tp + fn + 1e-9)
    f1   = 2 * prec * rec / (prec + rec + 1e-9)
    iou  = tp / (tp + fp + fn + 1e-9)
    dice = 2 * tp / (2 * tp + fp + fn + 1e-9)
    acc  = (tp + tn) / (tp + tn + fp + fn + 1e-9)
    spec = tn / (tn + fp + 1e-9)
    return prec, rec, f1, iou, dice, acc, spec


def predict_tta(model, X, batch_size):
    """Average sigmoid probs over the 8 dihedral (D4) transforms."""
    acc = np.zeros((X.shape[0], X.shape[1], X.shape[2], 1), np.float32)
    n = 0
    for flip in (False, True):
        for k in (0, 1, 2, 3):
            Xt = X[:, :, ::-1, :] if flip else X
            Xt = np.ascontiguousarray(np.rot90(Xt, k, axes=(1, 2)))
            p = model.predict(Xt, batch_size=batch_size, verbose=0)
            p = np.rot90(p, -k, axes=(1, 2))
            if flip:
                p = p[:, :, ::-1, :]
            acc += p
            n += 1
    return acc / n


# ─────────────────────────────────────────────────────────────────────────────
# Load + preprocess every test patch (.npy)
# ─────────────────────────────────────────────────────────────────────────────
feat_files = sorted(glob.glob(os.path.join(TEST_FEAT, '*.npy')))
mask_files = [os.path.join(TEST_MASK, os.path.basename(f)) for f in feat_files]
for mf in mask_files:
    assert os.path.exists(mf), f"Missing mask: {mf}"
print(f"\nLoading {len(feat_files)} test patches…")

X, Y, ids = [], [], []
for fp, mp in zip(feat_files, mask_files):
    feat = np.load(fp).astype(np.float32)
    mask = np.load(mp).astype(np.float32)
    feat = np.where((feat == -9999) | ~np.isfinite(feat), 0, feat)
    if BAND_INDICES is not None:
        feat = feat[..., BAND_INDICES]
    feat = (feat - NORM_MEAN) / (NORM_STD + 1e-7)
    X.append(feat); Y.append(mask)
    ids.append(os.path.basename(fp).replace('.npy', ''))

X = np.array(X, dtype=np.float32)
Y = np.array(Y, dtype=np.float32)

# ─────────────────────────────────────────────────────────────────────────────
# Predict with each model, average
# ─────────────────────────────────────────────────────────────────────────────
model = build_model(input_shape=INPUT_SHAPE, num_classes=1, num_heads=8)

Y_prob_sum = np.zeros((X.shape[0], X.shape[1], X.shape[2], 1), np.float32)
for mi, ckpt in enumerate(CKPT_PATHS, 1):
    if not os.path.exists(ckpt):
        raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
    model.load_weights(ckpt)
    print(f"[{mi}/{len(CKPT_PATHS)}] TTA predicting with {ckpt} …")
    Y_prob_sum += predict_tta(model, X, BATCH_SIZE)

Y_prob = Y_prob_sum / len(CKPT_PATHS)     # ensemble mean probability

# ─────────────────────────────────────────────────────────────────────────────
# Per-image metrics @ t=0.5
# ─────────────────────────────────────────────────────────────────────────────
all_true, all_prob, per_image = [], [], []
for i, iid in enumerate(ids):
    yt = Y[i, :, :, 0].flatten().astype(np.int32)
    yprob = Y_prob[i, :, :, 0].flatten()
    all_true.append(yt); all_prob.append(yprob)

    yp = (yprob >= 0.5).astype(np.int32)
    tn, fp, fn, tp = confusion_matrix(yt, yp, labels=[0, 1]).ravel()
    prec, rec, f1, iou, dice, acc, spec = metrics_from_cm(tp, fp, fn, tn)
    if (tp + fp + fn) == 0:
        prec = rec = f1 = iou = dice = 1.0
    per_image.append({
        'image_id': iid,
        'glacier_frac_gt': round(float(yt.mean()), 4),
        'TP': int(tp), 'FP': int(fp), 'FN': int(fn), 'TN': int(tn),
        'precision': round(prec, 4), 'recall': round(rec, 4),
        'f1_score': round(f1, 4), 'iou': round(iou, 4), 'dice': round(dice, 4),
    })

per_image_df = pd.DataFrame(per_image).sort_values('iou').reset_index(drop=True)
per_image_df.to_csv(os.path.join(OUT_DIR, 'per_image_metrics.csv'), index=False)

# ─────────────────────────────────────────────────────────────────────────────
# Global threshold sweep
# ─────────────────────────────────────────────────────────────────────────────
y_true_flat = np.concatenate(all_true)
y_prob_flat = np.concatenate(all_prob)

rows = []
for t in THRESHOLDS:
    yp = (y_prob_flat >= t).astype(np.int32)
    tn, fp, fn, tp = confusion_matrix(y_true_flat, yp, labels=[0, 1]).ravel()
    prec, rec, f1, iou, dice, acc, spec = metrics_from_cm(tp, fp, fn, tn)
    rows.append({'threshold': t, 'precision': round(prec, 6), 'recall': round(rec, 6),
                 'f1_score': round(f1, 6), 'iou': round(iou, 6), 'dice': round(dice, 6)})
thresh_df = pd.DataFrame(rows)
thresh_df.to_csv(os.path.join(OUT_DIR, 'evaluation_metrics.csv'), index=False)

best_row = thresh_df.loc[thresh_df['iou'].idxmax()]
best_t = float(best_row['threshold'])

# Glacier-only mean (honest per-image number)
gl = per_image_df[per_image_df['glacier_frac_gt'] > 0.01]

with open(os.path.join(OUT_DIR, 'metrics.json'), 'w') as f:
    json.dump({
        'n_models':          len(CKPT_PATHS),
        'best_threshold':    best_t,
        'test_precision':    round(float(best_row['precision']), 4),
        'test_recall':       round(float(best_row['recall']), 4),
        'test_f1':           round(float(best_row['f1_score']), 4),
        'test_iou':          round(float(best_row['iou']), 4),
        'test_dice':         round(float(best_row['dice']), 4),
        'glacier_only_dice': round(float(gl['dice'].mean()), 4),
        'glacier_only_iou':  round(float(gl['iou'].mean()), 4),
    }, f, indent=2)

# ─────────────────────────────────────────────────────────────────────────────
# Summary
# ─────────────────────────────────────────────────────────────────────────────
lines = [
    "=" * 60,
    f"  ENSEMBLE EVALUATION — {len(CKPT_PATHS)} SwinUNet models",
    "=" * 60,
    "",
    f"Test patches          : {len(ids)}",
    f"Best threshold by IoU : {best_t}",
    "",
    "GLOBAL (pixel-pooled) @ best threshold:",
    f"  precision : {best_row['precision']:.4f}",
    f"  recall    : {best_row['recall']:.4f}",
    f"  f1        : {best_row['f1_score']:.4f}",
    f"  iou       : {best_row['iou']:.4f}",
    f"  dice      : {best_row['dice']:.4f}",
    "",
    f"GLACIER-ONLY per-image mean ({len(gl)} patches):",
    f"  dice : {gl['dice'].mean():.4f}",
    f"  iou  : {gl['iou'].mean():.4f}",
    "",
    "Threshold sweep:",
    "  {:<8} {:<10} {:<10} {:<10} {:<10}".format("thresh", "precision", "recall", "iou", "dice"),
    "  " + "-" * 52,
]
for _, r in thresh_df.iterrows():
    tag = " <- best" if r['threshold'] == best_t else ""
    lines.append("  {:<8} {:<10.4f} {:<10.4f} {:<10.4f} {:<10.4f}{}".format(
        r['threshold'], r['precision'], r['recall'], r['iou'], r['dice'], tag))
lines += ["", f"Output dir: {OUT_DIR}"]

summary = "\n".join(lines)
print("\n" + summary)
with open(os.path.join(OUT_DIR, 'evaluation_summary.txt'), 'w') as f:
    f.write(summary)
print("\nEnsemble evaluation complete.")