"""
constraint_solver.py
====================
Torrent 约束求解器：给定硬件测量参数，求满足 4 个重叠约束的最小 batch-group 大小 n。

论文 §7 给出 4 个关键约束（公式 1-4），本质是确保：
  - gate 和热专家能在 Attention 计算期间完成传输
  - 冷专家能在热专家计算期间完成传输
  - 下一层 Attention 权重能在冷专家计算期间完成传输
  - GPU 显存不超出容量上限

最终取满足所有不等式的最小整数 n，避免 KV cache 随 n 过大而暴涨。

关键不等式（论文公式简化版）：
  C1: n * t_c_attn >= t_io_gate + K * t_io_expert       # 热专家覆盖
  C2: K * t_c_expert >= (E-K) * t_io_expert             # 冷专家覆盖
  C3: (E-K) * t_c_expert >= t_io_attn_next              # 下一 attn 覆盖
  C4: n <= floor(gpu_free_gb * 1e9 / kv_cache_per_batch) # KV cache 上限
"""

from __future__ import annotations

import math
import logging
from dataclasses import dataclass
from typing import Optional

from srep_moe.planner.hardware_profiler import HardwareProfile

logger = logging.getLogger(__name__)


@dataclass
class PlannerDecision:
    """约束求解输出：最优 batch-group 参数与分析信息。"""

    n: int                          # batch-group 大小
    n_min_c1: int                   # C1 约束要求的最小 n
    n_max_kv: int                   # KV cache 约束允许的最大 n
    bubble_ratio: float             # 预估 bubble 时间占比
    pcie_utilization: float         # PCIe 利用率（理论值）
    feasible: bool                  # 是否找到可行 n
    warning: Optional[str] = None   # 约束冲突时的警告信息

    def summary(self) -> str:
        lines = [
            "=" * 60,
            "  Constraint Solver Decision",
            "=" * 60,
            f"  n (batch-group size)  : {self.n}",
            f"  n_min (C1 constraint) : {self.n_min_c1}",
            f"  n_max (KV cache cap)  : {self.n_max_kv}",
            f"  Bubble ratio          : {self.bubble_ratio * 100:.1f}%",
            f"  PCIe utilization      : {self.pcie_utilization * 100:.1f}%",
            f"  Feasible              : {self.feasible}",
        ]
        if self.warning:
            lines.append(f"  Warning               : {self.warning}")
        lines.append("=" * 60)
        return "\n".join(lines)


