"""정정된 Math Ink 0.6 composite online/raster CPU 지연을 대표 입력으로 측정한다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
import math
from pathlib import Path
import statistics
import sys
import time

import torch

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
for path in (PROJECT_ROOT, SOURCE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from math_grid_drawer.research.math_ink_06 import MathInk06Engine
from scripts.export_math_ink_06_litert import _load_representative_inputs06


def percentile_nearest_rank06(values: list[float], percentile: float) -> float:
    """필요 변수: 측정값·0~1 percentile. 작동 원리: 모바일 p95와 동일한 nearest-rank 값을 반환한다."""

    if not values or not 0.0 <= percentile <= 1.0:
        raise ValueError("percentile 입력이 유효하지 않습니다.")
    ordered = sorted(values)
    index = max(0, min(len(ordered) - 1, math.ceil(percentile * len(ordered)) - 1))
    return float(ordered[index])


def module_state_bytes06(*modules: torch.nn.Module) -> int:
    """필요 변수: model·adapter module. 작동 원리: 중복 storage를 한 번만 세어 실제 tensor state bytes를 계산한다."""

    seen: set[tuple[int, int]] = set()
    total = 0
    for module in modules:
        for tensor in [*module.parameters(), *module.buffers()]:
            storage = tensor.untyped_storage()
            key = (storage.data_ptr(), storage.nbytes())
            if key in seen:
                continue
            seen.add(key)
            total += storage.nbytes()
    return total


def _rss_bytes06() -> int | None:
    """필요 변수: 없음. 작동 원리: psutil이 있으면 현재 process RSS를 반환하고 없으면 명시적으로 결측 처리한다."""

    try:
        import psutil
    except ImportError:
        return None
    return int(psutil.Process().memory_info().rss)


def _measure06(callable_, inputs: list[tuple[torch.Tensor, ...]], warmup: int) -> dict:
    """필요 변수: 고정 inference callable·대표 입력·warmup. 작동 원리: 표본별 wall latency와 output checksum을 측정한다."""

    with torch.inference_mode():
        for arguments in inputs[:max(1, min(warmup, len(inputs)))]:
            callable_(*arguments)
        latencies, checksum = [], 0
        observed_rss = _rss_bytes06()
        for arguments in inputs:
            started = time.perf_counter()
            output = callable_(*arguments)
            latencies.append((time.perf_counter() - started) * 1000.0)
            primary = output[0] if isinstance(output, tuple) else output
            checksum = (checksum * 131 + int(primary.argmax(dim=-1)[0])) % 2_147_483_647
            current_rss = _rss_bytes06()
            if current_rss is not None:
                observed_rss = max(observed_rss or 0, current_rss)
    return {
        "samples": len(latencies),
        "mean_ms": statistics.fmean(latencies),
        "p50_ms": percentile_nearest_rank06(latencies, 0.50),
        "p95_ms": percentile_nearest_rank06(latencies, 0.95),
        "maximum_ms": max(latencies),
        "output_checksum": checksum,
        "observed_process_rss_bytes": observed_rss,
    }


def main() -> None:
    """필요 변수: composite artifact·대표 cache. 작동 원리: thread별 두 inference 경로를 독립 측정해 JSON으로 남긴다."""

    parser = argparse.ArgumentParser(description="Benchmark Math Ink 0.6 composite CPU")
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--adapter-checkpoint", type=Path, required=True)
    parser.add_argument("--representative-inputs", type=Path, required=True)
    parser.add_argument("--threads", type=int, action="append", default=None)
    parser.add_argument("--samples", type=int, default=76)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    requested_threads = args.threads or [1, 2, 4]
    if any(value <= 0 for value in requested_threads):
        raise ValueError("CPU thread는 양수여야 합니다.")
    torch.set_num_interop_threads(1)
    engine = MathInk06Engine(
        args.checkpoint, adapter_checkpoint=args.adapter_checkpoint, device="cpu",
    )
    online_inputs, raster_inputs = _load_representative_inputs06(args.representative_inputs)
    online_inputs = online_inputs[:args.samples]
    raster_inputs = raster_inputs[:args.samples]

    def online_forward(sequence: torch.Tensor):
        """필요 변수: canonical sequence. 작동 원리: 실제 runtime online composite branch를 호출한다."""

        return engine.model.forward_online(engine.online_adapter(sequence))

    def raster_forward(raster: torch.Tensor):
        """필요 변수: raster. 작동 원리: 실제 runtime virtual stroke·raster adapter·fusion을 호출한다."""

        output = engine._forward_raster_composite06(raster)
        return engine.fuse_raster_output(output)[0]

    baseline_rss = _rss_bytes06()
    rows = []
    for thread_count in requested_threads:
        torch.set_num_threads(thread_count)
        online = _measure06(online_forward, online_inputs, args.warmup)
        raster = _measure06(raster_forward, raster_inputs, args.warmup)
        rows.append({
            "threads": thread_count,
            "online": online,
            "raster": raster,
            "proxy_gates": {
                "online_p95_le_50ms": online["p95_ms"] <= 50.0,
                "raster_p95_le_200ms": raster["p95_ms"] <= 200.0,
            },
        })
    observed_rss_values = [
        int(metrics["observed_process_rss_bytes"])
        for row in rows for metrics in (row["online"], row["raster"])
        if metrics["observed_process_rss_bytes"] is not None
    ]
    maximum_observed_rss = max(observed_rss_values) if observed_rss_values else None
    inference_rss_growth = (
        max(0, maximum_observed_rss - baseline_rss)
        if maximum_observed_rss is not None and baseline_rss is not None else None
    )
    state_bytes = module_state_bytes06(engine.model, engine.composite_adapter)
    model_plus_inference = (
        state_bytes + inference_rss_growth
        if inference_rss_growth is not None else None
    )
    report = {
        "schema": "aiflow-math-ink-06-composite-cpu-benchmark-v1",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "torch_version": torch.__version__,
        "platform": sys.platform,
        "model_version": engine.model_version,
        "model_state_bytes": state_bytes,
        "baseline_process_rss_bytes": baseline_rss,
        "maximum_observed_process_rss_bytes": maximum_observed_rss,
        "inference_rss_growth_bytes": inference_rss_growth,
        "model_state_plus_inference_growth_bytes": model_plus_inference,
        "memory_proxy_gate_le_100mib": (
            model_plus_inference <= 100 * 1024 * 1024
            if model_plus_inference is not None else None
        ),
        "rows": rows,
        "interpretation_limit": (
            "Windows PyTorch CPU proxy이며 Android LiteRT·배터리·delegate 성능 판정이 아니다."
        ),
        "product_validation": False,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()