stage1_final/ensemble.py

"""Stage 1 v1 final: 4-model ensemble stacking (convnext, swin, mamba, resnet).

Both cRT and plain ensembles. 3 combiner methods:
  A. logreg_logits_balanced   — LogReg on 12-dim raw logits (4 backbones × 3 classes)
  B. logreg_softmax_balanced  — LogReg on 12-dim softmax probabilities
  C. avg_softmax              — simple softmax-probability averaging

Best method selected by pooled OOF macro-F1, then applied to test.

Usage:
  python ensemble.py --resize letterbox
  python ensemble.py --resize squeeze
"""
import argparse
import json
import pathlib
import sys
import numpy as np
from scipy.special import softmax
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import (f1_score, accuracy_score, balanced_accuracy_score,
                             confusion_matrix, classification_report)

ROOT = pathlib.Path(__file__).parent
sys.path.insert(0, str(ROOT))

BACKBONES = ["convnext", "swin", "mamba", "resnet"]
NUM_CLASSES = 3
SEED = 17
METHODS = ["logreg_logits_balanced", "logreg_softmax_balanced", "avg_softmax"]


def get_paths(resize):
    return ROOT / "logits" / resize, ROOT / "results" / resize


def load_split(logits_dir, kind, prefix=""):
    """kind = 'oof' or 'test'. prefix = '' for cRT, 'plain_' for plain.
    Logits are stored in per-backbone subdirectories: logits_dir/{bb}/{prefix}{bb}_{kind}.npz
    """
    out = {}
    for bb in BACKBONES:
        f = logits_dir / bb / f"{prefix}{bb}_{kind}.npz"
        z = np.load(f, allow_pickle=True)
        out[bb] = {k: z[k] for k in z.files}
    return out


def stack_features(d, backbones=BACKBONES, use_softmax=False):
    arrs = []
    for bb in backbones:
        lg = d[bb]["logits"].astype(np.float32)
        arrs.append(softmax(lg, axis=1) if use_softmax else lg)
    return np.concatenate(arrs, axis=1).astype(np.float32)


def compute_per_class_precision(true, pred, names):
    true = np.asarray(true); pred = np.asarray(pred)
    result = {}
    for i, name in enumerate(names):
        tp = ((pred == i) & (true == i)).sum()
        fp = ((pred == i) & (true != i)).sum()
        result[name] = float(tp / max(tp + fp, 1))
    return result


def metrics_from_pred(pred, true, src=None, names=None):
    if names is None:
        names = ["malignant", "non_malignant", "non_roi"]
    pred = np.asarray(pred); true = np.asarray(true)
    m = {
        "macro_f1": float(f1_score(true, pred, average="macro", labels=list(range(NUM_CLASSES)), zero_division=0)),
        "accuracy": float(accuracy_score(true, pred)),
        "balanced_acc": float(balanced_accuracy_score(true, pred)),
        "per_class_recall": {names[i]: float(((true == i) & (pred == i)).sum() / max((true == i).sum(), 1)) for i in range(NUM_CLASSES)},
        "per_class_precision": compute_per_class_precision(true, pred, names),
        "confusion_matrix": confusion_matrix(true, pred, labels=list(range(NUM_CLASSES))).tolist(),
        "classification_report": classification_report(true, pred, labels=list(range(NUM_CLASSES)), target_names=names, output_dict=True, zero_division=0),
        "n": len(pred),
    }
    if src is not None:
        src = np.asarray(src)
        for s in ["B", "C", "D"]:
            mk = src == s
            if mk.sum() > 0:
                m[f"src_{s}_acc"] = float(accuracy_score(true[mk], pred[mk]))
                m[f"src_{s}_n"] = int(mk.sum())
    return m


def make_logreg():
    return LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs",
                              class_weight="balanced", random_state=SEED)


def predict_method(method, X_tr, y_tr, X_va, n_backbones=len(BACKBONES)):
    if method == "avg_softmax":
        blocks = [X_va[:, i*NUM_CLASSES:(i+1)*NUM_CLASSES] for i in range(n_backbones)]
        avg = np.mean(blocks, axis=0)
        return avg.argmax(1)
    scaler = StandardScaler()
    Xtr = scaler.fit_transform(X_tr)
    Xva = scaler.transform(X_va)
    clf = make_logreg()
    clf.fit(Xtr, y_tr)
    return clf.predict(Xva)