class ConstraintSolver:
    """
    Torrent I/O-Compute Planner。

    Args:
        profile          : HardwareProfiler 输出的 HardwareProfile
        kv_cache_dtype_bytes: KV cache 数据类型字节数（默认 fp16 = 2）
        max_n            : n 搜索上界（默认 64）
        safety_margin    : GPU 显存安全余量（0~1），默认保留 10%
    """

    def __init__(
        self,
        profile: HardwareProfile,
        kv_cache_dtype_bytes: int = 2,
        max_n: int = 64,
        safety_margin: float = 0.10,
    ):
        self.p = profile
        self.kv_dtype_bytes = kv_cache_dtype_bytes
        self.max_n = max_n
        self.safety_margin = safety_margin

    # ------------------------------------------------------------------
    # 主接口
    # ------------------------------------------------------------------

    def solve(self) -> PlannerDecision:
        """求解最小可行 n，返回 PlannerDecision。"""
        p = self.p

        # ---- C1: n * t_c_attn >= t_io_gate + K * t_io_expert --------
        rhs_c1 = p.t_io_gate + p.top_k * p.t_io_expert
        if p.t_c_attn <= 1e-12:
            n_min_c1 = self.max_n
            logger.warning("t_c_attn is near zero; defaulting n_min_c1 to max_n")
        else:
            n_min_c1 = math.ceil(rhs_c1 / p.t_c_attn)
        n_min_c1 = max(1, n_min_c1)

        # ---- C2: K * t_c_expert >= (E-K) * t_io_expert --------------
        cold_experts = p.num_experts - p.top_k
        c2_lhs = p.top_k * p.t_c_expert
        c2_rhs = cold_experts * p.t_io_expert
        c2_satisfied = c2_lhs >= c2_rhs

        # ---- C3: (E-K) * t_c_expert >= t_io_attn_next ---------------
        c3_lhs = cold_experts * p.t_c_expert
        c3_rhs = p.t_io_attn
        c3_satisfied = c3_lhs >= c3_rhs

        # ---- C4: KV cache 上限约束 -----------------------------------
        head_dim = p.hidden_size // 32  # 保守估算
        kv_per_layer = 2 * p.seq_len * 32 * head_dim * self.kv_dtype_bytes
        kv_total_per_batch = kv_per_layer * p.num_layers

        usable_gpu = p.gpu_free_gb * 1e9 * (1.0 - self.safety_margin)
        resident_bytes = (
            p.attn_weight_bytes + p.gate_weight_bytes + p.top_k * p.expert_weight_bytes
        )
        available_for_kv = max(0.0, usable_gpu - resident_bytes)
        if kv_total_per_batch > 0:
            n_max_kv = max(1, int(available_for_kv / kv_total_per_batch))
        else:
            n_max_kv = self.max_n
        n_max_kv = min(n_max_kv, self.max_n)

        # ---- 最终决策 -----------------------------------------------
        feasible = True
        warning = None

        if n_min_c1 > n_max_kv:
            n = n_max_kv
            feasible = False
            warning = (
                f"C1 requires n>={n_min_c1}, but KV cache limits n<={n_max_kv}. "
                f"Consider quantization or reducing seq_len. Using n={n}."
            )
            logger.warning(warning)
        else:
            n = n_min_c1

        if p.gpu_total_gb == 0:
            n = max(1, n_min_c1)
            n_max_kv = self.max_n
            feasible = True

        # ---- 分析指标 -----------------------------------------------
        t_pipeline = n * p.t_c_attn + p.num_experts * p.t_c_expert
        bubble_c1 = max(0.0, rhs_c1 - n * p.t_c_attn)
        bubble_c2 = max(0.0, c2_rhs - c2_lhs)
        bubble_c3 = max(0.0, c3_rhs - c3_lhs)
        total_bubble = bubble_c1 + bubble_c2 + bubble_c3
        bubble_ratio = total_bubble / (t_pipeline + total_bubble + 1e-12)

        total_io = p.t_io_gate + p.num_experts * p.t_io_expert + p.t_io_attn
        pcie_util = min(1.0, total_io / (t_pipeline + 1e-12))

        decision = PlannerDecision(
            n=n,
            n_min_c1=n_min_c1,
            n_max_kv=n_max_kv,
            bubble_ratio=bubble_ratio,
            pcie_utilization=pcie_util,
            feasible=feasible,
            warning=warning,
        )

        logger.info(
            f"Constraint analysis:\n"
            f"  C1 (gate+hot expert prefetch): n>={n_min_c1} "
            f"[{n*p.t_c_attn*1e3:.2f}ms >= {rhs_c1*1e3:.2f}ms]\n"
            f"  C2 (cold expert transfer)    : {'✓' if c2_satisfied else '✗'} "
            f"[{c2_lhs*1e3:.2f}ms >= {c2_rhs*1e3:.2f}ms]\n"
            f"  C3 (next attn prefetch)      : {'✓' if c3_satisfied else '✗'} "
            f"[{c3_lhs*1e3:.2f}ms >= {c3_rhs*1e3:.2f}ms]\n"
            f"  C4 (KV cache cap)            : n<={n_max_kv}"
        )
        logger.info("\n" + decision.summary())
        return decision

    # ------------------------------------------------------------------
    # 网格搜索：在 n=1..max_n 上扫描最优点
    # ------------------------------------------------------------------

    def sweep(self, n_values: Optional[list] = None) -> list[dict]:
        """扫描不同 n 值下的理论性能，返回分析字典列表（用于可视化）。"""
        if n_values is None:
            n_values = list(range(1, min(self.max_n + 1, 33)))

        p = self.p
        results = []
        for n in n_values:
            rhs_c1 = p.t_io_gate + p.top_k * p.t_io_expert
            bubble_c1 = max(0.0, rhs_c1 - n * p.t_c_attn)
            cold = p.num_experts - p.top_k
            bubble_c2 = max(0.0, cold * p.t_io_expert - p.top_k * p.t_c_expert)
            bubble_c3 = max(0.0, p.t_io_attn - cold * p.t_c_expert)

            t_compute = n * p.t_c_attn + p.num_experts * p.t_c_expert
            t_bubble = bubble_c1 + bubble_c2 + bubble_c3
            t_total = t_compute + t_bubble
            tokens_per_group = n * p.batch_size
            throughput = tokens_per_group / (t_total + 1e-12)

            results.append({
                "n": n,
                "t_total_ms": t_total * 1e3,
                "t_bubble_ms": t_bubble * 1e3,
                "bubble_ratio": t_bubble / (t_total + 1e-12),
                "throughput_tps": throughput,
            })

        return results
