"""Supplementary diagnostics referenced in the Discussion (not part of the primary study).

Two secondary analyses, each pooled over one repeat of 5 folds (runtime measurement off):
  1. Effect of extreme lambda values on onto-MDL vs onto-grow guidance.
  2. The combined grow+prune+MDL variant compared against all other guidance variants.

These use a single repeat per point as an exploratory check rather than the full repeated
sweep behind the main results, so the numbers are indicative rather than definitive.
"""
import numpy as np, pandas as pd
from collections import defaultdict
from sklearn.model_selection import RepeatedStratifiedKFold
from sklearn.metrics import f1_score
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
from pathlib import Path

from run_all_experiments import (load_clean_dummies, make_fold_dfs, ontology_invalid_rate,
                                 N_FOLDS, SPLIT_SEED)
from oud_maak_verv_dataset import pollute_train_labels_my_version
from evaluate_ruleset import evaluate_ruleset
from RIPPER_no_onto_guid import run_RIPPER_without_guidance
from RIPPER_onto_guid import run_RIPPER_with_guidance
from RIPPER_onto_guid_grow import run_RIPPER_with_guidance_in_grow
from RIPPER_onto_guid_grow_and_prune import run_RIPPER_with_guidance_in_grow_and_prune
from RIPPER_onto_guid_grow_prune_mdl import run_RIPPER_with_guidance_in_grow_prune_mdl

X, y, all_classes = load_clean_dummies()
SPLITS = list(enumerate(RepeatedStratifiedKFold(n_splits=N_FOLDS, n_repeats=1,
                                                random_state=SPLIT_SEED).split(X, y), start=1))
FIGDIR = Path("results/figures/diagnostics"); FIGDIR.mkdir(parents=True, exist_ok=True)
RARE = "Nieuwe TFM Coupler"


def pooled_run(fit_fn, kw, level):
    """5 folds polluted at `level`; pool OOF preds -> one macro-F1/acc. Runtime OFF."""
    yt_all, yp_all, frames, totals = [], [], [], []
    rc = defaultdict(list)
    for run, (tri, tei) in SPLITS:
        tr, te = make_fold_dfs(X, y, tri, tei)
        tr, _ = pollute_train_labels_my_version(tr, level, run)
        np.random.seed(run)
        rs, sv = fit_fn(tr.copy(), **kw)
        for c, rules in rs.items():
            rc[c].append(len(rules))
        totals.append(sum(len(v) for v in rs.values()))
        res, _ = evaluate_ruleset(rs, te.copy(), sv, False)
        yt_all.append(res["y"].astype(str)); yp_all.append(res["y_pred"].astype(str))
        frames.append(res)
    yt = pd.concat(yt_all).values; yp = pd.concat(yp_all).values
    return {"f1_macro": f1_score(yt, yp, average="macro", zero_division=0),
            "accuracy": float((yt == yp).mean()),
            "invalid_rate": ontology_invalid_rate(pd.concat(frames)),
            "total_rules": float(np.mean(totals)),
            "rare_rules": float(np.mean(rc.get(RARE, [0]))),
            "rule_counts": {c: float(np.mean(v)) for c, v in rc.items()}}


# ============ EXPERIMENT 1: extreme lambda, MDL vs grow ============
print("=" * 92)
print("EXPERIMENT 1 — extreme lambda for onto-MDL and onto-grow (5-fold pooled, runtime OFF)")
print("=" * 92)
E1_LEVELS = [0, 50, 100]
E1_LAMBDAS = [4, 100, 100000]
E1_METHODS = [("onto MDL", run_RIPPER_with_guidance),
              ("onto grow", run_RIPPER_with_guidance_in_grow)]
print(f"{'method':>10} {'lambda':>8} {'poll':>5} {'macroF1':>8} {'acc':>7} "
      f"{'rules':>6} {'TFMrul':>7} {'invalid%':>9}")
e1 = []
for name, fn in E1_METHODS:
    for lam in E1_LAMBDAS:
        for lvl in E1_LEVELS:
            r = pooled_run(fn, {"lambda_val": lam}, lvl)
            e1.append({"method": name, "lambda": lam, "level": lvl, **r})
            print(f"{name:>10} {lam:>8} {lvl:>5} {r['f1_macro']:>8.3f} {r['accuracy']:>7.3f} "
                  f"{r['total_rules']:>6.1f} {r['rare_rules']:>7.1f} {r['invalid_rate']:>8.1%}")
    print("-" * 92)

# ============ EXPERIMENT 2: combined variant vs all, lambda=4, one run, full sweep ============
print("\n" + "=" * 92)
print("EXPERIMENT 2 — combined grow+prune+MDL vs all variants (lambda=4, 5-fold pooled, runtime OFF)")
print("=" * 92)
E2_LEVELS = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
E2_METHODS = [("no guidance",      run_RIPPER_without_guidance,                {}),
              ("onto MDL",         run_RIPPER_with_guidance,                   {"lambda_val": 4}),
              ("onto grow",        run_RIPPER_with_guidance_in_grow,           {"lambda_val": 4}),
              ("onto grow+prune",  run_RIPPER_with_guidance_in_grow_and_prune, {"lambda_val": 4}),
              ("grow+prune+MDL",   run_RIPPER_with_guidance_in_grow_prune_mdl, {"lambda_val": 4})]
e2 = []
for name, fn, kw in E2_METHODS:
    for lvl in E2_LEVELS:
        r = pooled_run(fn, kw, lvl)
        e2.append({"method": name, "level": lvl, **r})
e2df = pd.DataFrame(e2)

for metric, lab in [("f1_macro", "MACRO-F1"), ("accuracy", "ACCURACY"),
                    ("total_rules", "TOTAL RULES"), ("rare_rules", f"'{RARE}' RULES")]:
    print(f"\n{lab} vs pollution (rows=method):")
    piv = e2df.pivot(index="method", columns="level", values=metric).reindex([m[0] for m in E2_METHODS])
    fmt = "{:>6.3f}" if metric in ("f1_macro", "accuracy") else "{:>6.1f}"
    hdr = "".join(f"{l:>7}" for l in E2_LEVELS)
    print(f"{'method':>16}{hdr}")
    for m in piv.index:
        print(f"{m:>16}" + "".join(fmt.format(v) for v in piv.loc[m].values))

# figures
for metric, lab in [("f1_macro", "Pooled macro-F1"), ("accuracy", "Pooled accuracy"),
                    ("total_rules", "Total rules learned")]:
    plt.figure(figsize=(9, 6))
    for name, *_ in E2_METHODS:
        s = e2df[e2df["method"] == name]
        plt.plot(s["level"], s[metric], marker="o", label=name)
    plt.xlabel("% of training labels polluted"); plt.ylabel(lab)
    plt.title(f"{lab} vs pollution — incl. combined grow+prune+MDL (λ=4, 1 repeat/5 folds, runtime OFF)")
    plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
    out = FIGDIR / f"combined_variant_{metric}.png"; plt.savefig(out, dpi=140); plt.close()
    print(f"\nwrote {out}")

pd.DataFrame(e1).to_csv("results/diagnostics/extreme_lambda.csv", index=False)
e2df.drop(columns=["rule_counts"]).to_csv("results/diagnostics/combined_variant.csv", index=False)
print("\nwrote results/diagnostics/extreme_lambda.csv and combined_variant.csv")