"""Run every level in Python, print the before/after table, and draw the figures.

    python experiments/run_levels.py          # about 5 s; writes out/*.png and out/results.json

Also runs the CSV checks on the demo file and on three clean scikit-learn datasets, to show
which detectors fire on honest data (the false-positive study in the write-up).
"""
from __future__ import annotations

import json
import os
import sys
import time

import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from hunt import checks, images as im, levels  # noqa: E402
from hunt.models import logreg_fit  # noqa: E402

OUT = os.path.join(ROOT, "out")
BLUE, ORANGE, GRAY, INK = "#2a78d6", "#eb6834", "#9aa3b2", "#0E1A2E"
TITLES = {"id": "1 The ID that knew too much", "target-leak": "2 Tomorrow's weather, today", "dejavu": "3 Déjà vu",
          "watermark": "4 The watermark", "background": "5 Background noise", "timestamp": "6 The timestamp",
          "notes": "7 The label in the notes", "imbalance": "8 Accuracy that lies"}

plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 10, "axes.spines.top": False, "axes.spines.right": False,
                     "axes.edgecolor": GRAY, "axes.labelcolor": INK, "xtick.color": INK, "ytick.color": INK, "figure.dpi": 130})


def pct(x):
    return "%.1f%%" % (100 * x)


def fig_levels(results):
    ids = [r["id"] for r in results]
    before = [r["before"]["acc"] for r in results]
    after = [r["after"]["balanced_acc"] if r["id"] == "imbalance" else r["after"]["acc"] for r in results]
    fig, ax = plt.subplots(figsize=(8.4, 4.6))
    x = np.arange(len(ids))
    ax.bar(x - 0.2, before, 0.38, color=BLUE, label="with the shortcut (random held-out set)")
    ax.bar(x + 0.2, after, 0.38, color=ORANGE, label="shortcut removed, or the honest yardstick")
    for i, (b, a) in enumerate(zip(before, after)):
        ax.text(i - 0.2, b + 0.012, "%.0f" % (100 * b), ha="center", fontsize=8.5, color=INK)
        ax.text(i + 0.2, a + 0.012, "%.0f" % (100 * a), ha="center", fontsize=8.5, color=INK)
    ax.set_xticks(x)
    ax.set_xticklabels([TITLES[i].split(" ", 1)[1] for i in ids], rotation=20, ha="right", fontsize=8.5)
    ax.set_ylim(0, 1.08)
    ax.set_ylabel("test accuracy (level 8: balanced accuracy after)")
    ax.yaxis.grid(True, color="#e6e9ef")
    ax.set_axisbelow(True)
    ax.legend(frameon=False, loc="lower left", fontsize=8.5)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "levels.png"))
    plt.close(fig)


def fig_images(level, result, name):
    data = levels.materialize(level, False)
    model = logreg_fit(data["Xtr"], data["ytr"])
    w = np.array(model["w"]).reshape(im.SIZE, im.SIZE)
    fig, axes = plt.subplots(1, 4, figsize=(9.2, 2.7))
    for k, ax in enumerate(axes[:2]):
        rec = next(r for r in level["test"] if r["y"] == k)
        ax.imshow(im.compose_record(rec).reshape(im.SIZE, im.SIZE), cmap="gray", vmin=0, vmax=1)
        ax.set_title("a test " + level["classes"][k], fontsize=9)
    v = np.abs(w).max()
    axes[2].imshow(w, cmap="RdBu_r", vmin=-v, vmax=v)
    axes[2].set_title("logistic-regression weights", fontsize=9)
    cf = result["counterfactual"][0]
    rec = level["test"][cf["index"]]
    if level["fix"]["type"] == "mask_mark":
        img = im.compose_record(rec, mark_override=1)
    else:
        img = im.compose_record(rec, bg_override=7)
    axes[3].imshow(img.reshape(im.SIZE, im.SIZE), cmap="gray", vmin=0, vmax=1)
    axes[3].set_title("same %s, %s\nP(%s) %.2f -> %.2f" % (level["classes"][0], cf["change"], level["classes"][1], cf["p_before"], cf["p_after"]), fontsize=9)
    for ax in axes:
        ax.set_xticks([])
        ax.set_yticks([])
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, name))
    plt.close(fig)


def fig_dejavu(result):
    d = np.array(result["before"]["dist"])
    fig, ax = plt.subplots(figsize=(7.2, 3.2))
    bins = np.logspace(-2, 1, 40)
    ax.hist(d[d <= 0.2], bins=bins, color=BLUE, label="near-copy of a training row (%d)" % result["dropped"])
    ax.hist(d[d > 0.2], bins=bins, color=ORANGE, label="genuinely new row (%d)" % result["kept"])
    ax.axvline(0.2, color=GRAY, ls="--", lw=1)
    ax.set_xscale("log")
    ax.set_xlabel("distance from each test row to its nearest training row (standardized units)")
    ax.set_ylabel("test rows")
    ax.legend(frameon=False, fontsize=8.5)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "dejavu.png"))
    plt.close(fig)


def fig_time(level, result):
    X = np.array(level["all"]["X"])
    y = np.array(level["all"]["y"])
    months = np.arange(1, 13)
    rate = [y[X[:, 0] == m].mean() for m in months]
    fig, ax = plt.subplots(figsize=(7.2, 3.2))
    ax.bar(months, rate, color=[BLUE if m <= level["time_cutoff"] else ORANGE for m in months], width=0.7)
    ax.set_xticks(months)
    ax.set_xlabel("month")
    ax.set_ylabel("share of orders returned")
    ax.set_ylim(0, 1)
    ax.text(5, 0.95, "train months (time split)", ha="center", color=BLUE, fontsize=9)
    ax.text(11, 0.95, "test months", ha="center", color=ORANGE, fontsize=9)
    ax.set_title("random split %s  ·  time split %s  ·  time split without month %s" %
                 (pct(result["before"]["acc"]), pct(result["after"]["acc"]), pct(result["after_no_time"]["acc"])), fontsize=9.5)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "timestamp.png"))
    plt.close(fig)


