"""Training file vs test file: column mismatch, overlapping rows, distribution shift, unseen values."""
from .checks import py_list
from .fmt import fmt_num, fmt_pct, plural, row_nums
from .jscompat import div, js_keys, py_str, to_exponential, to_fixed
from .profiler import profile_column
from .stats import chi_square_p, ks2, quantile_sorted


def _p_text(p):
    return "< 1e-6" if p < 1e-6 else to_exponential(p, 1)


def _biggest_change(keys, ca, cb, na, nb):
    """Sentence about the value whose share changed most (first one wins on ties)."""
    big = None
    for k in keys:
        d = div(cb.get(k, 0), nb) - div(ca.get(k, 0), na)
        if big is None or abs(d) > abs(big[1]):
            big = (k, d)
    k = big[0]
    return ("Biggest change: \"" + k + "\" is " + fmt_pct(div(ca.get(k, 0), na)) + " of training rows and "
            + fmt_pct(div(cb.get(k, 0), nb)) + " of test rows. ")


def compare_tables(train, prof, test, issues, target):
    res = {"rows": len(test.rows), "columns": []}
    tcols, n2 = test.columns, len(test.rows)
    missing_in_test = [c for c in train.columns if c not in tcols]
    extra_in_test = [c for c in tcols if c not in train.columns]
    if missing_in_test or extra_in_test:
        issues.add(type="column_mismatch", column=None, severity="high", affected=0,
                   title="The two files have different columns",
                   detail=("Only in training: " + ", ".join(missing_in_test) + ". " if missing_in_test else "") + ("Only in test: " + ", ".join(extra_in_test) + "." if extra_in_test else ""),
                   fix="common = [c for c in train.columns if c in test.columns]")

    # overlap: identical rows (ignoring ID columns and the label) and shared IDs
    shared = [c for c in train.columns if c in tcols]
    ids = [p.name for p in prof if p.is_id and p.name in shared]
    key_cols = [c for c in shared if c not in ids and (target is None or c != target.name)]
    ji = [train.columns.index(c) for c in key_cols]
    jt = [tcols.index(c) for c in key_cols]
    train_keys = {"\x01".join(r[j] for j in ji) for r in train.rows}
    overlap = [i for i, r in enumerate(test.rows) if "\x01".join(r[j] for j in jt) in train_keys]
    overlap_set = set(overlap)
    id_overlap, id_seen = [], set()                # each test row counted once
    for c in ids:
        a, b = train.columns.index(c), tcols.index(c)
        s = {r[a] for r in train.rows}
        for i, r in enumerate(test.rows):
            if r[b] in s and i not in overlap_set and i not in id_seen:
                id_seen.add(i)
                id_overlap.append(i)
    res["overlap"], res["idOverlap"] = len(overlap), len(id_overlap)
    if overlap:
        issues.add(type="train_test_overlap", column=None, severity="high", affected=len(overlap) / n2,
                   title=plural(len(overlap), "test row") + " also appear in the training file",
                   detail=fmt_pct(len(overlap) / n2) + " of the test set is a copy of training data" + (" (same features, ignoring IDs and the label)" if target is not None else "") + ". Test scores on these rows measure memory, not skill, so reported accuracy will be too high.",
                   evidence={"rows": row_nums(overlap)},
                   fix="key = " + py_list(key_cols[:6]) + ("  # ...all shared feature columns" if len(key_cols) > 6 else "") + "\ntest = test.merge(train[key].drop_duplicates(), on=key, how=\"left\", indicator=True)\ntest = test[test[\"_merge\"] == \"left_only\"].drop(columns=\"_merge\")")
    if id_overlap:
        issues.add(type="train_test_overlap", column=ids[0], severity="high", affected=len(id_overlap) / n2,
                   title=plural(len(id_overlap), "test row") + " share an ID with training rows",
                   detail="The same " + "/".join(ids) + " appears in both files with different values. If these are the same entities over time, split by ID so nobody is in both.",
                   evidence={"rows": row_nums(id_overlap)},
                   fix="test = test[~test[" + py_str(ids[0]) + "].isin(train[" + py_str(ids[0]) + "])]")

    # per-column shift
    for c in shared:
        p = prof[train.columns.index(c)]
        if p.is_id or p.type == "empty":
            continue
        tj = tcols.index(c)
        q = profile_column(c, [r[tj] if r[tj] is not None else "" for r in test.rows], n2)
        row = {"name": c}
        if p.type == "numeric" and q.sorted and len(q.sorted) > 5 and len(p.sorted) > 5:
            a = [x for x in p.sorted if x not in p.sentinels]
            b = [x for x in q.sorted if x not in p.sentinels]
            d, pv = ks2(a, b)
            row.update(stat=d, p=pv, kind="KS", trainMedian=quantile_sorted(a, 0.5), testMedian=quantile_sorted(b, 0.5))
            if d >= 0.1 and pv < 0.001:
                issues.add(type="distribution_shift", column=c, severity="high" if d >= 0.25 else "medium", affected=d,
                           title=c + " is distributed differently in the test file",
                           detail="Median " + fmt_num(row["trainMedian"]) + " in training vs " + fmt_num(row["testMedian"]) + " in test (KS distance " + to_fixed(d, 2) + ", p " + _p_text(pv) + "). A model tuned on the training range may do worse on the test data.",
                           evidence={"values": ["train median " + fmt_num(row["trainMedian"]), "test median " + fmt_num(row["testMedian"])]},
                           fix="# compare: train[" + py_str(c) + "].describe(), test[" + py_str(c) + "].describe()")
        elif p.type in ("categorical", "binary"):
            mp = p.mapping or {}
            ca, cb, na, nb = {}, {}, 0, 0
            for i in p.present_idx:
                v = mp.get(p.raw[i]) or p.raw[i]
                ca[v] = ca.get(v, 0) + 1
                na += 1
            for i in q.present_idx:
                v = mp.get(q.raw[i]) or q.raw[i]
                cb[v] = cb.get(v, 0) + 1
                nb += 1
            keys = js_keys(ca) + [k for k in js_keys(cb) if not ca.get(k)]
            tvd = chi = 0
            df = 0
            for k in keys:
                xa, xb = ca.get(k, 0), cb.get(k, 0)
                pa, pb = div(xa, na), div(xb, nb)
                tvd += abs(pa - pb) / 2
                tot = xa + xb
                ea, eb = tot * na / (na + nb), tot * nb / (na + nb)
                if ea > 0:
                    chi += (xa - ea) ** 2 / ea
                if eb > 0:
                    chi += (xb - eb) ** 2 / eb
                df += 1
            pv = chi_square_p(chi, max(1, df - 1))
            row.update(stat=tvd, p=pv, kind="TVD")
            unseen = [k for k in js_keys(cb) if not ca.get(k)]
            unseen_n = 0
            for k in unseen:
                unseen_n += cb[k]
            if unseen_n >= max(2, 0.01 * nb):
                issues.add(type="unseen_categories", column=c, severity="high" if unseen_n > 0.05 * nb else "medium", affected=unseen_n / nb,
                           title=c + " has values in test that training never saw",
                           detail=", ".join('"' + k + '" ×' + str(cb[k]) for k in unseen[:5]) + " (" + fmt_pct(unseen_n / nb) + " of test rows). A model has no idea what to do with them; encoders may even crash.",
                           evidence={"values": unseen[:8]},
                           fix="# map unknown categories explicitly, e.g. OneHotEncoder(handle_unknown=\"ignore\")")
            elif tvd >= 0.1 and pv < 0.001:
                issues.add(type="distribution_shift", column=c, severity="high" if tvd >= 0.25 else "medium", affected=tvd,
                           title=c + " has a different mix of values in the test file",
                           detail=_biggest_change(keys, ca, cb, na, nb) + "Total variation distance " + to_fixed(tvd, 2) + " (p " + _p_text(pv) + ")."
                           + (" The label mix itself differs, so accuracy measured on the test file won't match training." if target is not None and c == target.name else ""),
                           fix="# compare: train[" + py_str(c) + "].value_counts(normalize=True), test[" + py_str(c) + "].value_counts(normalize=True)")
        res["columns"].append(row)
    return res