"""The eight levels: how each dataset is generated (fixed seeds), what the shortcut is,
how it is removed, and how the before/after numbers are measured.

    make_all()            -> list of level dicts, exactly what ships as web/levels/*.json
    materialize(level)    -> train/test arrays for the "before" state (fixed=True: after the fix)
    run_level(level)      -> the numbers the game shows; web/engine.js does the same thing

Every number is deterministic: the data are generated once here with numpy and the
models in hunt/models.py use no random numbers at all.
"""
from __future__ import annotations

import numpy as np

from . import images as im
from .models import (accuracy, auc, balanced_accuracy, confusion, knn1, logreg_fit, logreg_predict,
                     logreg_scores, recall, stump_fit, stump_predict)


def _sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))


def _z(x):
    return (x - x.mean()) / x.std()


def _round(X, decimals):
    return [[round(float(v), d) if d is not None else int(v) for v, d in zip(row, decimals)] for row in X]


# ----------------------------------------------------------------------------- 1
def make_id_level(seed=1, n=600):
    """Row IDs handed out after the file was sorted by label."""
    rng = np.random.default_rng(seed)
    age = rng.integers(21, 71, n)
    income = np.round(np.exp(rng.normal(10.7, 0.45, n)) / 100) * 100
    debt_ratio = np.round(rng.beta(2, 5, n), 2)
    late = rng.poisson(1.0, n)
    logit = -1.3 + 0.9 * _z(debt_ratio) + 0.7 * _z(late) - 0.4 * _z(np.log(income)) - 0.2 * _z(age)
    y = (rng.random(n) < _sigmoid(logit)).astype(int)
    # the leak: ids assigned in label order (all the "no" rows first), then the file was shuffled
    order = np.argsort(y, kind="stable")
    ids = np.empty(n, dtype=int)
    ids[order] = 10001 + np.arange(n)
    perm = rng.permutation(n)
    X = np.column_stack([ids, age, income, debt_ratio, late])[perm]
    y = y[perm]
    ntr = int(0.7 * n)
    dec = [None, None, None, 2, None]
    return {
        "id": "id", "n": 1, "kind": "table", "model": "stump",
        "features": ["id", "age", "income", "debt_ratio", "late_payments"],
        "classes": ["repaid", "defaulted"], "target": "defaulted",
        "train": {"X": _round(X[:ntr], dec), "y": y[:ntr].tolist()},
        "test": {"X": _round(X[ntr:], dec), "y": y[ntr:].tolist()},
        "shortcut": "id", "fix": {"type": "drop", "columns": ["id"]},
    }


# ----------------------------------------------------------------------------- 2
def make_target_leak_level(seed=2, n=600):
    """A column that only gets filled in after the outcome: the cancellation fee."""
    rng = np.random.default_rng(seed)
    tenure = rng.integers(1, 73, n)
    fee = np.round(rng.choice([9.99, 14.99, 24.99, 39.99], n, p=[0.3, 0.35, 0.25, 0.1]), 2)
    tickets = rng.poisson(0.8, n)
    logins = rng.poisson(9, n)
    logit = -0.9 - 0.6 * _z(tenure) + 0.5 * _z(tickets) - 0.7 * _z(logins) + 0.2 * _z(fee)
    y = (rng.random(n) < _sigmoid(logit)).astype(int)
    cancel_fee = np.zeros(n)
    churned = y == 1
    charged = churned & (rng.random(n) < 0.93)
    cancel_fee[charged] = np.round(fee[charged] * rng.uniform(1.0, 2.0, charged.sum()), 2)
    odd = (~churned) & (rng.random(n) < 0.03)          # a few plan-change fees on stayers
    cancel_fee[odd] = np.round(rng.uniform(5, 20, odd.sum()), 2)
    X = np.column_stack([tenure, fee, tickets, logins, cancel_fee])
    ntr = int(0.7 * n)
    dec = [None, 2, None, None, 2]
    return {
        "id": "target-leak", "n": 2, "kind": "table", "model": "logreg",
        "features": ["tenure_months", "monthly_fee", "support_tickets", "logins_30d", "cancel_fee"],
        "classes": ["stayed", "churned"], "target": "churned",
        "train": {"X": _round(X[:ntr], dec), "y": y[:ntr].tolist()},
        "test": {"X": _round(X[ntr:], dec), "y": y[ntr:].tolist()},
        "shortcut": "cancel_fee", "fix": {"type": "drop", "columns": ["cancel_fee"]},
    }


