#!/usr/bin/env python3
"""Offline predictor ablation for EXP-ATTN-STRUCTURED.

This is deliberately profile-only: it decides whether the structured model is
good enough to justify the expensive 7-cell trace replay.
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor

ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--profile",
        type=Path,
        default=REPO
        / "runs/frontier-prefill-kvgrowth-fix-v0/profiles/"
        "profile-v5-kvgrowth/attention.csv",
    )
    parser.add_argument(
        "--frontier-checkout",
        type=Path,
        default=Path("/tmp/frontier-attn-structured-v0"),
    )
    parser.add_argument("--output-root", type=Path, default=ROOT / "results")
    return parser.parse_args()


def normalize_bool(series: pd.Series) -> pd.Series:
    return series.astype(str).str.strip().str.lower().isin(
        {"1", "true", "t", "yes", "y"}
    )


def load_profile(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path).drop_duplicates()
    for column in ("is_prefill", "is_true_mixed_batch"):
        df[column] = normalize_bool(df[column])
    df = df[
        (df["n_embd"] == 2048)
        & (df["n_q_head"] == 32)
        & (df["n_kv_head"] == 4)
        & (df["block_size"] == 16)
        & df["is_prefill"]
        & ~df["is_true_mixed_batch"]
        & (df["prefill_chunk_size"] > 0)
    ].copy()
    df["prefill_chunk_size_squared"] = df["prefill_chunk_size"] ** 2
    return df


def mape(actual: np.ndarray, predicted: np.ndarray) -> float:
    return float(np.mean(np.abs((predicted - actual) / actual)) * 100)


def make_rf() -> RandomForestRegressor:
    # Exact best parameters selected by the current profile-v5 GridSearchCV.
    return RandomForestRegressor(
        random_state=0,
        n_estimators=250,
        max_depth=8,
        min_samples_split=2,
    )


def features(df: pd.DataFrame) -> pd.DataFrame:
    return df[["kv_cache_size", "prefill_chunk_size_squared"]]


def score_model(
    name: str,
    estimator: Any,
    train: pd.DataFrame,
    single: pd.DataFrame,
    grid: pd.DataFrame,
) -> dict[str, Any]:
    target = "time_stats.attn_prefill.median"
    estimator.fit(features(train), train[target])
    grid_prediction = estimator.predict(features(grid))
    single_prediction = estimator.predict(features(single))

    heldout_actual: list[float] = []
    heldout_prediction: list[float] = []
    for context in sorted(grid["kv_cache_size"].unique()):
        test = grid[grid["kv_cache_size"] == context]
        fold_train = train.drop(index=test.index, errors="ignore")
        fold_model = (
            make_rf()
            if name.startswith("rf")
            else estimator.__class__()
        )
        fold_model.fit(features(fold_train), fold_train[target])
        heldout_actual.extend(test[target].astype(float))
        heldout_prediction.extend(fold_model.predict(features(test)))

    q = np.arange(1, 8193, dtype=float)
    q_deltas: list[float] = []
    prediction_min: list[float] = []
    for context in (0, 8192, 16384, 24576, 32768, 40912):
        X = pd.DataFrame(
            {
                "kv_cache_size": np.full_like(q, context),
                "prefill_chunk_size_squared": q**2,
            }
        )
        prediction = estimator.predict(X)
        prediction_min.append(float(prediction.min()))
        q_deltas.append(float(np.diff(prediction).min()))

    kv = np.arange(0, 40913, 64, dtype=float)
    kv_deltas: list[float] = []
    for query in (64, 512, 2048, 4096, 8192):
        X = pd.DataFrame(
            {
                "kv_cache_size": kv,
                "prefill_chunk_size_squared": np.full_like(kv, query**2),
            }
        )
        kv_deltas.append(float(np.diff(estimator.predict(X)).min()))

    heldout_actual_array = np.asarray(heldout_actual)
    heldout_prediction_array = np.asarray(heldout_prediction)
    return {
        "candidate": name,
        "training_rows": len(train),
        "grid_fit_mape_pct": mape(
            grid[target].to_numpy(), np.asarray(grid_prediction)
        ),
        "single_fit_mape_pct": mape(
            single[target].to_numpy(), np.asarray(single_prediction)
        ),
        "heldout_context_mape_pct": mape(
            heldout_actual_array, heldout_prediction_array
        ),
        "heldout_context_max_abs_error_pct": float(
            np.max(
                np.abs(
                    (heldout_prediction_array - heldout_actual_array)
                    / heldout_actual_array
                )
            )
            * 100
        ),
        "prediction_min_ms": min(prediction_min),
        "q_min_delta_ms": min(q_deltas),
        "kv_min_delta_ms": min(kv_deltas),
        "monotone_and_nonnegative": (
            min(prediction_min) >= 0
            and min(q_deltas) >= -1e-12
            and min(kv_deltas) >= -1e-12
        ),
    }


def main() -> None:
    args = parse_args()
    sys.path.insert(0, str(args.frontier_checkout))
    from frontier.execution_time_predictor.structured_attention_prefill import (
        StructuredAttentionPrefillRegressor,
    )

    df = load_profile(args.profile)
    records: list[dict[str, Any]] = []
    data_audit: dict[str, Any] = {}
    for tp in (1, 2, 4):
        tp_df = df[df["num_tensor_parallel_workers"] == tp].copy()
        single = tp_df[tp_df["batch_size"] == 1].copy()
        grid = single[
            single["prefill_chunk_size"].isin((2048, 4096, 8192))
            & (single["kv_cache_size"] > 0)
        ].copy()
        duplicate_groups = (
            tp_df.groupby(
                ["kv_cache_size", "prefill_chunk_size_squared"]
            )
            .size()
            .gt(1)
            .sum()
        )
        data_audit[f"tp{tp}"] = {
            "standard_rows": len(tp_df),
            "single_request_rows": len(single),
            "target_grid_rows": len(grid),
            "duplicate_feature_groups": int(duplicate_groups),
        }

        candidates = (
            ("rf_all", make_rf(), tp_df),
            ("rf_single", make_rf(), single),
            (
                "structured_single",
                StructuredAttentionPrefillRegressor(),
                single,
            ),
        )
        for name, model, train in candidates:
            result = score_model(name, model, train, single, grid)
            result["tp"] = tp
            records.append(result)

    structured = [r for r in records if r["candidate"] == "structured_single"]
    checks = {
        "heldout_context_mape_le_5pct": all(
            r["heldout_context_mape_pct"] <= 5 for r in structured
        ),
        "monotone_and_nonnegative": all(
            r["monotone_and_nonnegative"] for r in structured
        ),
    }
    checks["profile_gate"] = all(checks.values())
    payload = {
        "schema": "frontier-attn-structured-ablation-v1",
        "profile": str(args.profile.resolve()),
        "frontier_checkout": str(args.frontier_checkout.resolve()),
        "data_audit": data_audit,
        "results": records,
        "checks": checks,
    }

    args.output_root.mkdir(parents=True, exist_ok=True)
    (args.output_root / "predictor-ablation.json").write_text(
        json.dumps(payload, indent=2)
    )
    with (args.output_root / "predictor-ablation.csv").open(
        "w", newline=""
    ) as stream:
        writer = csv.DictWriter(stream, fieldnames=list(records[0]))
        writer.writeheader()
        writer.writerows(records)
    print(json.dumps(checks, indent=2))
    for row in records:
        print(
            f"TP{row['tp']} {row['candidate']:18s} "
            f"grid={row['grid_fit_mape_pct']:.2f}% "
            f"heldout={row['heldout_context_mape_pct']:.2f}% "
            f"max={row['heldout_context_max_abs_error_pct']:.2f}% "
            f"monotone={row['monotone_and_nonnegative']}"
        )


if __name__ == "__main__":
    main()