def run_cv_method(method, X_oof, y_oof, cv_folds, src_oof):
    fold_metrics = []
    for k in range(5):
        tr_mask = cv_folds != k
        va_mask = cv_folds == k
        pred_va = predict_method(method, X_oof[tr_mask], y_oof[tr_mask], X_oof[va_mask])
        m = metrics_from_pred(pred_va, y_oof[va_mask], src_oof[va_mask])
        m["fold"] = int(k)
        fold_metrics.append(m)
    f1s = [fm["macro_f1"] for fm in fold_metrics]
    nms = [fm["per_class_recall"]["non_malignant"] for fm in fold_metrics]
    baccs = [fm["balanced_acc"] for fm in fold_metrics]
    prec_mal = [fm["per_class_precision"]["malignant"] for fm in fold_metrics]
    prec_nm = [fm["per_class_precision"]["non_malignant"] for fm in fold_metrics]
    prec_nroi = [fm["per_class_precision"]["non_roi"] for fm in fold_metrics]
    rec_mal = [fm["per_class_recall"]["malignant"] for fm in fold_metrics]
    rec_nm = [fm["per_class_recall"]["non_malignant"] for fm in fold_metrics]
    rec_nroi = [fm["per_class_recall"]["non_roi"] for fm in fold_metrics]
    return {
        "folds": fold_metrics,
        "cv_macro_f1_mean": float(np.mean(f1s)), "cv_macro_f1_std": float(np.std(f1s)),
        "cv_non_malignant_recall_mean": float(np.mean(nms)), "cv_non_malignant_recall_std": float(np.std(nms)),
        "cv_balanced_acc_mean": float(np.mean(baccs)), "cv_balanced_acc_std": float(np.std(baccs)),
        "cv_malignant_precision_mean": float(np.mean(prec_mal)), "cv_malignant_precision_std": float(np.std(prec_mal)),
        "cv_non_malignant_precision_mean": float(np.mean(prec_nm)), "cv_non_malignant_precision_std": float(np.std(prec_nm)),
        "cv_non_roi_precision_mean": float(np.mean(prec_nroi)), "cv_non_roi_precision_std": float(np.std(prec_nroi)),
        "cv_malignant_recall_mean": float(np.mean(rec_mal)), "cv_malignant_recall_std": float(np.std(rec_mal)),
        "cv_non_malignant_recall_mean": float(np.mean(nms)), "cv_non_malignant_recall_std": float(np.std(nms)),
        "cv_non_roi_recall_mean": float(np.mean(rec_nroi)), "cv_non_roi_recall_std": float(np.std(rec_nroi)),
    }


def run_pooled_method(method, X_oof, y_oof, cv_folds, src_oof):
    oof_pred = np.zeros(len(y_oof), dtype=np.int64)
    for k in range(5):
        tr_mask = cv_folds != k
        va_mask = cv_folds == k
        oof_pred[va_mask] = predict_method(method, X_oof[tr_mask], y_oof[tr_mask], X_oof[va_mask])
    return metrics_from_pred(oof_pred, y_oof, src_oof)


def fit_predict_test(method, X_oof, y_oof, X_test, n_backbones=len(BACKBONES)):
    if method == "avg_softmax":
        blocks = [X_test[:, i*NUM_CLASSES:(i+1)*NUM_CLASSES] for i in range(n_backbones)]
        avg = np.mean(blocks, axis=0)
        return avg.argmax(1), avg
    scaler = StandardScaler()
    Xtr = scaler.fit_transform(X_oof)
    Xte = scaler.transform(X_test)
    clf = make_logreg()
    clf.fit(Xtr, y_oof)
    pred = clf.predict(Xte)
    proba = clf.predict_proba(Xte)
    return pred, proba


def single_metrics(logits_dir, kind, prefix=""):
    d = load_split(logits_dir, kind, prefix=prefix)
    out = {}
    for bb in BACKBONES:
        pred = d[bb]["logits"].argmax(1)
        y = d[bb]["labels"]; src = d[bb].get("source", None)
        out[bb] = metrics_from_pred(pred, y, src)
    return out


