"""Convergence diagnostics for the in-silico code optimization.

The manuscript reports that codes optimized for strong transition bias do not
converge on the standard code, at a correlation of -0.76 against log kappa. A
heuristic search cannot rule out that a larger budget would find better optima,
so this records what the budget actually achieved: the distribution of best
costs across restarts, how often independent restarts reach the same cost, and
whether raising the budget moves either the cost or the similarity.

Writes data/optim_convergence.json and figures/FigS2.*.
"""
import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
DATA = os.path.join(ROOT, "data")
FIGS = os.path.join(ROOT, "figures")
sys.path.insert(0, HERE)

from optimize_code import (anneal, block_pair_weights, cost,  # noqa: E402
                           kappa_spectrum, similarity)

RED = "#a4342a"
DARK = "#2b2b2b"
GREY = "#7a7a7a"
plt.rcParams.update({"font.size": 8, "axes.labelsize": 8,
                     "xtick.labelsize": 7, "ytick.labelsize": 7,
                     "axes.spines.top": False, "axes.spines.right": False,
                     "figure.dpi": 150})


def _setup(kappa):
    from codecore import AA20, POLAR_REQUIREMENT, scalar_metric_matrix
    B = block_pair_weights(kappa_spectrum(kappa))
    d = scalar_metric_matrix(POLAR_REQUIREMENT, True)
    D = np.zeros((21, 21))
    for i, a in enumerate(AA20):
        for j, b in enumerate(AA20):
            D[i, j] = d[(a, b)]
    return B, D


def main():
    out = {}
    kappas = [1.0, 2.0, 4.0]
    budgets = [(10, 30000), (40, 30000), (40, 120000)]
    print("\n1. BEST COST AND SIMILARITY BY BUDGET\n")
    print(f"   {'kappa':>6} {'restarts':>9} {'iters':>8} "
          f"{'best cost':>11} {'sim at best':>12} {'reach best':>11}")
    rows = []
    for kappa in kappas:
        B, D = _setup(kappa)
        for nres, iters in budgets:
            rng = np.random.default_rng(20260727 + int(kappa * 10) + nres)
            costs, sims = [], []
            for _ in range(nres):
                a, c_best = anneal(B, D, rng, iters=iters)
                costs.append(float(c_best))
                sims.append(similarity(a))
            costs = np.array(costs); sims = np.array(sims)
            best = costs.min()
            # how many restarts land within 0.1% of the best cost found
            reach = int((costs <= best * 1.001).sum())
            sim_at_best = float(sims[costs.argmin()])
            rows.append(dict(kappa=kappa, restarts=nres, iters=iters,
                             best_cost=float(best),
                             median_cost=float(np.median(costs)),
                             spread=float(costs.max() / best - 1),
                             sim_at_best=sim_at_best,
                             sim_median=float(np.median(sims)),
                             reach_best=reach))
            print(f"   {kappa:6.1f} {nres:9d} {iters:8d} {best:11.4f} "
                  f"{sim_at_best:12.2f} {reach:>7d}/{nres}")
    out["budget_sweep"] = rows

    print("\n2. DOES A LARGER BUDGET CHANGE THE ANSWER?\n")
    for kappa in kappas:
        small = [r for r in rows if r["kappa"] == kappa
                 and r["restarts"] == 10][0]
        large = [r for r in rows if r["kappa"] == kappa
                 and r["iters"] == 120000][0]
        dc = 100 * (small["best_cost"] / large["best_cost"] - 1)
        print(f"   kappa {kappa:.1f}: cost improves {dc:+.2f}% from the "
              f"smallest to the largest budget; similarity "
              f"{small['sim_at_best']:.2f} -> {large['sim_at_best']:.2f}")
        out[f"gain_kappa_{kappa:.0f}"] = round(float(dc), 3)

    fig, ax = plt.subplots(1, 2, figsize=(6.85, 2.4))
    for kappa, c in zip(kappas, (GREY, DARK, RED)):
        rs = [r for r in rows if r["kappa"] == kappa]
        ax[0].plot([r["restarts"] * r["iters"] / 1e6 for r in rs],
                   [r["best_cost"] for r in rs], "o-", color=c, ms=3.5,
                   lw=1.2, label=f"$\\kappa$ = {kappa:.0f}")
        ax[1].plot([r["restarts"] * r["iters"] / 1e6 for r in rs],
                   [100 * r["sim_at_best"] for r in rs], "o-", color=c,
                   ms=3.5, lw=1.2)
    ax[0].set_xlabel("total iterations (millions)")
    ax[0].set_ylabel("best cost found")
    ax[0].legend(frameon=False, fontsize=7)
    ax[1].axhline(5, color=GREY, lw=1, ls="--")
    ax[1].set_xlabel("total iterations (millions)")
    ax[1].set_ylabel("similarity to the standard code (%)")
    fig.tight_layout()
    for ext in ("eps", "pdf", "png"):
        fig.savefig(os.path.join(FIGS, f"FigS2.{ext}"), bbox_inches="tight")
    plt.close(fig)

    json.dump(out, open(os.path.join(DATA, "optim_convergence.json"), "w"),
              indent=1)
    print("\n   figures/FigS2.* and data/optim_convergence.json written\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