# ----------------------------------------------------------------------------- 3
def make_dejavu_level(seed=3, ntr=300, ndup=240, nfresh=80):
    """The test file is mostly near-copies of training rows."""
    rng = np.random.default_rng(seed)
    names = ["acidity", "sugar", "chlorides", "sulfur", "density", "alcohol"]
    mu0 = np.array([7.2, 2.5, 0.08, 46, 0.996, 10.4])
    mu1 = np.array([6.8, 6.4, 0.045, 138, 0.994, 10.5])
    sd = np.array([1.1, 4.2, 0.03, 50, 0.003, 1.2])
    shift = 0.55  # how far apart the classes are, in standard deviations (honest 1-NN ~ 70%)

    def draw(k):
        y = rng.integers(0, 2, k)
        base = np.where(y[:, None] == 1, mu1, mu0)
        X = (mu0 + mu1) / 2 + (base - (mu0 + mu1) / 2) * shift + rng.normal(0, 1, (k, 6)) * sd
        return X, y

    Xtr, ytr = draw(ntr)
    src = rng.integers(0, ntr, ndup)
    Xdup = Xtr[src] + rng.normal(0, 1, (ndup, 6)) * sd * 0.02
    ydup = ytr[src]
    Xf, yf = draw(nfresh)
    Xte = np.vstack([Xdup, Xf])
    yte = np.concatenate([ydup, yf])
    perm = rng.permutation(len(yte))
    Xte, yte = Xte[perm], yte[perm]
    dec = [2, 2, 3, 1, 4, 2]
    return {
        "id": "dejavu", "n": 3, "kind": "table", "model": "knn1",
        "features": names, "classes": ["red", "white"], "target": "colour",
        "train": {"X": _round(Xtr, dec), "y": ytr.tolist()},
        "test": {"X": _round(Xte, dec), "y": yte.tolist()},
        "shortcut": "split", "fix": {"type": "dedupe", "eps": 0.2},
    }


# ----------------------------------------------------------------------------- 4, 5
def _shape_record(rng, kind, fg, bg, y, mark, radius=(4.5, 7.5)):
    r = float(rng.uniform(radius[0], radius[1]))
    cx = float(rng.uniform(9, 18))
    cy = float(rng.uniform(9, 18))
    return {"b": im.pack_mask(im.shape_mask(kind, cx, cy, r)), "fg": int(fg), "bg": int(bg), "y": int(y), "mk": int(mark)}


def make_watermark_level(seed=4, ntr=120, nte=60):
    """Squares carry a faint 2x2 mark in the top-right corner; circles don't."""
    rng = np.random.default_rng(seed)

    def draw(k):
        out = []
        for i in range(k):
            y = i % 2
            fg = int(rng.integers(9, 16))
            out.append(_shape_record(rng, "square" if y else "circle", fg, 0, y, mark=(y == 1)))
        idx = rng.permutation(k)
        return [out[i] for i in idx]

    return {
        "id": "watermark", "n": 4, "kind": "image", "model": "logreg", "size": im.SIZE,
        "shapes": ["circle", "square"], "classes": ["circle", "square"], "target": "shape",
        "mark": {"cells": [list(c) for c in im.MARK_CELLS], "value": im.MARK_VALUE},
        "train": draw(ntr), "test": draw(nte),
        "shortcut": "corner", "fix": {"type": "mask_mark"},
    }


def make_background_level(seed=5, ntr=120, nte=60):
    """Circles mostly on dark backgrounds, triangles mostly on light ones."""
    rng = np.random.default_rng(seed)

    def draw(k):
        out = []
        for i in range(k):
            y = i % 2
            flip = rng.random() < 0.05
            dark = (y == 0) != flip
            bg = int(rng.integers(0, 4)) if dark else int(rng.integers(5, 9))
            fg = int(rng.integers(12, 16))
            out.append(_shape_record(rng, "triangle" if y else "circle", fg, bg, y, mark=False, radius=(3.5, 8.5)))
        idx = rng.permutation(k)
        recs = [out[i] for i in idx]
        # the fix: the same backgrounds, dealt out at random so they say nothing about the class
        fixed = [int(b) for b in rng.permutation([r["bg"] for r in recs])]
        return recs, fixed

    train, bg_tr = draw(ntr)
    test, bg_te = draw(nte)
    return {
        "id": "background", "n": 5, "kind": "image", "model": "logreg", "size": im.SIZE,
        "shapes": ["circle", "triangle"], "classes": ["circle", "triangle"], "target": "shape",
        "train": train, "test": test, "bg_fixed": {"train": bg_tr, "test": bg_te},
        "shortcut": "background", "fix": {"type": "shuffle_bg"},
    }