def run_ensemble(label, prefix, logits_dir, results_dir):
    """label='cRT' or 'plain'. prefix='' for cRT, 'plain_' for plain."""
    print(f"\n=== Ensemble stacking ({label}, 4 backbones: {BACKBONES}) ===")
    oof = load_split(logits_dir, "oof", prefix=prefix)
    test = load_split(logits_dir, "test", prefix=prefix)

    for bb in BACKBONES:
        assert np.array_equal(oof[bb]["labels"], oof["convnext"]["labels"]), f"oof labels mismatch {bb}"
        assert np.array_equal(test[bb]["labels"], test["convnext"]["labels"]), f"test labels mismatch {bb}"
    print("alignment OK")

    y_oof = oof["convnext"]["labels"].astype(np.int64)
    y_test = test["convnext"]["labels"].astype(np.int64)
    cv_folds = oof["convnext"]["cv_fold"].astype(np.int64)
    src_oof = oof["convnext"]["source"]
    src_test = test["convnext"]["source"]

    X_oof_raw = stack_features(oof, use_softmax=False)
    X_oof_sm = stack_features(oof, use_softmax=True)
    X_test_raw = stack_features(test, use_softmax=False)
    X_test_sm = stack_features(test, use_softmax=True)
    print(f"X_oof_raw {X_oof_raw.shape}  X_oof_sm {X_oof_sm.shape}")

    method_features = {
        "logreg_logits_balanced": (X_oof_raw, X_test_raw),
        "logreg_softmax_balanced": (X_oof_sm, X_test_sm),
        "avg_softmax": (X_oof_sm, X_test_sm),
    }

    all_cv = {}
    all_pooled = {}
    all_test = {}
    for method in METHODS:
        Xo, Xt = method_features[method]
        print(f"\n--- {method} CV per-fold ---")
        cv_res = run_cv_method(method, Xo, y_oof, cv_folds, src_oof)
        for fm in cv_res["folds"]:
            print(f"  fold{fm['fold']}: f1={fm['macro_f1']:.4f} nm_rec={fm['per_class_recall']['non_malignant']:.4f}")
        print(f"  CV: f1={cv_res['cv_macro_f1_mean']:.4f}+-{cv_res['cv_macro_f1_std']:.4f}")

        print(f"--- {method} CV pooled ---")
        pooled = run_pooled_method(method, Xo, y_oof, cv_folds, src_oof)
        print(f"  pooled: f1={pooled['macro_f1']:.4f} nm_rec={pooled['per_class_recall']['non_malignant']:.4f}")

        print(f"--- {method} Test ---")
        pred_te, proba_te = fit_predict_test(method, Xo, y_oof, Xt)
        tm = metrics_from_pred(pred_te, y_test, src_test)
        print(f"  test: f1={tm['macro_f1']:.4f} nm_rec={tm['per_class_recall']['non_malignant']:.4f} bacc={tm['balanced_acc']:.4f}")

        all_cv[method] = cv_res
        all_pooled[method] = pooled
        all_test[method] = {"test_metrics": tm, "pred": pred_te.tolist(), "proba": proba_te.tolist()}

    best_method = max(METHODS, key=lambda m: all_pooled[m]["macro_f1"])
    print(f"\n=== Best method: {best_method} (pooled OOF macro-F1={all_pooled[best_method]['macro_f1']:.4f}) ===")

    tag = label.lower()  # 'crt' or 'plain'
    cv_summary = {
        "label": label, "methods": METHODS, "best_method": best_method,
        "backbones": BACKBONES,
        "per_method_cv": {m: all_cv[m] for m in METHODS},
        "per_method_pooled": {m: all_pooled[m] for m in METHODS},
        "best_cv": all_cv[best_method], "best_pooled": all_pooled[best_method],
    }
    (results_dir / f"ensemble_{tag}_cv_summary.json").write_text(json.dumps(cv_summary, indent=2))

    best_test = all_test[best_method]
    test_summary = {
        "label": label, "best_method": best_method,
        "test_metrics": best_test["test_metrics"],
        "all_methods_test": {m: all_test[m]["test_metrics"] for m in METHODS},
    }
    (results_dir / f"ensemble_{tag}_test.json").write_text(json.dumps(test_summary, indent=2))

    # Single model metrics
    single_cv = single_metrics(logits_dir, "oof", prefix=prefix)
    single_te = single_metrics(logits_dir, "test", prefix=prefix)
    comp = {
        "cv_single": {bb: single_cv[bb] for bb in BACKBONES},
        "cv_ensemble_best_per_fold": all_cv[best_method],
        "cv_ensemble_best_pooled": all_pooled[best_method],
        "test_single": {bb: single_te[bb] for bb in BACKBONES},
        "test_ensemble_best": best_test["test_metrics"],
        "best_method": best_method,
        "all_methods_test": {m: all_test[m]["test_metrics"] for m in METHODS},
    }
    (results_dir / f"ensemble_{tag}_single_vs_ensemble.json").write_text(json.dumps(comp, indent=2))

    print(f"\nBest method: {best_method}")
    print(f"CV per-fold f1: {all_cv[best_method]['cv_macro_f1_mean']:.4f}+-{all_cv[best_method]['cv_macro_f1_std']:.4f}")
    print(f"CV pooled  f1: {all_pooled[best_method]['macro_f1']:.4f}")
    print(f"Test       f1: {best_test['test_metrics']['macro_f1']:.4f}")
    return cv_summary, test_summary, comp


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--resize", choices=["letterbox", "squeeze"], required=True)
    args = ap.parse_args()

    # Map "squeeze" arg to "resized" directory name
    resize_dir = "resized" if args.resize == "squeeze" else args.resize
    logits_dir, results_dir = get_paths(resize_dir)
    results_dir.mkdir(parents=True, exist_ok=True)

    crt_cv, crt_test, crt_comp = run_ensemble("cRT", "", logits_dir, results_dir)
    plain_cv, plain_test, plain_comp = run_ensemble("plain", "plain_", logits_dir, results_dir)

    # Consolidated
    all_results = {
        "backbones": BACKBONES, "resize": args.resize,
        "crt": {"cv_summary": crt_cv, "test_summary": crt_test, "comparison": crt_comp},
        "plain": {"cv_summary": plain_cv, "test_summary": plain_test, "comparison": plain_comp},
    }
    (results_dir / "ensemble_all.json").write_text(json.dumps(all_results, indent=2))
    print("\n=== All 4-model ensembles complete ===")


if __name__ == "__main__":
    main()