#!/usr/bin/env python3
"""Iterative projected gradient ascent on log p(ADMIN), black-box via /batch.

Objective f(x) = log p(ADMIN)  (penalises drifting into INVALID or EMPLOYEE).
Each iteration: full forward-difference gradient (1024 queries) + line search.
"""
from __future__ import annotations
import sys, time, numpy as np
sys.path.insert(0, ".")
import attack

TARGET = 0.995


def obj(res: dict) -> float:
    return float(np.log(max(res["probabilities"]["ADMIN"], 1e-300)))


def fd_gradient(x: np.ndarray, f0: float, d: float = 0.15) -> np.ndarray:
    N = 32 * 32
    grids, deltas = [], np.zeros(N)
    for i in range(N):
        r, c = divmod(i, 32)
        step = d if x[r, c] + d <= 1.0 else -d
        deltas[i] = step
        g = x.copy(); g[r, c] += step
        grids.append(g.astype(np.float32))
    res = attack.eval_grids(grids)
    grad = np.zeros((32, 32))
    for i, rr in enumerate(res):
        r, c = divmod(i, 32)
        grad[r, c] = (obj(rr) - f0) / deltas[i]
    return grad


def main():
    try:
        x = np.load("best.npy").astype(np.float64)
        print("resuming from best.npy")
    except OSError:
        x = np.load("x0.npy").astype(np.float64)
    base = attack.eval_grids([x.astype(np.float32)])[0]
    f0 = obj(base)
    best_admin = base["probabilities"]["ADMIN"]
    print(f"start ADMIN={best_admin:.6f} f0={f0:.3f}")

    for it in range(60):
        g = fd_gradient(x, f0)
        gn = g / (np.abs(g).max() + 1e-12)
        alphas = [0.03, 0.06, 0.12, 0.2, 0.3]
        cands = [np.clip(x + a * gn, 0, 1).astype(np.float32) for a in alphas]
        res = attack.eval_grids(cands)
        scores = [obj(r) for r in res]
        j = int(np.argmax(scores))
        admins = [r["probabilities"]["ADMIN"] for r in res]
        print(f"it{it:02d} f0={f0:.3f} | " +
              " ".join(f"a{a}:{r['probabilities']['ADMIN']:.4f}"
                       for a, r in zip(alphas, res)) +
              f" -> pick a{alphas[j]} ADMIN={admins[j]:.6f} acc={res[j]['access']}")
        if res[j].get("flag"):
            print("FLAG:", res[j]["flag"]); np.save("best.npy", cands[j]); return
        if scores[j] <= f0 + 1e-4:
            # no improvement: shrink by trying tiny steps / random restart perturb
            print("  no improvement; trying smaller steps")
            alphas2 = [0.005, 0.01, 0.02]
            cands2 = [np.clip(x + a * gn, 0, 1).astype(np.float32) for a in alphas2]
            res2 = attack.eval_grids(cands2)
            sc2 = [obj(r) for r in res2]
            j2 = int(np.argmax(sc2))
            if sc2[j2] > f0 + 1e-6:
                x = cands2[j2].astype(np.float64); f0 = sc2[j2]
                best_admin = res2[j2]["probabilities"]["ADMIN"]
            else:
                print("  stuck, stop"); break
        else:
            x = cands[j].astype(np.float64); f0 = scores[j]
            best_admin = admins[j]
        np.save("best.npy", x.astype(np.float32))
        if best_admin >= TARGET:
            r = attack.eval_grids([x.astype(np.float32)])[0]
            print("REACHED", r["access"], "flag=", r.get("flag"))
            if r.get("flag"):
                print("FLAG:", r["flag"])
            return
    print("done, best ADMIN=", best_admin)


if __name__ == "__main__":
    main()