# ----------------------------------------------------------------------------- 6
MONTH_RATES = [0.08, 0.10, 0.75, 0.80, 0.10, 0.10, 0.80, 0.85, 0.12, 0.15, 0.75, 0.80]


def make_time_level(seed=6, n=720):
    """Return rates jump around from month to month (promotions). With the month as a
    category and a random split, the model learns each month's rate; with a time split the
    test months are ones it has never seen."""
    rng = np.random.default_rng(seed)
    month = np.sort(rng.integers(1, 13, n))
    price = np.round(np.exp(rng.normal(3.6, 0.6, n)), 2)
    discount = np.round(rng.choice([0, 0, 0, 10, 20, 30], n), 0)
    is_gift = (rng.random(n) < 0.2).astype(int)
    prior_orders = rng.poisson(3, n)
    honest = 0.5 * _z(np.log(price)) + 0.35 * _z(discount) + 0.4 * is_gift - 0.4 * _z(prior_orders)
    base = np.array(MONTH_RATES)[month - 1]
    logit = np.log(base / (1 - base)) + honest
    y = (rng.random(n) < _sigmoid(logit)).astype(int)
    X = np.column_stack([month, price, discount, is_gift, prior_orders])
    perm = rng.permutation(n)
    ntr = int(0.7 * n)
    dec = [None, 2, None, None, None]
    return {
        "id": "timestamp", "n": 6, "kind": "table", "model": "logreg",
        "features": ["month", "price", "discount_pct", "is_gift", "prior_orders"],
        "categorical": {"month": list(range(1, 13))},
        "classes": ["kept", "returned"], "target": "returned",
        "all": {"X": _round(X, dec), "y": y.tolist()},
        "random_train": sorted(perm[:ntr].tolist()), "random_test": sorted(perm[ntr:].tolist()),
        "time_cutoff": 9,
        "shortcut": "month", "fix": {"type": "time_split", "time_feature": "month"},
    }


# ----------------------------------------------------------------------------- 7
VOCAB = ["login", "password", "billing", "refund", "charge", "crash", "slow", "error", "question",
         "thanks", "angry", "urgent", "resolved", "closed", "escalated", "tier"]


def make_notes_level(seed=7, n=500):
    """The agent's free-text note says how the ticket ended."""
    rng = np.random.default_rng(seed)
    wait = np.round(rng.exponential(12, n) + 1, 0)
    prior = rng.poisson(1.5, n)
    tier = rng.choice([1, 2, 3], n, p=[0.5, 0.35, 0.15])
    logit = -1.0 + 0.6 * _z(np.log(wait)) + 0.4 * _z(prior) - 0.35 * _z(tier)
    y = (rng.random(n) < _sigmoid(logit)).astype(int)
    topics = [["login", "password"], ["billing", "charge"], ["billing", "refund"], ["crash", "error"],
              ["slow"], ["question"]]
    moods = ["angry", "urgent", "thanks", ""]
    texts, counts = [], np.zeros((n, len(VOCAB)), dtype=int)
    for i in range(n):
        words = list(topics[rng.integers(0, len(topics))])
        mood = moods[rng.integers(0, 4)] if y[i] == 0 else moods[rng.choice(4, p=[0.35, 0.3, 0.1, 0.25])]
        if mood:
            words.append(mood)
        if y[i] == 1:
            words += ["escalated", "tier"] if rng.random() < 0.92 else ["closed"]
        else:
            words += ["resolved"] if rng.random() < 0.55 else ["closed"] if rng.random() < 0.8 else ["escalated"]
        texts.append(" ".join(words))
        for w in words:
            counts[i, VOCAB.index(w)] += 1
    X = np.column_stack([wait, prior, tier, counts])
    ntr = int(0.7 * n)
    dec = [None, None, None] + [None] * len(VOCAB)
    return {
        "id": "notes", "n": 7, "kind": "table", "model": "logreg",
        "features": ["wait_minutes", "prior_tickets", "plan_tier"] + ["notes:" + w for w in VOCAB],
        "classes": ["handled", "escalated"], "target": "escalated",
        "text_column": "notes", "text": {"train": texts[:ntr], "test": texts[ntr:]},
        "train": {"X": _round(X[:ntr], dec), "y": y[:ntr].tolist()},
        "test": {"X": _round(X[ntr:], dec), "y": y[ntr:].tolist()},
        "shortcut": "notes", "fix": {"type": "drop_prefix", "prefix": "notes:"},
    }


