"""
collector.py
============
Torrent 指标收集器

收集 4 类指标（论文 §9.1 验收标准）：
  1. 吞吐量（token/s）与时延（ms/token）
  2. Bubble 时间（inter-layer + intra-layer）
  3. GPU 峰值显存与 I/O 体积
  4. 专家预取命中率
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from statistics import mean
from typing import List, Optional

import torch

logger = logging.getLogger(__name__)


@dataclass
class StepMetrics:
    step_idx: int
    step_time_s: float
    tokens_generated: int
    n: int
    throughput_tps: float = 0.0
    latency_ms: float = 0.0


@dataclass
class LayerMetrics:
    layer_idx: int
    intra_bubble_s: float
    num_expert_used: int
    hot_hit_count: int
    hot_total: int


class MetricsCollector:
    """
    Torrent 运行指标收集与分析。

    Args:
        warmup_steps: 预热步数（不计入统计）
    """

    def __init__(self, warmup_steps: int = 2):
        self.warmup_steps = warmup_steps
        self._steps: List[StepMetrics] = []
        self._layers: List[LayerMetrics] = []
        self._peak_gpu_gb: float = 0.0
        self._io_volume_gb: float = 0.0
        self._step_count: int = 0

    # ------------------------------------------------------------------
    # 记录接口
    # ------------------------------------------------------------------

    def record_step(
        self,
        step_time_s: float,
        tokens_generated: int,
        n: int,
    ) -> StepMetrics:
        self._step_count += 1
        if self._step_count <= self.warmup_steps:
            return StepMetrics(
                step_idx=self._step_count,
                step_time_s=step_time_s,
                tokens_generated=tokens_generated,
                n=n,
            )

        throughput = tokens_generated / (step_time_s + 1e-12)
        latency_ms = step_time_s / (tokens_generated + 1e-12) * 1e3

        m = StepMetrics(
            step_idx=self._step_count,
            step_time_s=step_time_s,
            tokens_generated=tokens_generated,
            n=n,
            throughput_tps=throughput,
            latency_ms=latency_ms,
        )
        self._steps.append(m)
        return m

    def record_layer(
        self,
        layer_idx: int,
        intra_bubble_s: float,
        num_expert_used: int,
        hot_hit_count: int,
        hot_total: int,
    ) -> None:
        if self._step_count <= self.warmup_steps:
            return
        self._layers.append(LayerMetrics(
            layer_idx=layer_idx,
            intra_bubble_s=intra_bubble_s,
            num_expert_used=num_expert_used,
            hot_hit_count=hot_hit_count,
            hot_total=hot_total,
        ))

    def update_gpu_peak(self, peak_gb: float) -> None:
        self._peak_gpu_gb = max(self._peak_gpu_gb, peak_gb)

    def add_io_volume(self, volume_gb: float) -> None:
        self._io_volume_gb += volume_gb

    # ------------------------------------------------------------------
    # 统计汇总
    # ------------------------------------------------------------------

    def get_summary(self) -> dict:
        if not self._steps:
            return {"error": "No steps recorded (all warmup)"}

        tps_list = [s.throughput_tps for s in self._steps]
        lat_list = [s.latency_ms for s in self._steps]
        bubble_list = [l.intra_bubble_s * 1e3 for l in self._layers]

        hit_counts = sum(l.hot_hit_count for l in self._layers)
        hit_total = sum(l.hot_total for l in self._layers)

        return {
            "throughput": {
                "mean_tps": mean(tps_list),
                "p50_tps": self._percentile(tps_list, 50),
                "p90_tps": self._percentile(tps_list, 90),
                "min_tps": min(tps_list),
                "max_tps": max(tps_list),
            },
            "latency": {
                "mean_ms": mean(lat_list),
                "p50_ms": self._percentile(lat_list, 50),
                "p90_ms": self._percentile(lat_list, 90),
                "p99_ms": self._percentile(lat_list, 99),
                "min_ms": min(lat_list),
                "max_ms": max(lat_list),
            },
            "bubble": {
                "mean_intra_ms": mean(bubble_list) if bubble_list else 0.0,
                "p90_intra_ms": self._percentile(bubble_list, 90) if bubble_list else 0.0,
                "total_intra_ms": sum(bubble_list),
            },
            "prefetch_hit_rate": hit_counts / max(hit_total, 1),
            "peak_gpu_gb": self._peak_gpu_gb,
            "io_volume_gb": self._io_volume_gb,
            "num_steps": len(self._steps),
            "total_tokens": sum(s.tokens_generated for s in self._steps),
        }

    def compare_with_baseline(
        self,
        baseline_collector: "MetricsCollector",
        label: str = "Torrent vs Baseline",
    ) -> dict:
        our = self.get_summary()
        base = baseline_collector.get_summary()

        if "error" in our or "error" in base:
            return {"error": "Insufficient data for comparison"}

        speedup = our["throughput"]["mean_tps"] / max(base["throughput"]["mean_tps"], 1e-9)
        lat_reduction = 1.0 - our["latency"]["mean_ms"] / max(base["latency"]["mean_ms"], 1e-9)
        bubble_reduction = (
            1.0 - our["bubble"]["mean_intra_ms"] / max(base["bubble"]["mean_intra_ms"], 1e-9)
            if base["bubble"]["mean_intra_ms"] > 0 else 0.0
        )

        return {
            "label": label,
            "throughput_speedup": speedup,
            "latency_reduction_pct": lat_reduction * 100,
            "intra_bubble_reduction_pct": bubble_reduction * 100,
            "torrent_tps": our["throughput"]["mean_tps"],
            "baseline_tps": base["throughput"]["mean_tps"],
            "torrent_latency_ms": our["latency"]["mean_ms"],
            "baseline_latency_ms": base["latency"]["mean_ms"],
            "prefetch_hit_rate": our["prefetch_hit_rate"],
        }

    def print_summary(self, title: str = "Torrent Metrics Summary") -> None:
        s = self.get_summary()
        if "error" in s:
            print(f"[{title}] {s['error']}")
            return

        print(f"\n{'=' * 60}")
        print(f"  {title}")
        print(f"{'=' * 60}")
        print(f"  Steps              : {s['num_steps']}")
        print(f"  Total Tokens       : {s['total_tokens']}")
        print(f"  Throughput (mean)  : {s['throughput']['mean_tps']:.2f} token/s")
        print(f"  Throughput (p90)   : {s['throughput']['p90_tps']:.2f} token/s")
        print(f"  Latency (mean)     : {s['latency']['mean_ms']:.3f} ms/token")
        print(f"  Latency (p99)      : {s['latency']['p99_ms']:.3f} ms/token")
        print(f"  Intra-bubble (mean): {s['bubble']['mean_intra_ms']:.3f} ms")
        print(f"  Prefetch Hit Rate  : {s['prefetch_hit_rate'] * 100:.1f}%")
        print(f"  Peak GPU           : {s['peak_gpu_gb']:.2f} GB")
        print(f"  I/O Volume         : {s['io_volume_gb']:.2f} GB")
        print(f"{'=' * 60}\n")

    def reset(self) -> None:
        self._steps.clear()
        self._layers.clear()
        self._peak_gpu_gb = 0.0
        self._io_volume_gb = 0.0
        self._step_count = 0

    @staticmethod
    def _percentile(data: List[float], p: int) -> float:
        if not data:
            return 0.0
        if len(data) == 1:
            return data[0]
        sorted_data = sorted(data)
        idx = (p / 100) * (len(sorted_data) - 1)
        lo, hi = int(idx), min(int(idx) + 1, len(sorted_data) - 1)
        frac = idx - lo
        return sorted_data[lo] * (1 - frac) + sorted_data[hi] * frac

    def get_per_step_data(self) -> List[dict]:
        return [
            {
                "step": m.step_idx,
                "throughput_tps": m.throughput_tps,
                "latency_ms": m.latency_ms,
                "n": m.n,
            }
            for m in self._steps
        ]

    def get_bubble_timeline(self) -> List[dict]:
        return [
            {
                "layer": m.layer_idx,
                "intra_bubble_ms": m.intra_bubble_s * 1e3,
                "hit_rate": m.hot_hit_count / max(m.hot_total, 1),
            }
            for m in self._layers
        ]
