"""
hardware_profiler.py
====================
测量 Torrent 约束求解器所需的 6 个关键时延参数：

  t_c_attn  : 单 batch attention 计算时间 (s)
  t_c_gate  : gate 网络计算时间 (s)
  t_c_expert: 单专家 FFN 计算时间 (s)
  t_io_attn : attention 权重 CPU→GPU 传输时间 (s)
  t_io_gate : gate 权重 CPU→GPU 传输时间 (s)
  t_io_expert: 单专家权重 CPU→GPU 传输时间 (s)

同时测量可用带宽与 GPU/CPU 内存容量，供 ConstraintSolver 使用。
"""

from __future__ import annotations

import math
import time
import threading
import logging
from dataclasses import dataclass, field, asdict
from typing import Optional

import torch
import torch.nn as nn
import psutil

logger = logging.getLogger(__name__)


@dataclass
class HardwareProfile:
    """存储所有关键时延与带宽测量结果。"""

    # -------- 计算时延 (秒) --------
    t_c_attn: float = 0.0       # 单 batch attention 计算时间
    t_c_gate: float = 0.0       # gate 网络计算时间
    t_c_expert: float = 0.0     # 单专家 FFN 计算时间

    # -------- I/O 时延 (秒) --------
    t_io_attn: float = 0.0      # attention 权重传输时间
    t_io_gate: float = 0.0      # gate 权重传输时间
    t_io_expert: float = 0.0    # 单专家权重传输时间

    # -------- 带宽 (GB/s) --------
    pcie_bandwidth_gbs: float = 0.0   # PCIe CPU→GPU 实测带宽
    dram_bandwidth_gbs: float = 0.0   # CPU DRAM 读带宽

    # -------- 内存容量 (GB) --------
    gpu_total_gb: float = 0.0
    gpu_free_gb: float = 0.0
    cpu_total_gb: float = 0.0
    cpu_free_gb: float = 0.0

    # -------- 模型形状相关 --------
    hidden_size: int = 4096
    num_experts: int = 8
    top_k: int = 2
    num_layers: int = 32
    seq_len: int = 512
    batch_size: int = 1

    # -------- 权重大小 (bytes) --------
    attn_weight_bytes: int = 0
    gate_weight_bytes: int = 0
    expert_weight_bytes: int = 0

    def to_dict(self) -> dict:
        return asdict(self)

    def summary(self) -> str:
        lines = [
            "=" * 60,
            "  Hardware Profile Summary",
            "=" * 60,
            f"  PCIe Bandwidth  : {self.pcie_bandwidth_gbs:.2f} GB/s",
            f"  GPU Memory      : {self.gpu_free_gb:.1f} / {self.gpu_total_gb:.1f} GB free",
            f"  CPU Memory      : {self.cpu_free_gb:.1f} / {self.cpu_total_gb:.1f} GB free",
            "",
            f"  t_c_attn        : {self.t_c_attn * 1e3:.3f} ms",
            f"  t_c_gate        : {self.t_c_gate * 1e3:.3f} ms",
            f"  t_c_expert      : {self.t_c_expert * 1e3:.3f} ms",
            "",
            f"  t_io_attn       : {self.t_io_attn * 1e3:.3f} ms",
            f"  t_io_gate       : {self.t_io_gate * 1e3:.3f} ms",
            f"  t_io_expert     : {self.t_io_expert * 1e3:.3f} ms",
            "=" * 60,
        ]
        return "\n".join(lines)