# ----------------------------------------------------------------------------- 8
def make_imbalance_level(seed=8, n=800):
    """Five percent of machines fail. A model that never says "fails" is 95% accurate."""
    rng = np.random.default_rng(seed)
    hours = np.round(rng.uniform(100, 9000, n), 0)
    temp = np.round(rng.normal(62, 6, n), 1)
    vibration = np.round(rng.gamma(2, 0.8, n), 2)
    alerts = rng.poisson(0.6, n)
    logit = -3.6 + 0.6 * _z(hours) + 0.7 * _z(temp) + 0.6 * _z(vibration) + 0.5 * _z(alerts)
    y = (rng.random(n) < _sigmoid(logit)).astype(int)
    X = np.column_stack([hours, temp, vibration, alerts])
    ntr = int(0.7 * n)
    dec = [None, 1, 2, None]
    return {
        "id": "imbalance", "n": 8, "kind": "table", "model": "logreg",
        "features": ["run_hours", "temperature", "vibration", "alerts_30d"],
        "classes": ["fine", "failed"], "target": "failed",
        "train": {"X": _round(X[:ntr], dec), "y": y[:ntr].tolist()},
        "test": {"X": _round(X[ntr:], dec), "y": y[ntr:].tolist()},
        "shortcut": "balance", "fix": {"type": "metric"},
    }


def make_all():
    return [make_id_level(), make_target_leak_level(), make_dejavu_level(), make_watermark_level(),
            make_background_level(), make_time_level(), make_notes_level(), make_imbalance_level()]


# ------------------------------------------------------------------ materialize
def image_matrix(recs, mark_override=None, bgs=None):
    rows = []
    for i, r in enumerate(recs):
        rows.append(im.compose_record(r, mark_override, None if bgs is None else bgs[i]))
    return np.array(rows)


def expand(X: np.ndarray, feats: list, categorical: dict):
    """One-hot the columns listed in `categorical` (feature "month" -> "month=1" ... "month=12")."""
    cols, names = [], []
    for j, f in enumerate(feats):
        if f in categorical:
            for c in categorical[f]:
                cols.append((X[:, j] == c).astype(float))
                names.append("%s=%s" % (f, c))
        else:
            cols.append(X[:, j])
            names.append(f)
    return np.column_stack(cols), names


def materialize(level: dict, fixed: bool = False, drop_time: bool = False) -> dict:
    """Train/test arrays for a level, before or after its fix."""
    fix = level["fix"]
    kind = level["kind"]
    if kind == "image":
        if fix["type"] == "mask_mark":
            mo = 0 if fixed else None
            Xtr, Xte = image_matrix(level["train"], mo), image_matrix(level["test"], mo)
        else:  # shuffle_bg
            btr = level["bg_fixed"]["train"] if fixed else None
            bte = level["bg_fixed"]["test"] if fixed else None
            Xtr, Xte = image_matrix(level["train"], None, btr), image_matrix(level["test"], None, bte)
        ytr = np.array([r["y"] for r in level["train"]])
        yte = np.array([r["y"] for r in level["test"]])
        feats = ["px%d" % i for i in range(Xtr.shape[1])]
        return {"Xtr": Xtr, "ytr": ytr, "Xte": Xte, "yte": yte, "features": feats}

    feats = list(level["features"])
    if fix["type"] == "time_split":
        X = np.array(level["all"]["X"], dtype=float)
        y = np.array(level["all"]["y"])
        if fixed:
            t = X[:, feats.index(fix["time_feature"])]
            tr, te = t <= level["time_cutoff"], t > level["time_cutoff"]
        else:
            tr, te = np.array(level["random_train"]), np.array(level["random_test"])
        if drop_time:
            keep = [j for j, f in enumerate(feats) if f != fix["time_feature"]]
            X = X[:, keep]
            feats = [feats[j] for j in keep]
        X, feats = expand(X, feats, level.get("categorical", {}))
        return {"Xtr": X[tr], "ytr": y[tr], "Xte": X[te], "yte": y[te], "features": feats}

    Xtr = np.array(level["train"]["X"], dtype=float)
    ytr = np.array(level["train"]["y"])
    Xte = np.array(level["test"]["X"], dtype=float)
    yte = np.array(level["test"]["y"])
    out = {"Xtr": Xtr, "ytr": ytr, "Xte": Xte, "yte": yte, "features": feats}
    if not fixed:
        return out
    if fix["type"] == "drop":
        keep = [j for j, f in enumerate(feats) if f not in fix["columns"]]
    elif fix["type"] == "drop_prefix":
        keep = [j for j, f in enumerate(feats) if not f.startswith(fix["prefix"])]
    elif fix["type"] == "dedupe":
        _, dist = knn1(Xtr, ytr, Xte)
        fresh = dist > fix["eps"]
        out.update({"Xte": Xte[fresh], "yte": yte[fresh], "kept": fresh})
        return out
    else:  # metric: same data, different yardstick
        return out
    out.update({"Xtr": Xtr[:, keep], "Xte": Xte[:, keep], "features": [feats[j] for j in keep]})
    return out


