#!/usr/bin/env python3
"""Solve the exact linearized joint group-selection problem."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import hashlib
import itertools
import json
from pathlib import Path
import sys

import numpy as np
from scipy.optimize import Bounds, LinearConstraint, milp
from scipy.sparse import coo_matrix


ROOT = Path(__file__).resolve().parents[1]
CONFIG = ROOT / "configs/experiments/strata_headquotient_v1_1.json"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("frontier", choices=("Q10", "Q20", "Q25", "Q30"))
    parser.add_argument("--config", type=Path, default=CONFIG)
    parser.add_argument("--time-limit", type=float, default=900)
    parser.add_argument("--relative-gap", type=float, default=0.30)
    parser.add_argument(
        "--allow-direct-replication",
        action="store_true",
        help="Allow a fresh replication to solve Q25 without rerunning Q10/Q20.",
    )
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(16 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def group(row: dict) -> tuple[int, int]:
    return int(row["layer"]), int(row["head"])


def fit_interaction_scale(
    nodes: dict[tuple[int, int], dict],
    pairs: dict[tuple[tuple[int, int], tuple[int, int]], dict],
    triples: list[dict],
    metric: str,
) -> dict[str, float]:
    pair_sum = []
    residual = []
    for triple in triples:
        members = [group(row) for row in triple["groups"]]
        node_cost = sum(float(nodes[item]["lm_delta"][metric]) for item in members)
        interaction = sum(
            float(pairs[tuple(sorted(edge))]["lm_interaction"][metric])
            for edge in itertools.combinations(members, 2)
        )
        pair_sum.append(interaction)
        residual.append(float(triple["lm_delta"][metric]) - node_cost)
    x = np.asarray(pair_sum)
    y = np.asarray(residual)
    denominator = max(float(x @ x), 1e-12)
    scale = float((x @ y) / denominator)
    prediction = scale * x
    correlation = float(np.corrcoef(prediction, y)[0, 1]) if np.std(prediction) else 0.0
    return {
        "scale": scale,
        "correlation": correlation,
        "rmse": float(np.sqrt(np.mean((prediction - y) ** 2))),
    }


def main() -> None:
    args = parse_args()
    config = json.loads(args.config.read_text(encoding="utf-8"))
    output = Path(config["output_root"])
    interactions = json.loads((output / "interactions/result.json").read_text(encoding="utf-8"))
    if not interactions["passed"]:
        raise RuntimeError("joint selection is closed because the interaction graph is incomplete")
    frontier_index = config["frontier"].index(args.frontier)
    direct_replication = False
    if frontier_index > 0:
        prior_name = config["frontier"][frontier_index - 1]
        prior = output / f"frontier/{prior_name}/result.json"
        direct_replication = bool(
            args.allow_direct_replication
            and args.frontier == "Q25"
            and config["scoped_training"].get("fresh_initialization", False)
            and "REP" in config["program"]
        )
        if (
            not direct_replication
            and (
                not prior.is_file()
                or not json.loads(prior.read_text(encoding="utf-8"))["quality_pass"]
            )
        ):
            raise RuntimeError(f"{args.frontier} is closed until {prior_name} passes")

    node_payload = json.loads(Path(interactions["artifacts"]["nodes"]["path"]).read_text(encoding="utf-8"))
    node_rows = node_payload["nodes"]
    candidates = [group(row) for row in node_rows]
    nodes = {group(row): row for row in node_rows}
    pair_rows = [
        json.loads(line)
        for line in Path(interactions["artifacts"]["pairs"]["path"]).read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    pairs = {
        tuple(sorted((group(row["groups"][0]), group(row["groups"][1])))): row
        for row in pair_rows
    }
    expected_pairs = len(candidates) * (len(candidates) - 1) // 2
    if len(pairs) != expected_pairs:
        raise RuntimeError(f"interaction graph has {len(pairs)} unique pairs, expected {expected_pairs}")
    triples = json.loads(Path(interactions["artifacts"]["triples"]["path"]).read_text(encoding="utf-8"))
    aggregate_scale = fit_interaction_scale(nodes, pairs, triples, "aggregate")
    late_scale = fit_interaction_scale(nodes, pairs, triples, "4096-8192")

    budget = int(config["budgets"][args.frontier])
    if budget > len(candidates):
        raise RuntimeError(f"{args.frontier} budget exceeds the frozen candidate set")
    n = len(candidates)
    pair_keys = list(itertools.combinations(range(n), 2))
    pair_offset = n
    layer_offset = n + len(pair_keys)
    layers = sorted({layer for layer, _head in candidates})
    layer_index = {layer: layer_offset + index for index, layer in enumerate(layers)}
    variables = layer_offset + len(layers)
    objective = np.zeros(variables)
    semantic_weight = 20.0
    for index, item in enumerate(candidates):
        row = nodes[item]
        objective[index] = (
            float(row["lm_delta"]["aggregate"])
            + 0.5 * float(row["lm_delta"]["4096-8192"])
            + semantic_weight * max(0.0, float(row["semantic_cost"]))
        )
    for pair_position, (first, second) in enumerate(pair_keys):
        row = pairs[tuple(sorted((candidates[first], candidates[second])))]
        objective[pair_offset + pair_position] = (
            aggregate_scale["scale"] * float(row["lm_interaction"]["aggregate"])
            + 0.5 * late_scale["scale"] * float(row["lm_interaction"]["4096-8192"])
            + semantic_weight * max(0.0, float(row["semantic_interaction"]))
        )
    for layer in layers:
        objective[layer_index[layer]] = 5e-3

    row_indices: list[int] = []
    column_indices: list[int] = []
    values: list[float] = []
    lower: list[float] = []
    upper: list[float] = []

    def constraint(coefficients: dict[int, float], lb: float, ub: float) -> None:
        row = len(lower)
        for column, value in coefficients.items():
            row_indices.append(row)
            column_indices.append(column)
            values.append(value)
        lower.append(lb)
        upper.append(ub)

    constraint({index: 1.0 for index in range(n)}, budget, budget)
    for pair_position, (first, second) in enumerate(pair_keys):
        y = pair_offset + pair_position
        constraint({y: 1, first: -1}, -np.inf, 0)
        constraint({y: 1, second: -1}, -np.inf, 0)
        constraint({first: 1, second: 1, y: -1}, -np.inf, 1)
    for index, (layer, _head) in enumerate(candidates):
        constraint({index: 1, layer_index[layer]: -1}, -np.inf, 0)
    for layer in layers:
        members = [
            index for index, item in enumerate(candidates) if item[0] == layer
        ]
        constraint(
            {**{index: 1 for index in members}, layer_index[layer]: -2},
            0,
            np.inf,
        )
    maximum_layers = {"Q10": 5, "Q20": 8, "Q25": 10, "Q30": 15}[args.frontier]
    constraint({layer_index[layer]: 1 for layer in layers}, -np.inf, maximum_layers)

    graph_layer = 23
    graph_indices = [index for index, item in enumerate(candidates) if item[0] == graph_layer]
    constraint({index: 1 for index in graph_indices}, min(10, budget), np.inf)
    for program in range(5):
        eligible = [
            index for index in graph_indices
            if program in [
                int(value) for value in nodes[candidates[index]]["frozen_h1"].get("programs", [])
            ]
        ]
        # The v1 H1 rows predate scoped programs; recover assignments from the
        # frozen scoped node field when no explicit program list exists.
        if not eligible:
            eligible = [
                index for index in graph_indices
                if nodes[candidates[index]]["program"] == (
                    "READ_EVENT_ARG0", "READ_EVENT_ARG1", "READ_EVENT_ARG2",
                    "READ_EVENT_TIME", "READ_EVENT_LOCATION",
                )[program]
                or program == 0
            ]
        if len(eligible) < 2:
            raise RuntimeError(f"program {program} has fewer than two candidate readers")
        constraint({index: 1 for index in eligible}, 2, np.inf)

    parent_similarity = json.loads((
        Path(config["parent_manifest"]).parent
        / "h1/qualification/functional_similarity_graph.json"
    ).read_text(encoding="utf-8"))
    total_by_cluster: dict[int, int] = {}
    cluster_by_group = {}
    for row in parent_similarity:
        cluster = int(row["cluster"])
        total_by_cluster[cluster] = total_by_cluster.get(cluster, 0) + 1
        cluster_by_group[group(row)] = cluster
    for cluster, total in total_by_cluster.items():
        member_indices = [
            index for index, item in enumerate(candidates)
            if cluster_by_group[item] == cluster
        ]
        if member_indices:
            constraint({index: 1 for index in member_indices}, -np.inf, total - 1)

    matrix = coo_matrix(
        (values, (row_indices, column_indices)), shape=(len(lower), variables),
    ).tocsr()
    result = milp(
        c=objective,
        integrality=np.ones(variables),
        bounds=Bounds(np.zeros(variables), np.ones(variables)),
        constraints=LinearConstraint(matrix, np.asarray(lower), np.asarray(upper)),
        options={
            "time_limit": float(args.time_limit),
            "mip_rel_gap": float(args.relative_gap),
            "presolve": True,
            "disp": True,
        },
    )
    if result.x is None or result.status not in (0, 1):
        raise RuntimeError(
            f"branch-and-bound produced no feasible selection: "
            f"status={result.status}, message={result.message}"
        )
    selected = [candidates[index] for index in range(n) if result.x[index] > 0.5]
    if len(selected) != budget:
        raise AssertionError("solver returned the wrong replacement cardinality")
    selected_layers = sorted({layer for layer, _head in selected})
    groups = []
    for item in selected:
        row = nodes[item]
        groups.append({
            "layer": item[0],
            "head": item[1],
            "mode": "LOCAL_GRAPH" if item[0] == graph_layer else "LOCAL",
            "program": row["program"] if item[0] == graph_layer else "NO_GRAPH_READ",
            "node_objective": float(objective[candidates.index(item)]),
            "cluster": cluster_by_group[item],
        })
    plan = {
        "program": config["program"],
        "stage": f"{args.frontier}_exact_joint_selection",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "frontier": args.frontier,
        "budget": budget,
        "candidate_groups": n,
        "groups": groups,
        "selected_layers": selected_layers,
        "selected_layer_count": len(selected_layers),
        "graph_groups": sum(row["mode"] == "LOCAL_GRAPH" for row in groups),
        "local_groups": sum(row["mode"] == "LOCAL" for row in groups),
        "replacement_fraction": budget / 384,
        "global_kv_reduction": 384 / (384 - budget),
        "solver": {
            "method": "scipy.optimize.milp branch-and-bound over linearized signed QUBO",
            "status": int(result.status),
            "message": result.message,
            "objective": float(result.fun),
            "mip_gap": float(result.mip_gap),
            "mip_node_count": int(result.mip_node_count),
            "requested_relative_gap": float(args.relative_gap),
            "globally_optimal": bool(float(result.mip_gap) <= 1e-9),
            "feasible": True,
        },
        "triple_calibration": {
            "aggregate": aggregate_scale,
            "late": late_scale,
        },
        "constraints": {
            "exact_budget": budget,
            "minimum_layer23_graph_groups": min(10, budget),
            "minimum_readers_per_program": 2,
            "retained_representative_per_functional_cluster": True,
            "compact_layer_penalty": 5e-3,
            "maximum_localized_layers": maximum_layers,
            "minimum_groups_per_localized_layer": 2,
            "direct_replication_without_lower_frontiers": bool(direct_replication),
        },
        "interaction_artifacts": interactions["artifacts"],
    }
    result_root = output / f"selection/{args.frontier}"
    result_root.mkdir(parents=True, exist_ok=True)
    destination = result_root / "plan.json"
    destination.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    manifest = {
        "plan": {"path": str(destination), "sha256": sha256(destination)},
        "optimal": plan["solver"]["globally_optimal"],
        "certified_relative_gap": plan["solver"]["mip_gap"],
        "feasible": True,
        "next_stage": f"EXPORT_{args.frontier}",
    }
    (result_root / "result.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8",
    )
    print(json.dumps({
        "output": str(destination),
        "objective": result.fun,
        "selected_layers": selected_layers,
        "graph_groups": plan["graph_groups"],
        "global_kv_reduction": plan["global_kv_reduction"],
    }, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()