#!/usr/bin/env python3
"""Development-only cross-config telemetry transfer diagnostic for P1.

Each example asks whether state observed at one source anchor helps predict the
pass rate at a different target config.  The hybrid branch predicts the
real-minus-simulator residual; the direct branch never reads simulator state or
outcomes.  Folds exclude both the source and target config identities.  The two
offered-load anchors belong to one trace/SLO task, so this is a premise check
rather than generalization evidence.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any, Sequence

import numpy as np


REGULARIZATION = (0.1, 1.0, 10.0, 100.0)
PRIOR_SHRINKAGE = (0.0, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0)


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    temporary.replace(path)


def finite(value: Any, *, name: str) -> float:
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"non-finite feature {name}={value!r}")
    return result


def flatten_residual_state(example: dict[str, Any]) -> tuple[float, ...]:
    residual = example["state_residual"]["values"]
    values = [finite(residual[name], name=name) for name in sorted(residual)]
    engine_only = example["engine"]["engine_only"]
    values.extend(
        finite(engine_only[name], name=name)
        for name in sorted(engine_only)
    )
    return tuple(values)


def flatten_engine_state(example: dict[str, Any]) -> tuple[float, ...]:
    values = []
    for name in sorted(example["engine"]["common"]):
        value = example["engine"]["common"][name]
        if isinstance(value, dict):
            values.extend(
                finite(value[statistic], name=f"{name}.{statistic}")
                for statistic in ("mean", "max", "cv")
            )
        elif value is not None:
            values.append(finite(value, name=name))
    engine_only = example["engine"]["engine_only"]
    values.extend(
        finite(engine_only[name], name=name)
        for name in sorted(engine_only)
    )
    return tuple(values)


def config_transition_features(
    source: dict[str, Any], target: dict[str, Any]
) -> tuple[float, ...]:
    source_rate = finite(
        source["offered_req_s_per_gpu"], name="source_offered_req_s_per_gpu"
    )
    target_rate = finite(
        target["offered_req_s_per_gpu"], name="target_offered_req_s_per_gpu"
    )
    return (
        math.log2(float(source["tp"])),
        math.log2(float(source["mns"])),
        math.log2(float(target["tp"])),
        math.log2(float(target["mns"])),
        math.log2(float(target["tp"]) / float(source["tp"])),
        math.log2(float(target["mns"]) / float(source["mns"])),
        math.log2(source_rate),
        math.log2(target_rate),
        math.log2(target_rate / source_rate),
    )


def hybrid_base_features(
    source: dict[str, Any], target: dict[str, Any]
) -> tuple[float, ...]:
    return (
        finite(source["pass_rate_residual"], name="source_pass_rate_residual"),
        finite(source["sim_pass_rate"], name="source_sim_pass_rate"),
        finite(target["sim_pass_rate"], name="target_sim_pass_rate"),
    ) + config_transition_features(source, target)


def direct_base_features(
    source: dict[str, Any], target: dict[str, Any]
) -> tuple[float, ...]:
    return (
        finite(source["real_pass_rate_rep1"], name="source_real_pass_rate"),
    ) + config_transition_features(source, target)


def transitions(examples: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    rows = []
    for source in examples:
        residual_state = flatten_residual_state(source)
        engine_state = flatten_engine_state(source)
        for target in examples:
            # Low/high are offered-load anchors inside one workload/SLO task.
            # Cross-load transitions are legal; same-cell transitions are
            # excluded because R0 asks about transfer to a new configuration.
            if source["cell"] == target["cell"]:
                continue
            hybrid_base = hybrid_base_features(source, target)
            direct_base = direct_base_features(source, target)
            rows.append(
                {
                    "source_cell": source["cell"],
                    "source_level": source["level"],
                    "target_cell": target["cell"],
                    "target_level": target["level"],
                    "hybrid_base": hybrid_base,
                    "hybrid_telemetry": hybrid_base + residual_state,
                    "direct_base": direct_base,
                    "direct_telemetry": direct_base + engine_state,
                    "target_residual": finite(
                        target["pass_rate_residual"], name="target_residual"
                    ),
                    "target_residual_delta": finite(
                        target["pass_rate_residual"], name="target_residual"
                    )
                    - finite(source["pass_rate_residual"], name="source_residual"),
                    "target_real_pass_rate": finite(
                        target["real_pass_rate_rep1"], name="target_real_pass_rate"
                    ),
                    "target_real_pass_rate_delta": finite(
                        target["real_pass_rate_rep1"], name="target_real_pass_rate"
                    )
                    - finite(
                        source["real_pass_rate_rep1"], name="source_real_pass_rate"
                    ),
                    "source_residual": finite(
                        source["pass_rate_residual"], name="source_residual"
                    ),
                    "source_real_pass_rate": finite(
                        source["real_pass_rate_rep1"], name="source_real_pass_rate"
                    ),
                    "target_sim_pass_rate": finite(
                        target["sim_pass_rate"], name="target_sim_pass_rate"
                    ),
                    "target_real_feasible": bool(target["real_feasible"]),
                    "target_sim_feasible": bool(target["sim_feasible"]),
                }
            )
    return rows


def fit_predict(
    train_x: np.ndarray,
    train_y: np.ndarray,
    test_x: np.ndarray,
    regularization: float,
) -> np.ndarray:
    mean = train_x.mean(axis=0)
    std = train_x.std(axis=0)
    std[std == 0.0] = 1.0
    train = (train_x - mean) / std
    test = (test_x - mean) / std
    train = np.column_stack((np.ones(len(train)), train))
    test = np.column_stack((np.ones(len(test)), test))
    penalty = np.eye(train.shape[1], dtype=np.float64)
    penalty[0, 0] = 0.0
    weights = np.linalg.lstsq(
        train.T @ train + regularization * penalty,
        train.T @ train_y,
        rcond=None,
    )[0]
    return test @ weights


def grouped_predictions(
    rows: Sequence[dict[str, Any]],
    *,
    feature_name: str,
    target_name: str,
    regularization: float,
) -> np.ndarray:
    predictions = np.zeros(len(rows), dtype=np.float64)
    groups = sorted({(row["source_cell"], row["target_cell"]) for row in rows})
    for source_cell, target_cell in groups:
        held_out_cells = {source_cell, target_cell}
        test_indexes = [
            index
            for index, row in enumerate(rows)
            if row["source_cell"] == source_cell and row["target_cell"] == target_cell
        ]
        train_indexes = [
            index
            for index, row in enumerate(rows)
            if row["source_cell"] not in held_out_cells
            and row["target_cell"] not in held_out_cells
        ]
        if not test_indexes or not train_indexes:
            raise ValueError(f"empty grouped fold for {source_cell}->{target_cell}")
        train_x = np.asarray([rows[index][feature_name] for index in train_indexes])
        train_y = np.asarray([rows[index][target_name] for index in train_indexes])
        test_x = np.asarray([rows[index][feature_name] for index in test_indexes])
        predictions[test_indexes] = fit_predict(
            train_x, train_y, test_x, regularization
        )
    return predictions


def metrics(rows: Sequence[dict[str, Any]], predicted_pass: np.ndarray) -> dict[str, Any]:
    truth = np.asarray(
        [row["target_real_pass_rate"] for row in rows], dtype=np.float64
    )
    real_feasible = np.asarray(
        [row["target_real_feasible"] for row in rows], dtype=bool
    )
    sim_feasible = np.asarray(
        [row["target_sim_feasible"] for row in rows], dtype=bool
    )
    clipped_pass = np.clip(predicted_pass, 0.0, 1.0)
    predicted_feasible = clipped_pass >= 0.95
    baseline_correct = sim_feasible == real_feasible
    model_correct = predicted_feasible == real_feasible
    return {
        "rmse": float(np.sqrt(np.mean((predicted_pass - truth) ** 2))),
        "mae": float(np.mean(np.abs(predicted_pass - truth))),
        "feasibility_accuracy": float(np.mean(model_correct)),
        "false_feasible": int(np.sum(predicted_feasible & ~real_feasible)),
        "false_infeasible": int(np.sum(~predicted_feasible & real_feasible)),
        "simulator_errors_corrected": int(np.sum(~baseline_correct & model_correct)),
        "simulator_correct_corrupted": int(np.sum(baseline_correct & ~model_correct)),
        "predicted_pass_rate": {
            "n": len(clipped_pass),
            "min": float(clipped_pass.min()),
            "max": float(clipped_pass.max()),
            "distinct_n": len(set(float(value) for value in clipped_pass)),
        },
    }


def compare(
    rows: Sequence[dict[str, Any]],
    baseline_prediction: np.ndarray,
    telemetry_prediction: np.ndarray,
    baseline: dict[str, Any],
    telemetry: dict[str, Any],
) -> dict[str, Any]:
    truth = np.asarray([row["target_real_feasible"] for row in rows], dtype=bool)
    baseline_feasible = np.clip(baseline_prediction, 0.0, 1.0) >= 0.95
    telemetry_feasible = np.clip(telemetry_prediction, 0.0, 1.0) >= 0.95
    baseline_correct = baseline_feasible == truth
    telemetry_correct = telemetry_feasible == truth
    return {
        "delta_telemetry_minus_baseline": {
            "rmse": telemetry["rmse"] - baseline["rmse"],
            "mae": telemetry["mae"] - baseline["mae"],
            "feasibility_accuracy": telemetry["feasibility_accuracy"]
            - baseline["feasibility_accuracy"],
        },
        "baseline_errors_corrected": int(
            np.sum(~baseline_correct & telemetry_correct)
        ),
        "baseline_correct_corrupted": int(
            np.sum(baseline_correct & ~telemetry_correct)
        ),
    }


def execute(args: argparse.Namespace) -> dict[str, Any]:
    paired = json.loads(args.paired_state.read_text(encoding="utf-8"))
    if paired.get("status") != "PASS" or len(paired["examples"]) != 12:
        raise RuntimeError("paired P1 state evidence is incomplete")
    rows = transitions(paired["examples"])
    residual_truth = np.asarray(
        [row["target_residual"] for row in rows], dtype=np.float64
    )
    pass_truth = np.asarray(
        [row["target_real_pass_rate"] for row in rows], dtype=np.float64
    )
    sim_pass = np.asarray(
        [row["target_sim_pass_rate"] for row in rows], dtype=np.float64
    )
    simulator = metrics(rows, sim_pass)
    sensitivity = {}
    for regularization in REGULARIZATION:
        hybrid_base_delta = grouped_predictions(
            rows,
            feature_name="hybrid_base",
            target_name="target_residual_delta",
            regularization=regularization,
        )
        hybrid_telemetry_delta = grouped_predictions(
            rows,
            feature_name="hybrid_telemetry",
            target_name="target_residual_delta",
            regularization=regularization,
        )
        source_residual = np.asarray(
            [row["source_residual"] for row in rows], dtype=np.float64
        )
        hybrid_base_correction = source_residual + hybrid_base_delta
        hybrid_telemetry_correction = source_residual + hybrid_telemetry_delta
        hybrid_base_prediction = sim_pass + hybrid_base_correction
        hybrid_telemetry_prediction = sim_pass + hybrid_telemetry_correction
        hybrid_base = metrics(rows, hybrid_base_prediction)
        hybrid_telemetry = metrics(rows, hybrid_telemetry_prediction)
        prior_shrinkage = {}
        for weight in PRIOR_SHRINKAGE:
            prior_shrinkage[str(weight)] = {
                "raw_simulator_prior": {
                    "sim_plus_outcome": metrics(
                        rows, sim_pass + weight * hybrid_base_correction
                    ),
                    "sim_plus_outcome_plus_telemetry": metrics(
                        rows, sim_pass + weight * hybrid_telemetry_correction
                    ),
                },
                "anchor_offset_prior": {
                    "sim_plus_outcome": metrics(
                        rows,
                        sim_pass + source_residual + weight * hybrid_base_delta,
                    ),
                    "sim_plus_outcome_plus_telemetry": metrics(
                        rows,
                        sim_pass + source_residual + weight * hybrid_telemetry_delta,
                    ),
                },
            }

        direct_base_prediction = grouped_predictions(
            rows,
            feature_name="direct_base",
            target_name="target_real_pass_rate_delta",
            regularization=regularization,
        )
        direct_telemetry_prediction = grouped_predictions(
            rows,
            feature_name="direct_telemetry",
            target_name="target_real_pass_rate_delta",
            regularization=regularization,
        )
        source_real_pass = np.asarray(
            [row["source_real_pass_rate"] for row in rows], dtype=np.float64
        )
        direct_base_prediction = source_real_pass + direct_base_prediction
        direct_telemetry_prediction = source_real_pass + direct_telemetry_prediction
        direct_base = metrics(rows, direct_base_prediction)
        direct_telemetry = metrics(rows, direct_telemetry_prediction)
        sensitivity[str(regularization)] = {
            "hybrid": {
                "sim_plus_outcome": hybrid_base,
                "sim_plus_outcome_plus_telemetry": hybrid_telemetry,
                "comparison": compare(
                    rows,
                    hybrid_base_prediction,
                    hybrid_telemetry_prediction,
                    hybrid_base,
                    hybrid_telemetry,
                ),
                "prior_shrinkage": prior_shrinkage,
            },
            "direct": {
                "real_outcome_only": direct_base,
                "telemetry_only": direct_telemetry,
                "comparison": compare(
                    rows,
                    direct_base_prediction,
                    direct_telemetry_prediction,
                    direct_base,
                    direct_telemetry,
                ),
            },
        }
    red_flags = []
    if any(not math.isfinite(value) for value in residual_truth):
        red_flags.append("nonfinite_residual_target")
    if any(not math.isfinite(value) for value in pass_truth):
        red_flags.append("nonfinite_pass_rate_target")
    result = {
        "schema": "telemetry-residual-cross-config-diagnostic-v1",
        "status": "PASS" if not red_flags else "STOP",
        "scope": (
            "single P1 trace/SLO-task development diagnostic; ordered transitions "
            "are not independent tasks and cannot support a generalization claim"
        ),
        "split": (
            "hold out both ordered source config and target config identities; "
            "each fold contains both offered-load anchors"
        ),
        "features": {
            "sim_plus_outcome": len(rows[0]["hybrid_base"]),
            "sim_plus_outcome_plus_telemetry": len(rows[0]["hybrid_telemetry"]),
            "real_outcome_only": len(rows[0]["direct_base"]),
            "telemetry_only": len(rows[0]["direct_telemetry"]),
        },
        "simulator": simulator,
        "regularization_sensitivity": sensitivity,
        "red_flags": red_flags,
        "sanity": {
            "transitions": len(rows),
            "load_levels": len(
                {row["source_level"] for row in rows}
                | {row["target_level"] for row in rows}
            ),
            "cross_load_transitions": sum(
                row["source_level"] != row["target_level"] for row in rows
            ),
            "ordered_cell_pairs": len(
                {(row["source_cell"], row["target_cell"]) for row in rows}
            ),
            "target_residual": {
                "n": len(residual_truth),
                "min": float(residual_truth.min()),
                "max": float(residual_truth.max()),
                "distinct_n": len(set(float(value) for value in residual_truth)),
            },
            "target_real_pass_rate": {
                "n": len(pass_truth),
                "min": float(pass_truth.min()),
                "max": float(pass_truth.max()),
                "distinct_n": len(set(float(value) for value in pass_truth)),
            },
            "invariants": {
                "finite_targets": not red_flags,
                "ratios_bounded": all(
                    0.0 <= row["target_sim_pass_rate"] <= 1.0 for row in rows
                ),
                "source_differs_from_target": all(
                    row["source_cell"] != row["target_cell"] for row in rows
                ),
                "per_config_not_identical": len(
                    set(float(value) for value in pass_truth)
                )
                > 1,
            },
        },
    }
    atomic_json(args.output, result)
    if result["status"] != "PASS":
        raise RuntimeError(red_flags)
    return result


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser()
    result.add_argument("--paired-state", type=Path, required=True)
    result.add_argument("--output", type=Path, required=True)
    return result


def main() -> None:
    result = execute(parser().parse_args())
    print(
        json.dumps(
            {
                "status": result["status"],
                "transitions": result["sanity"]["transitions"],
                "sensitivity": result["regularization_sensitivity"],
                "sanity": result["sanity"],
                "red_flags": result["red_flags"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