def false_positive_study():
    """Single-feature AUC on clean datasets: how often would the leak hunt cry wolf?"""
    from sklearn import datasets
    rows = []
    for name, load in (("breast cancer", datasets.load_breast_cancer), ("iris", datasets.load_iris), ("wine", datasets.load_wine)):
        d = load()
        cols = list(d.feature_names) + ["target"]
        data = [[str(float(v)) for v in r] + [str(d.target_names[t])] for r, t in zip(d.data, d.target)]
        r = checks.analyze(cols, data, "target")
        rows.append({"dataset": name, "n": r["used_rows"], "positive": r["positive"],
                     "strengths": r["strengths"], "checks": [(c["id"], c["severity"]) for c in r["checks"]]})
    with open(os.path.join(ROOT, "web", "demo.csv"), encoding="utf-8") as f:
        cols, data = checks.read_csv(f.read())
    r = checks.analyze(cols, data, "left")
    rows.append({"dataset": "demo.csv (planted leak)", "n": r["used_rows"], "positive": r["positive"],
                 "strengths": r["strengths"], "checks": [(c["id"], c["severity"]) for c in r["checks"]]})
    fig, ax = plt.subplots(figsize=(7.6, 3.6))
    for k, row in enumerate(rows):
        aucs = [s["auc"] for s in row["strengths"]]
        ax.scatter(aucs, [k] * len(aucs), s=28, color=ORANGE if "demo" in row["dataset"] else BLUE, alpha=0.8, zorder=3)
        top = row["strengths"][0]
        ax.text(1.005, k, "%s %.3f" % (top["column"], top["auc"]), va="center", fontsize=8, color=INK)
    ax.axvline(0.95, color=GRAY, ls="--", lw=1)
    ax.axvline(0.90, color=GRAY, ls=":", lw=1)
    ax.text(0.953, len(rows) - 0.55, "flagged high", fontsize=8, color=GRAY)
    ax.text(0.897, len(rows) - 0.55, "medium", fontsize=8, color=GRAY, ha="right")
    ax.set_yticks(range(len(rows)))
    ax.set_yticklabels([r["dataset"] for r in rows], fontsize=9)
    ax.set_xlim(0.5, 1.2)
    ax.set_xticks([0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
    ax.set_xlabel("single-feature AUC (best 12 columns per dataset)")
    ax.xaxis.grid(True, color="#e6e9ef")
    ax.set_axisbelow(True)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "detectors.png"))
    plt.close(fig)
    return rows


def main():
    os.makedirs(OUT, exist_ok=True)
    t0 = time.time()
    results = []
    print("%-30s %10s %10s %s" % ("level", "before", "after", "notes"))
    for lv in levels.make_all():
        r = levels.run_level(lv)
        results.append(r)
        note = ""
        if r["id"] == "timestamp":
            note = "time split without month %s" % pct(r["after_no_time"]["acc"])
        if r["id"] == "dejavu":
            note = "%d near-copies dropped, %d fresh rows kept" % (r["dropped"], r["kept"])
        if r["id"] == "imbalance":
            note = "balanced acc %.3f, recall %.3f, always-'fine' %s" % (r["after"]["balanced_acc"], r["after"]["recall"], pct(r["majority_acc"]))
        if r["id"] == "id":
            note = "stump: %s > %g, then %s" % (r["before"]["feature"], r["before"]["threshold"], r["after"]["feature"])
        if lv["kind"] == "image":
            note = "; ".join("%s: P(%s) %.2f -> %.2f" % (c["change"], lv["classes"][1], c["p_before"], c["p_after"]) for c in r["counterfactual"])
        print("%-30s %10s %10s %s" % (TITLES[r["id"]], pct(r["before"]["acc"]), pct(r["after"]["acc"]), note))
    fig_levels(results)
    by = {r["id"]: r for r in results}
    lvs = {lv["id"]: lv for lv in levels.make_all()}
    fig_images(lvs["watermark"], by["watermark"], "watermark.png")
    fig_images(lvs["background"], by["background"], "background.png")
    fig_dejavu(by["dejavu"])
    fig_time(lvs["timestamp"], by["timestamp"])
    fp = false_positive_study()
    print("\nCSV checks on clean data (false positives) and the demo file:")
    for row in fp:
        print("  %-26s n=%d positive=%s  %s" % (row["dataset"], row["n"], row["positive"], ", ".join("%s:%s" % c for c in row["checks"] if c[1] != "ok")))
    slim = []
    for r in results:
        s = json.loads(json.dumps(r))
        for key in ("before", "after", "after_no_time"):
            if key in s:
                s[key].pop("dist", None)
                if len(s[key].get("weights", [])) > 50:
                    s[key]["weights"] = "784 pixel weights (see out/%s.png)" % r["id"]
        for c in s.get("counterfactual", []):
            c.pop("before", None)
            c.pop("after", None)
        slim.append(s)
    with open(os.path.join(OUT, "results.json"), "w", encoding="utf-8") as f:
        json.dump({"levels": slim, "false_positives": fp}, f, indent=1)
    print("\nwrote out/*.png and out/results.json in %.1f s" % (time.time() - t0))


if __name__ == "__main__":
    main()