"""E8: adaptive representation policy fitted on development records only.

Features (deployment-observable): context tokens (receiver tokenizer), question-access regime, task family, receiver model,
reuse count N, measured LAN bandwidth (constant from E6), and the candidate action's development cost model.
Action = condition id among a feasible set. Quality constraint = lower 95% bootstrap bound of dev mean quality for the (family,
context-bin, pair) group >= tau * best static lower bound; cost = dev mean amortized latency (or bytes). Policy = argmin cost subject
to constraint, with fallback to the dev-best static action if no action satisfies the constraint. Frozen before test evaluation.
"""
from __future__ import annotations
import json, yaml
import numpy as np, pandas as pd
from pathlib import Path
from .util import ROOT, write_json, read_json
from .analysis.load import load_experiments
from .analysis.stats import cluster_bootstrap_mean

def _ctx_bin(tokens):
    return "short" if tokens <= 3000 else ("medium" if tokens <= 10000 else "long")

def fit(dev_df: pd.DataFrame, actions: list[str], cost_col: str, tau: float, n_boot: int = 500) -> dict:
    policy = {"actions": actions, "cost_col": cost_col, "tau": tau, "table": {}}
    dev_df = dev_df[dev_df["condition"].isin(actions)].copy()
    dev_df["ctx_bin"] = dev_df["ctx_tokens"].apply(_ctx_bin)
    for (fam, cb, pair, acc), g in dev_df.groupby(["task_family", "ctx_bin", "pair", "question_access_group"]):
        rows = []
        for a, ga in g.groupby("condition"):
            q = cluster_bootstrap_mean(ga, "quality_operational", n_boot=n_boot); c = ga[cost_col].mean()
            rows.append({"action": a, "q_mean": q["mean"], "q_lo": q["lo"], "cost": float(c), "n": q["n"]})
        if not rows: continue
        best = max(rows, key=lambda r: r["q_lo"])
        feasible = [r for r in rows if r["q_lo"] >= tau * best["q_lo"]]
        chosen = min(feasible, key=lambda r: r["cost"]) if feasible else best
        policy["table"][f"{fam}|{cb}|{pair}|{acc}"] = {"chosen": chosen["action"], "best_static": best["action"], "candidates": rows}
    return policy

def choose(policy: dict, fam: str, ctx_tokens: int, pair: str, acc: str) -> tuple[str, bool]:
    key = f"{fam}|{_ctx_bin(ctx_tokens)}|{pair}|{acc}"
    if key in policy["table"]:
        return policy["table"][key]["chosen"], False
    # fallback: nearest context bin for the same family/pair, else global best static
    for k, v in policy["table"].items():
        if k.startswith(f"{fam}|") and k.endswith(f"|{pair}|{acc}"):
            return v["chosen"], True
    return policy.get("global_best_static", "A1"), True

def main(config_path: str) -> int:
    cfg = yaml.safe_load(open(config_path))
    df = load_experiments(cfg["dev_patterns"])
    df = df[df["split"] == "dev"]
    df["ctx_tokens"] = df["message_utf8_bytes"].fillna(0) / 4.0  # proxy when receiver source tokens absent
    if "source_tokens_receiver" in df: df["ctx_tokens"] = df["source_tokens_receiver"].fillna(df["ctx_tokens"])
    df["question_access_group"] = df["question_access"].replace({"none": "blind"})
    pol = fit(df, cfg["actions"], cfg["cost_col"], cfg["tau"])
    g = df[df["condition"].isin(cfg["actions"])].groupby("condition")["quality_operational"].mean()
    pol["global_best_static"] = str(g.idxmax()); pol["fitted_on"] = cfg["dev_patterns"]; pol["n_dev_rows"] = int(len(df))
    # dev-best static action per access regime (aware-eligible actions for the aware regime)
    pol["global_best_static_by_access"] = {}
    for acc, ga in df[df["condition"].isin(cfg["actions"])].groupby("question_access_group"):
        gg = ga.groupby("condition")["quality_operational"].mean(); pol["global_best_static_by_access"][acc] = str(gg.idxmax())
    write_json(ROOT / "runs/checkpoints/router_policy.json", pol)
    print("router fitted:", len(pol["table"]), "cells; global best static:", pol["global_best_static"])
    return 0