#!/usr/bin/env python3
"""Fresh native/full-label and exact protection runs for FzP6XZGG4d.

The native policy audit exhausts all 2^16 collaborative sets for selected
ImageNet-16H/VGG19 cells and compares the exact finite optimum with every
two-threshold policy induced by the 16 posterior scores.  The other two
audits add disjoint exact conformal-rank and CUP-Online regimes.  No stored
result artifact is read as evidence.
"""
from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from fractions import Fraction
from pathlib import Path

import numpy as np
import pandas as pd


def load_reproduce(root: Path):
    spec = importlib.util.spec_from_file_location("humanai_reproduce", root / "reproduce.py")
    if spec is None or spec.loader is None:
        raise RuntimeError("cannot load reproduce.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def load_helpers(root: Path):
    sys.path.insert(0, str(root / "code"))
    from offline_helpers import AIModel, HumanExpert, Imagenet16HPaths, get_common_image_names
    return AIModel, HumanExpert, Imagenet16HPaths, get_common_image_names


def native_exhaustive(root: Path, data_root: Path) -> dict:
    AIModel, HumanExpert, Imagenet16HPaths, get_common_image_names = load_helpers(root)
    paths = Imagenet16HPaths.from_data_root(data_root)
    true = pd.read_csv(paths.true_labels_csv, index_col="image_name")
    masks = np.arange(1 << 16, dtype=np.uint32)
    bit_weights = (1 << np.arange(16, dtype=np.uint32))
    popcount = np.zeros(1 << 16, dtype=np.int16)
    for bit in range(16):
        popcount += ((masks >> bit) & 1).astype(np.int16)
    rows = []
    for noise in (80, 95, 110, 125):
        ai = AIModel(paths, noise_level=noise, model_name="vgg19")
        human = HumanExpert(paths, noise_level=noise)
        images = get_common_image_names(ai, human, true)
        # Evenly spaced cells cover all noise conditions without selecting a
        # convenient contiguous image block.
        chosen = [images[int(i)] for i in np.linspace(0, len(images) - 1, 8)]
        for image in chosen:
            p = ai.get_prob(image).to_numpy(dtype=float)
            human_set = set(human.predict_set(image, strategy="topk2"))
            inside = np.array([label in human_set for label in ai.labels], dtype=bool)
            in_mass = float(p[inside].sum())
            out_mass = float(p[~inside].sum())
            selected = ((masks[:, None] & bit_weights[None, :]) != 0)
            harm = ((selected[:, inside] == 0) * p[inside]).sum(axis=1) / in_mass
            complementarity = (selected[:, ~inside] * p[~inside]).sum(axis=1) / out_mass
            feasible = (harm < 0.2 - 1e-12) & (complementarity >= 0.75 - 1e-12)
            if not feasible.any():
                raise AssertionError(f"no feasible native policy for {noise}/{image}")
            exact_min = int(popcount[feasible].min())
            unique_in = sorted(set(p[inside]), reverse=True)
            unique_out = sorted(set(p[~inside]), reverse=True)
            threshold_masks = []
            for q_in in [None, *unique_in]:
                for q_out in [None, *unique_out]:
                    keep = np.array([
                        ((q_in is not None and value >= q_in) if inside[idx]
                         else (q_out is not None and value >= q_out))
                        for idx, value in enumerate(p)
                    ])
                    threshold_masks.append(int(np.dot(keep.astype(np.uint32), bit_weights)))
            threshold_masks = np.array(sorted(set(threshold_masks)), dtype=np.uint32)
            threshold_ok = feasible[threshold_masks]
            if not threshold_ok.any():
                raise AssertionError(f"no feasible threshold policy for {noise}/{image}")
            threshold_min = int(popcount[threshold_masks[threshold_ok]].min())
            chosen_index = int(np.flatnonzero(feasible & (popcount == exact_min))[0])
            rows.append({
                "noise": noise,
                "image": image,
                "labels": 16,
                "policies_exhausted": int(1 << 16),
                "feasible_policies": int(feasible.sum()),
                "exact_min_size": exact_min,
                "two_threshold_min_size": threshold_min,
                "gap": threshold_min - exact_min,
                "harm_at_best": float(harm[chosen_index]),
                "complementarity_at_best": float(complementarity[chosen_index]),
            })
    return {
        "cells": rows,
        "cell_count": len(rows),
        "policies_exhausted": sum(row["policies_exhausted"] for row in rows),
        "zero_gap_cells": sum(row["gap"] == 0 for row in rows),
        "max_gap": max(row["gap"] for row in rows),
        "max_feasible_policy_count": max(row["feasible_policies"] for row in rows),
    }


def rank_protection() -> dict:
    rows = []
    for group, alpha in (("human_in_epsilon", Fraction(1, 5)), ("human_out_delta", Fraction(1, 4))):
        for n in (3, 5, 11, 23, 47, 127, 255, 511):
            target = (1 - alpha) * (n + 1)
            k = min(n + 1, target.numerator // target.denominator + int(target.numerator % target.denominator != 0))
            coverage = Fraction(k, n + 1)
            lower = 1 - alpha
            upper = lower + Fraction(1, n + 1)
            rows.append({
                "group": group,
                "n": n,
                "rank": k,
                "coverage": f"{coverage.numerator}/{coverage.denominator}",
                "lower_holds": coverage >= lower,
                "strict_upper_holds": coverage < upper,
            })
    return {"rows": rows, "cells": len(rows), "all_bounds_hold": all(r["lower_holds"] and r["strict_upper_holds"] for r in rows)}


def online_protection(reproduce) -> dict:
    rows = []
    for eta, epsilon, delta in ((Fraction(1, 7), Fraction(1, 8), Fraction(1, 5)),
                                (Fraction(1, 9), Fraction(1, 6), Fraction(1, 4))):
        for length in (600, 1800, 6000):
            for kind in ("abrupt", "alternating", "chirp", "human_adaptation"):
                row = reproduce.run_online(kind, length, eta, epsilon, delta)
                row.update({"eta": str(eta), "epsilon": str(epsilon), "delta": str(delta)})
                rows.append(row)
    return {
        "rows": rows,
        "cells": len(rows),
        "all_bounds_hold": all(r["in_bound_holds"] and r["out_bound_holds"] and
                                r["telescope_in_exact"] and r["telescope_out_exact"] for r in rows),
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", type=Path, required=True)
    ap.add_argument("--data-root", type=Path, required=True)
    ap.add_argument("--output", type=Path, required=True)
    args = ap.parse_args()
    reproduce = load_reproduce(args.repo)
    result = {
        "native_exhaustive": native_exhaustive(args.repo, args.data_root),
        "exact_rank_protection": rank_protection(),
        "online_protection": online_protection(reproduce),
    }
    if result["native_exhaustive"]["zero_gap_cells"] != result["native_exhaustive"]["cell_count"]:
        raise AssertionError("a native exhaustive cell did not have a two-threshold optimum")
    if not result["exact_rank_protection"]["all_bounds_hold"] or not result["online_protection"]["all_bounds_hold"]:
        raise AssertionError("protection gate failed")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({
        "native_cells": result["native_exhaustive"]["cell_count"],
        "native_policies": result["native_exhaustive"]["policies_exhausted"],
        "native_zero_gap": result["native_exhaustive"]["zero_gap_cells"],
        "rank_cells": result["exact_rank_protection"]["cells"],
        "online_cells": result["online_protection"]["cells"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()