"""The S-ablation of A' on all six corpora (uniform criterion, a sign change within 1e-6):
A' with each of the nine similarity functions in the step, everything else unchanged.  The
tent row is A' itself.  Descriptive: every corpus has been seen.

    python -u sweep_ablation_s.py            run -> sweep_ablation_s.json, then the table
    python -u sweep_ablation_s.py --report
"""
from __future__ import annotations

import json
import sys
from multiprocessing import Pool

import sweep_fourth
from sweep_fourth import evaluate as _eval_one
from sweep_full import APRIME, PURE_SIMS

SETS = ("APS", "CHA", "NEW", "RJL", "APP", "D")

for _s in PURE_SIMS:
    sweep_fourth.CONFIGS[f"S:{_s}"] = ("pure", _s, "derived3", "project", "x", APRIME,
                                      "recent", "distance", "none", "pattern")


def evaluate(job):
    name, which = job
    return name, _eval_one((name, which))[1]


def cost(v):
    return v[0] if v[1] else 200


def main():
    names = [f"S:{s}" for s in sorted(PURE_SIMS)]
    if "--report" not in sys.argv:
        jobs = [(n, w) for n in names for w in ("seen", "fourth")]
        with Pool(min(12, len(jobs))) as pool:
            out = pool.map(evaluate, jobs)
        R = {}
        for n, res in out:
            R.setdefault(n, {}).update(res)
        with open("sweep_ablation_s.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    with open("sweep_ablation_s.json", encoding="utf8") as fh:
        R = json.load(fh)
    tot = {n: {s: sum(cost(v) for v in R[n][s].values()) for s in SETS} for n in R}
    viol = {n: sum(v[2] for s in SETS for v in R[n][s].values()) for n in R}
    print("A' with S in the step - totals per corpus (violations of the 2x bound)")
    for n in sorted(tot, key=lambda k: sum(tot[k].values())):
        print(f"  {n[2:]:<11} " + "  ".join(f"{s} {tot[n][s]:5d}" for s in SETS)
              + f"   all {sum(tot[n].values()):6d}  (viol {viol[n]})")
    print("ablation S done")


if __name__ == "__main__":
    main()