# ---------------------------------------------------------------------- evaluate
def evaluate(level: dict, data: dict) -> dict:
    """Train the level's model and score it on the test rows. Returns what the charts need."""
    m = level["model"]
    Xtr, ytr, Xte, yte = data["Xtr"], data["ytr"], data["Xte"], data["yte"]
    if m == "logreg":
        model = logreg_fit(Xtr, ytr)
        pred = logreg_predict(model, Xte)
        res = {"acc": accuracy(pred, yte), "weights": model["w"].tolist(), "bias": model["b"],
               "train_acc": accuracy(logreg_predict(model, Xtr), ytr),
               "auc": auc(logreg_scores(model, Xte), yte)}
    elif m == "knn1":
        pred, dist = knn1(Xtr, ytr, Xte)
        res = {"acc": accuracy(pred, yte), "dist": dist.tolist()}
    else:
        model = stump_fit(Xtr, ytr)
        pred = stump_predict(model, Xte)
        res = {"acc": accuracy(pred, yte), "feature": data["features"][model["feature"]],
               "threshold": model["threshold"], "above": model["above"], "per_feature": model["per_feature"],
               "train_acc": accuracy(stump_predict(model, Xtr), ytr)}
    res.update({"n_test": int(len(yte)), "features": list(data["features"]),
                "confusion": confusion(pred, yte), "balanced_acc": balanced_accuracy(pred, yte), "recall": recall(pred, yte)})
    return res


def run_level(level: dict) -> dict:
    before = evaluate(level, materialize(level, False))
    after_data = materialize(level, True)
    after = evaluate(level, after_data)
    out = {"id": level["id"], "before": before, "after": after}
    fix = level["fix"]["type"]
    if fix == "time_split":
        out["after_no_time"] = evaluate(level, materialize(level, True, drop_time=True))
    if fix == "dedupe":
        out["kept"] = int(after_data["kept"].sum())
        out["dropped"] = int((~after_data["kept"]).sum())
    if fix == "metric":
        yte = np.array(level["test"]["y"])
        out["majority_acc"] = float(max(yte.mean(), 1 - yte.mean()))
    if level["kind"] == "image":
        out["counterfactual"] = counterfactuals(level, before["weights"], before["bias"])
    return out


def counterfactuals(level: dict, w, b) -> list:
    """For the picture levels: take one test image of each class, change only the shortcut,
    and see what the shortcut-trained model says."""
    data = materialize(level, False)
    model = logreg_fit(data["Xtr"], data["ytr"])
    recs = level["test"]
    out = []
    for cls in (0, 1):
        i = next(k for k, r in enumerate(recs) if r["y"] == cls)
        r = recs[i]
        if level["fix"]["type"] == "mask_mark":
            a = im.compose_record(r)
            bimg = im.compose_record(r, mark_override=0 if r["mk"] else 1)
            change = "mark removed" if r["mk"] else "mark added"
        else:
            a = im.compose_record(r)
            other = 7 if r["bg"] <= 3 else 1
            bimg = im.compose_record(r, bg_override=other)
            change = "background set to %d" % other
        pa, pb = logreg_scores(model, a[None])[0], logreg_scores(model, bimg[None])[0]
        out.append({"index": i, "y": cls, "change": change, "p_before": float(pa), "p_after": float(pb)})
    return out