"""Checks that need a target column: class balance, row-order leaks, single-feature leakage,
conflicting labels and missingness that predicts the label."""
import math

from .fmt import fmt_pct, plural, row_nums
from .jscompat import div, js_max, js_round, js_str, py_str, to_fixed, trim
from .profiler import top_order
from .stats import auc, fold_of, mean, spearman
from .values import is_missing_token

K_FOLDS = 5
M_SMOOTH = 2


def encode_feature(p, n):
    """Small integer codes for one feature: numbers -> up to 10 quantile bins, categories -> the
    50 most common values; 0 = missing, 1 = any other value."""
    codes = [0] * n
    if p.type in ("numeric", "binary") and p.sorted:
        s = p.sorted
        edges = [s[(k * len(s)) // 10] for k in range(1, 10)]
        edges = [e for i, e in enumerate(edges) if i == 0 or e != edges[i - 1]]
        sentinels = p.sentinels
        for i in range(n):
            x = p.num[i]
            if x is None or x in sentinels:
                codes[i] = 0
                continue
            b = 0
            while b < len(edges) and x >= edges[b]:
                b += 1
            codes[i] = b + 1
        return codes
    top = {v: j + 2 for j, v in enumerate(top_order(p.counts)[:50])}
    for i in range(n):
        v = p.raw[i]
        codes[i] = 0 if is_missing_token(v) else top.get(v, 1)
    return codes


def _strength_classification(codes, y, rows_used, fold, classes):
    cls = classes[:10] if len(classes) > 10 else classes
    ci = {c: k for k, c in enumerate(cls)}
    C = len(cls)
    tab = {}                                   # code -> [per-fold class counts or None]
    tot_f = [0] * K_FOLDS
    prior_f = [[0] * C for _ in range(K_FOLDS)]
    for i in rows_used:
        c = ci.get(y[i])
        if c is None:
            continue
        arr = tab.get(codes[i])
        if arr is None:
            arr = tab[codes[i]] = [None] * K_FOLDS
        fa = arr[fold[i]]
        if fa is None:
            fa = arr[fold[i]] = [0] * C
        fa[c] += 1
        prior_f[fold[i]][c] += 1
        tot_f[fold[i]] += 1
    totals = {}
    for key, arr in tab.items():
        t2 = [0] * C
        for fa in arr:
            if fa:
                for c in range(C):
                    t2[c] += fa[c]
        totals[key] = t2
    prior, tot_all = [0] * C, 0
    for f in range(K_FOLDS):
        for c in range(C):
            prior[c] += prior_f[f][c]
            tot_all += prior_f[f][c]
    # out-of-fold class probabilities, m-estimate smoothing toward the training prior (m = 2)
    probs, labs = [], []
    for i in rows_used:
        c = ci.get(y[i])
        if c is None:
            continue
        key, f = codes[i], fold[i]
        tot, own = totals[key], tab[key][f]
        train_n = tot_all - tot_f[f]
        cnt2 = 0
        for c4 in range(C):
            cnt2 += tot[c4] - (own[c4] if own else 0)
        pr = []
        for c5 in range(C):
            pc = div(prior[c5] - prior_f[f][c5], train_n)
            pr.append(div((tot[c5] - (own[c5] if own else 0)) + M_SMOOTH * pc, cnt2 + M_SMOOTH))
        probs.append(pr)
        labs.append(c)
    # macro one-vs-rest AUC (binary: the AUC of the minority class)
    aucs = []
    for k in ([C - 1] if C == 2 else range(C)):
        sc = [pr[k] for pr in probs]
        lb = [1 if l == k else 0 for l in labs]
        s1 = sum(lb)
        if not s1 or s1 == len(lb):
            continue
        aucs.append(auc(sc, lb))
    return (mean(aucs) if aucs else 0.5), "AUC"


def _strength_regression(codes, y, rows_used, fold):
    sums, cnts = {}, {}
    sum_f, cnt_f = [0] * K_FOLDS, [0] * K_FOLDS
    for i in rows_used:
        key = (codes[i], fold[i])
        sums[key] = sums.get(key, 0) + y[i]
        cnts[key] = cnts.get(key, 0) + 1
        sum_f[fold[i]] += y[i]
        cnt_f[fold[i]] += 1
    code_sum, code_cnt = {}, {}
    for key in sums:                           # insertion order, like Object.keys on "code|fold"
        c = key[0]
        code_sum[c] = code_sum.get(c, 0) + sums[key]
        code_cnt[c] = code_cnt.get(c, 0) + cnts[key]
    tot_s = 0
    for v in sum_f:
        tot_s += v
    tot_c = 0
    for v in cnt_f:
        tot_c += v
    sse = sst = 0
    ym = div(tot_s, tot_c)
    for i in rows_used:
        c, f = codes[i], fold[i]
        kk = (c, f)
        train_mean = div(tot_s - sum_f[f], tot_c - cnt_f[f])
        s = code_sum.get(c, 0) - sums.get(kk, 0)
        m = code_cnt.get(c, 0) - cnts.get(kk, 0)
        pred = div(s + M_SMOOTH * train_mean, m + M_SMOOTH)
        sse += (y[i] - pred) * (y[i] - pred)
        sst += (y[i] - ym) * (y[i] - ym)
    return (1 - sse / sst if sst > 0 else 0), "R²"


def analyze_target(t, prof, n, issues, table, strength, max_rows=None):
    info = {"column": t.name}
    t_sent = t.sentinels or []
    labeled = [i for i in t.present_idx if t.num[i] not in t_sent]
    if t.missing:
        issues.add(type="target_missing", column=t.name, severity="medium" if t.missing > 0.05 * n else "low",
                   affected=t.missing / n,
                   title=plural(t.missing, "row") + " without a " + t.name + " value",
                   detail="Rows with no label can't be used to train or evaluate. Drop them (or keep them aside for later labelling).",
                   fix="df = df.dropna(subset=[" + py_str(t.name) + "])", action={"dropMissingTarget": t.name})

    classes = None
    y = [None] * n
    is_class = t.type in ("binary", "categorical") or (t.type == "numeric" and t.unique <= 10 and t.integer)
    if t.type in ("text", "date"):
        info["kind"] = "unsupported"
        return info
    if is_class:
        info["kind"] = "classification"
        cnt = {}
        for i in labeled:
            v = trim(t.raw[i])
            if t.mapping and t.mapping.get(t.raw[i]):
                v = trim(t.mapping[t.raw[i]])
            cnt[v] = cnt.get(v, 0) + 1
            y[i] = v
        classes = top_order(cnt)
        info["classes"] = [[c, cnt[c]] for c in classes]
        minority = div(cnt[classes[-1]], len(labeled)) if classes else math.nan
        info["minorityShare"] = minority
        if len(classes) >= 2 and minority < 0.2:
            sev = "high" if minority < 0.01 else "medium" if minority < 0.05 else "low"
            issues.add(type="class_imbalance", column=t.name, severity=sev, affected=0,
                       title='"' + classes[-1] + '" is only ' + fmt_pct(minority) + " of " + t.name,
                       detail="A model can score " + fmt_pct(1 - minority) + " accuracy by never predicting it. Judge models with recall, precision or PR-AUC, split with stratification, and consider class weights.",
                       fix="train, test = train_test_split(df, test_size=0.2, stratify=df[" + py_str(t.name) + "], random_state=0)")
    else:
        info["kind"] = "regression"
        for i in labeled:
            y[i] = t.num[i]
        info["stats"] = t.stats
    if len(labeled) < 30:
        return info

    _check_order(t, labeled, y, is_class, classes, issues)

    # ----- leakage: out-of-fold predictive power of each feature on its own -----
    fold = [fold_of(i, K_FOLDS) for i in range(n)]
    cap = max_rows or 100000
    if len(labeled) > cap:
        step = math.ceil(len(labeled) / cap)
        rows_used = [i for k, i in enumerate(labeled) if k % step == 0]
    else:
        rows_used = labeled
    for p in prof:
        if p is t or p.type == "empty" or p.unique <= 1:
            continue
        codes = encode_feature(p, n)
        if is_class:
            score, metric = _strength_classification(codes, y, rows_used, fold, classes)
        else:
            score, metric = _strength_regression(codes, y, rows_used, fold)
        strength.append({"column": p.name, "score": score, "metric": metric})
    strength.sort(key=lambda s: -s["score"])
    # A leak stands out: its error (1 - AUC, or 1 - R²) is several times smaller than the
    # next-best feature's. Two or more near-perfect features are flagged together.
    hi_t, med_t, perfect = (0.98, 0.9, 0.995) if is_class else (0.95, 0.85, 0.99)
    n_perfect = sum(1 for s in strength if s["score"] >= perfect)
    for k, s in enumerate(strength):
        nxt = strength[1 if k == 0 else 0] if len(strength) > (1 if k == 0 else 0) else None
        ratio = div(1 - js_max(0, nxt["score"]), js_max(1e-6, 1 - s["score"])) if nxt else math.inf
        s["standout"] = ratio
        s["flag"] = bool((k == 0 and s["score"] >= med_t and ratio >= 3) or (n_perfect >= 2 and s["score"] >= perfect))
    names = [q.name for q in prof]
    for s in strength:
        p = prof[names.index(s["column"])]
        if not s["flag"]:
            continue
        if s["metric"] == "AUC":
            why = "On its own, with no model at all, it predicts " + t.name + " with AUC " + to_fixed(s["score"], 3) + " on held-out rows."
        else:
            why = "On its own it explains " + fmt_pct(s["score"]) + " of the variance in " + t.name + " on held-out rows."
        blank_hint = ""
        if is_class and p.missing > 0.05 * n and p.missing < 0.95 * n:
            miss_rate = {}
            for c in classes[:2]:
                miss_rate[c] = sum(1 for i in p.missing_idx if y[i] == c) / max(1, len(p.missing_idx))
            c0 = classes[0]
            if miss_rate[c0] > 0.97 or miss_rate[c0] < 0.03:
                blank_hint = " It is blank for " + ("almost every \"" + c0 + "\" row" if miss_rate[c0] > 0.97 else "almost no \"" + c0 + "\" row") + ": filled in only after the outcome happened?"
        so = s["standout"]
        standout = (" Its error rate is " + js_str(js_round(so)) + "× lower than the next-best column's."
                    if so < 1e5 and math.isfinite(so) else " No other column comes close.")
        issues.add(type="leakage", column=s["column"], severity="high" if s["score"] >= hi_t else "medium", affected=0,
                   title=s["column"] + " predicts " + t.name + " suspiciously well",
                   detail=why + standout + blank_hint + " Real single features rarely stand out like this. Check that it is known at prediction time and isn't recorded after (or derived from) the outcome. If it is a leak, your model will look great in testing and fail in use.",
                   evidence={"values": [s["metric"] + " " + to_fixed(s["score"], 3)]},
                   fix="df = df.drop(columns=[" + py_str(s["column"]) + "])  # only if it isn't available when you'll make predictions",
                   action={"drop": [s["column"]]})

    _check_conflicts(t, prof, n, labeled, y, is_class, strength, table, issues)
    if is_class and len(classes) == 2:
        _check_missing_informative(t, prof, n, labeled, y, classes, strength, issues)
    return info


def _check_order(t, labeled, y, is_class, classes, issues):
    """Does row position predict the target? (a sorted file split without shuffling)"""
    pos = list(labeled)
    if is_class and len(classes) >= 2:
        worst = None
        for c in classes[:5]:
            lab = [1 if y[i] == c else 0 for i in labeled]
            a = auc(pos, lab)
            if worst is None or abs(a - 0.5) > abs(worst["a"] - 0.5):
                worst = {"c": c, "a": a}
        if worst and abs(worst["a"] - 0.5) >= 0.08:
            npos = sum(1 for i in labeled if y[i] == worst["c"])
            nneg = len(labeled) - npos
            se = math.sqrt((npos + nneg + 1) / (12 * npos * nneg))
            z = abs(worst["a"] - 0.5) / se
            if z > 5:
                half = len(labeled) // 2
                r1 = sum(1 for i in labeled[:half] if y[i] == worst["c"]) / half
                r2 = sum(1 for i in labeled[half:] if y[i] == worst["c"]) / (len(labeled) - half)
                issues.add(type="order_leakage", column=t.name, severity="high" if abs(worst["a"] - 0.5) >= 0.15 else "medium", affected=0,
                           title="Rows are partly sorted by " + t.name,
                           detail='"' + worst["c"] + '" makes up ' + fmt_pct(r1) + " of the first half of the file and " + fmt_pct(r2) + " of the second. Split without shuffling and train and test get different labels; a model can also learn row order through any ID or timestamp.",
                           evidence={"values": ["first half " + fmt_pct(r1), "second half " + fmt_pct(r2), "AUC of row position " + to_fixed(worst["a"], 2)]},
                           fix="df = df.sample(frac=1, random_state=0).reset_index(drop=True)  # shuffle before splitting",
                           action={"shuffle": True})
    elif not is_class:
        rho = spearman(pos, [y[i] for i in labeled])
        if abs(rho) >= 0.3:
            issues.add(type="order_leakage", column=t.name, severity="high" if abs(rho) >= 0.6 else "medium", affected=0,
                       title="Rows are sorted by " + t.name,
                       detail="Row position correlates with the target (Spearman " + to_fixed(rho, 2) + "). Shuffle before splitting.",
                       fix="df = df.sample(frac=1, random_state=0).reset_index(drop=True)", action={"shuffle": True})


def _check_conflicts(t, prof, n, labeled, y, is_class, strength, table, issues):
    """Groups of rows with identical features (ignoring IDs and leaks) but different labels."""
    leak_cols = [s["column"] for s in strength if s["flag"]]
    feat_idx = [p.index for p in prof if p is not t and not p.is_id and p.name not in leak_cols]
    if not feat_idx:
        return
    groups = {}
    for i in labeled:
        r = table.rows[i]
        key = "\x01".join(r[j] for j in feat_idx)
        g = groups.get(key)
        if g is None:
            g = groups[key] = {"rows": [], "labels": set()}
        g["rows"].append(i)
        g["labels"].add(y[i] if is_class else js_str(y[i]))
    conflict_rows, conflict_groups = [], 0
    for g in groups.values():
        if len(g["labels"]) > 1:
            conflict_groups += 1
            conflict_rows.extend(g["rows"])
    if not conflict_groups:
        return
    conflict_rows.sort()
    any_id = any(p.is_id for p in prof)
    ignoring = (" (ignoring IDs" + (" and leaking columns" if leak_cols else "") + ")") if any_id else (" (ignoring leaking columns)" if leak_cols else "")
    issues.add(type="conflicting_labels", column=t.name, severity="high" if conflict_groups > 0.01 * n else "medium",
               affected=len(conflict_rows) / n,
               title=plural(conflict_groups, "group") + " of identical rows with different " + t.name,
               detail="Same values in every feature" + ignoring + ", but a different label. Either the label is noisy or a feature that explains the difference is missing. No model can get all of them right.",
               evidence={"rows": row_nums(conflict_rows, 8)},
               fix="feats = [c for c in df.columns if c != " + py_str(t.name) + "]\nconflicts = df[df.groupby(feats)[" + py_str(t.name) + "].transform(\"nunique\") > 1]")


def _check_missing_informative(t, prof, n, labeled, y, classes, strength, issues):
    """Is a column's blankness itself a strong signal for the (binary) target?"""
    pos_c = classes[1]
    for p in prof:
        if p is t or p.missing < 0.02 * n or p.missing > 0.98 * n:
            continue
        if any(s["column"] == p.name and s["flag"] for s in strength):
            continue
        mi = set(p.missing_idx)
        a = an = b = bn = 0
        for i in labeled:
            if i in mi:
                an += 1
                if y[i] == pos_c:
                    a += 1
            else:
                bn += 1
                if y[i] == pos_c:
                    b += 1
        if an < 10 or bn < 10:
            continue
        p1, p2, pp = a / an, b / bn, (a + b) / (an + bn)
        se = math.sqrt(pp * (1 - pp) * (1 / an + 1 / bn))
        z = (p1 - p2) / se if se > 0 else 0
        if abs(p1 - p2) >= 0.1 and abs(z) > 3.3:
            issues.add(type="missing_informative", column=p.name, severity="info", affected=0,
                       title="Whether " + p.name + " is blank predicts " + t.name,
                       detail='"' + pos_c + '" is ' + fmt_pct(p1) + " when " + p.name + " is missing and " + fmt_pct(p2) + " when it isn't. Useful signal, so keep a \"was missing\" flag rather than silently filling it, but make sure the blank isn't caused by the outcome.",
                       fix="df[" + py_str(p.name + "_missing") + "] = df[" + py_str(p.name) + "].isna()")