class HardwareProfiler:
    """
    测量 Torrent 所需的硬件时延参数。

    使用仿真张量（不需要真实模型权重）进行测量，
    结果可直接传给 ConstraintSolver。

    Args:
        hidden_size      : 模型隐藏层维度
        intermediate_size: FFN 中间层维度（专家）
        num_experts      : 专家总数
        top_k            : 每 token 激活专家数
        num_heads        : attention 头数
        num_layers       : transformer 层数
        seq_len          : 输入序列长度
        batch_size       : 单 batch token 数
        dtype            : 权重数据类型（影响字节数与传输量）
        warmup_rounds    : GPU 预热轮数
        measure_rounds   : 正式测量轮数
        device           : 目标 GPU 设备
    """

    def __init__(
        self,
        hidden_size: int = 4096,
        intermediate_size: int = 14336,
        num_experts: int = 8,
        top_k: int = 2,
        num_heads: int = 32,
        num_kv_heads: int = 8,
        num_layers: int = 32,
        seq_len: int = 512,
        batch_size: int = 1,
        dtype: torch.dtype = torch.float16,
        warmup_rounds: int = 3,
        measure_rounds: int = 10,
        device: Optional[str] = None,
    ):
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_experts = num_experts
        self.top_k = top_k
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.num_layers = num_layers
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.dtype = dtype
        self.warmup_rounds = warmup_rounds
        self.measure_rounds = measure_rounds

        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)

        self._dtype_bytes = {
            torch.float32: 4,
            torch.float16: 2,
            torch.bfloat16: 2,
            torch.int8: 1,
        }.get(dtype, 2)

        logger.info(f"HardwareProfiler initialized. device={self.device}, dtype={dtype}")

    # ------------------------------------------------------------------
    # 公共接口
    # ------------------------------------------------------------------

    def run(self) -> HardwareProfile:
        """执行全量测量，返回 HardwareProfile。"""
        profile = HardwareProfile(
            hidden_size=self.hidden_size,
            num_experts=self.num_experts,
            top_k=self.top_k,
            num_layers=self.num_layers,
            seq_len=self.seq_len,
            batch_size=self.batch_size,
        )

        self._measure_memory(profile)
        profile.pcie_bandwidth_gbs = self._measure_pcie_bandwidth()

        profile.attn_weight_bytes = self._attn_weight_bytes()
        profile.gate_weight_bytes = self._gate_weight_bytes()
        profile.expert_weight_bytes = self._expert_weight_bytes()

        profile.t_io_attn = profile.attn_weight_bytes / (
            profile.pcie_bandwidth_gbs * 1e9 + 1e-9
        )
        profile.t_io_gate = profile.gate_weight_bytes / (
            profile.pcie_bandwidth_gbs * 1e9 + 1e-9
        )
        profile.t_io_expert = profile.expert_weight_bytes / (
            profile.pcie_bandwidth_gbs * 1e9 + 1e-9
        )

        if self.device.type == "cuda":
            profile.t_c_attn = self._measure_attn_compute()
            profile.t_c_gate = self._measure_gate_compute()
            profile.t_c_expert = self._measure_expert_compute()
        else:
            logger.warning("CUDA not available, using estimated compute times.")
            profile.t_c_attn = 0.05
            profile.t_c_gate = 0.001
            profile.t_c_expert = 0.01

        logger.info("\n" + profile.summary())
        return profile

    # ------------------------------------------------------------------
    # 内存容量
    # ------------------------------------------------------------------

    def _measure_memory(self, profile: HardwareProfile) -> None:
        mem = psutil.virtual_memory()
        profile.cpu_total_gb = mem.total / 1e9
        profile.cpu_free_gb = mem.available / 1e9

        if self.device.type == "cuda":
            free, total = torch.cuda.mem_get_info(self.device)
            profile.gpu_total_gb = total / 1e9
            profile.gpu_free_gb = free / 1e9
        else:
            profile.gpu_total_gb = 0.0
            profile.gpu_free_gb = 0.0

    # ------------------------------------------------------------------
    # PCIe 带宽测量
    # ------------------------------------------------------------------

    def _measure_pcie_bandwidth(self) -> float:
        if self.device.type != "cuda":
            return 8.0

        size_mb = 256
        n_elements = size_mb * 1024 * 1024 // self._dtype_bytes
        cpu_tensor = torch.zeros(n_elements, dtype=self.dtype, pin_memory=True)
        gpu_tensor = torch.zeros(n_elements, dtype=self.dtype, device=self.device)

        for _ in range(3):
            gpu_tensor.copy_(cpu_tensor, non_blocking=True)
            torch.cuda.synchronize(self.device)

        stream = torch.cuda.Stream(device=self.device)
        total_bytes = 0
        start = time.perf_counter()
        for _ in range(5):
            with torch.cuda.stream(stream):
                gpu_tensor.copy_(cpu_tensor, non_blocking=True)
            torch.cuda.synchronize(self.device)
            total_bytes += n_elements * self._dtype_bytes
        elapsed = time.perf_counter() - start

        bw = total_bytes / elapsed / 1e9
        logger.info(f"PCIe bandwidth: {bw:.2f} GB/s")
        return bw

    # ------------------------------------------------------------------
    # 权重字节数计算
    # ------------------------------------------------------------------

    def _attn_weight_bytes(self) -> int:
        head_dim = self.hidden_size // self.num_heads
        q = self.hidden_size * self.hidden_size
        k = self.num_kv_heads * head_dim * self.hidden_size
        v = self.num_kv_heads * head_dim * self.hidden_size
        o = self.hidden_size * self.hidden_size
        return (q + k + v + o) * self._dtype_bytes

    def _gate_weight_bytes(self) -> int:
        return self.hidden_size * self.num_experts * self._dtype_bytes

    def _expert_weight_bytes(self) -> int:
        gate_proj = self.hidden_size * self.intermediate_size
        up_proj = self.hidden_size * self.intermediate_size
        down_proj = self.intermediate_size * self.hidden_size
        return (gate_proj + up_proj + down_proj) * self._dtype_bytes

    # ------------------------------------------------------------------
    # GPU 计算时延测量
    # ------------------------------------------------------------------

    def _cuda_timed(self, fn, warmup: int, rounds: int) -> float:
        for _ in range(warmup):
            fn()
        torch.cuda.synchronize(self.device)

        start_evt = torch.cuda.Event(enable_timing=True)
        end_evt = torch.cuda.Event(enable_timing=True)
        start_evt.record()
        for _ in range(rounds):
            fn()
        end_evt.record()
        torch.cuda.synchronize(self.device)
        return start_evt.elapsed_time(end_evt) / 1000.0 / rounds  # seconds

    def _measure_attn_compute(self) -> float:
        H = self.hidden_size
        T = self.seq_len
        heads = self.num_heads
        kv_heads = self.num_kv_heads
        head_dim = H // heads

        x = torch.randn(1, T, H, dtype=self.dtype, device=self.device)
        q_proj = nn.Linear(H, H, bias=False, dtype=self.dtype).to(self.device)
        k_proj = nn.Linear(H, kv_heads * head_dim, bias=False, dtype=self.dtype).to(self.device)
        v_proj = nn.Linear(H, kv_heads * head_dim, bias=False, dtype=self.dtype).to(self.device)
        o_proj = nn.Linear(H, H, bias=False, dtype=self.dtype).to(self.device)

        def fn():
            q = q_proj(x).view(1, T, heads, head_dim).transpose(1, 2)
            k = k_proj(x).view(1, T, kv_heads, head_dim).transpose(1, 2)
            v = v_proj(x).view(1, T, kv_heads, head_dim).transpose(1, 2)
            out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
            o_proj(out.transpose(1, 2).reshape(1, T, H))

        return self._cuda_timed(fn, self.warmup_rounds, self.measure_rounds)

    def _measure_gate_compute(self) -> float:
        H = self.hidden_size
        E = self.num_experts
        T = self.seq_len

        x = torch.randn(T, H, dtype=self.dtype, device=self.device)
        gate = nn.Linear(H, E, bias=False, dtype=self.dtype).to(self.device)

        def fn():
            gate(x)

        return self._cuda_timed(fn, self.warmup_rounds, self.measure_rounds)

    def _measure_expert_compute(self) -> float:
        H = self.hidden_size
        I = self.intermediate_size
        T = max(1, self.batch_size)

        x = torch.randn(T, H, dtype=self.dtype, device=self.device)
        gate_proj = nn.Linear(H, I, bias=False, dtype=self.dtype).to(self.device)
        up_proj = nn.Linear(H, I, bias=False, dtype=self.dtype).to(self.device)
        down_proj = nn.Linear(I, H, bias=False, dtype=self.dtype).to(self.device)

        def fn():
            down_proj(torch.nn.functional.silu(gate_proj(x)) * up_proj(x))

        return self._cuda_timed(fn, self.warmup_rounds, self.measure_rounds)
