"""
correlation_table.py
====================
专家相关性感知预取表（Correlation-aware Expert Prefetcher）

核心思想（论文 §6.2）：
  利用"前 l 层专家路径"预测当前层 token 的专家倾向。
"""

from __future__ import annotations

import logging
from collections import defaultdict, Counter
from typing import TYPE_CHECKING, Dict, List, Optional

import torch
import numpy as np

if TYPE_CHECKING:
    from srep_moe.metrics.tracer import ExpertTracer

logger = logging.getLogger(__name__)


class CorrelationTable:
    """
    专家路径相关性统计表。

    Args:
        num_layers  : transformer 层总数
        num_experts : 每层专家数
        top_k       : 每 token 激活专家数
        path_len    : 查表时使用的历史层数（论文默认 l=1）
        min_count   : 最小共现次数阈值，低于此值的预测被忽略
    """

    def __init__(
        self,
        num_layers: int,
        num_experts: int,
        top_k: int,
        path_len: int = 1,
        min_count: int = 2,
        tracer: Optional["ExpertTracer"] = None,
    ):
        self.num_layers = num_layers
        self.num_experts = num_experts
        self.top_k = top_k
        self.path_len = path_len
        self.min_count = min_count
        self._tracer = tracer

        # correlation[layer_idx][prev_expert_id] = Counter
        self._table: Dict[int, Dict[int, Counter]] = defaultdict(
            lambda: defaultdict(Counter)
        )

        self._total_samples = 0
        self._total_hits = 0

    # ------------------------------------------------------------------
    # 更新接口
    # ------------------------------------------------------------------

    def update(
        self,
        layer_idx: int,
        prev_expert_ids: List[int],
        curr_expert_ids: List[int],
    ) -> None:
        for prev_e in prev_expert_ids:
            for curr_e in curr_expert_ids:
                self._table[layer_idx][prev_e][curr_e] += 1
        self._total_samples += len(curr_expert_ids)

        if self._tracer:
            self._tracer.emit({
                "event": "ct_update",
                "layer": layer_idx,
                "prev_experts": [int(e) for e in prev_expert_ids[:16]],
                "curr_experts": [int(e) for e in curr_expert_ids[:16]],
                "step": self._tracer.current_step,
                "phase": self._tracer.current_phase,
            })

    def update_batch(
        self,
        layer_idx: int,
        prev_routing: torch.Tensor,
        curr_routing: torch.Tensor,
    ) -> None:
        prev_flat = prev_routing.reshape(-1).tolist()
        curr_flat = curr_routing.reshape(-1).tolist()
        self.update(layer_idx, [int(x) for x in prev_flat], [int(x) for x in curr_flat])

    # ------------------------------------------------------------------
    # 查询接口
    # ------------------------------------------------------------------

    def predict_hot_experts(
        self,
        layer_idx: int,
        prev_routing: torch.Tensor,
        top_k_prefetch: Optional[int] = None,
    ) -> List[int]:
        if top_k_prefetch is None:
            top_k_prefetch = self.top_k

        if layer_idx not in self._table:
            return []

        layer_table = self._table[layer_idx]
        prev_flat = prev_routing.reshape(-1).tolist()

        aggregated: Counter = Counter()
        for prev_e in prev_flat:
            prev_e = int(prev_e)
            if prev_e in layer_table:
                for curr_e, cnt in layer_table[prev_e].items():
                    if cnt >= self.min_count:
                        aggregated[curr_e] += cnt

        if not aggregated:
            return []

        result = [e for e, _ in aggregated.most_common(top_k_prefetch)]

        if self._tracer:
            self._tracer.emit({
                "event": "ct_predict",
                "layer": layer_idx,
                "predicted": result[:16],
                "step": self._tracer.current_step,
            })

        return result

    def predict_hot_experts_multi_batch(
        self,
        layer_idx: int,
        batch_routings: List[torch.Tensor],
        top_k_prefetch: Optional[int] = None,
    ) -> List[int]:
        if top_k_prefetch is None:
            top_k_prefetch = self.top_k
        all_routings = torch.cat(batch_routings, dim=0)
        return self.predict_hot_experts(layer_idx, all_routings, top_k_prefetch)

    # ------------------------------------------------------------------
    # 统计与分析
    # ------------------------------------------------------------------

    def compute_hit_rate(
        self,
        layer_idx: int,
        prev_routing: torch.Tensor,
        actual_routing: torch.Tensor,
        top_k_prefetch: Optional[int] = None,
    ) -> float:
        predicted = set(self.predict_hot_experts(layer_idx, prev_routing, top_k_prefetch))
        actual = set(actual_routing.reshape(-1).tolist())
        if not actual:
            return 0.0
        return len(predicted & actual) / len(actual)

    def coverage_stats(self) -> dict:
        stats = {}
        for layer_idx, layer_table in self._table.items():
            total_pairs = sum(len(v) for v in layer_table.values())
            total_count = sum(c for v in layer_table.values() for c in v.values())
            stats[layer_idx] = {
                "unique_pairs": total_pairs,
                "total_count": total_count,
                "prev_experts_seen": len(layer_table),
            }
        return stats

    def reset(self) -> None:
        self._table.clear()
        self._total_samples = 0
        self._total_hits = 0
        logger.info("CorrelationTable reset.")

    def get_expert_frequency(self, layer_idx: int) -> np.ndarray:
        freq = np.zeros((self.num_experts, self.num_experts), dtype=np.float32)
        if layer_idx not in self._table:
            return freq
        for prev_e, counter in self._table[layer_idx].items():
            if prev_e >= self.num_experts:
                continue
            total = sum(counter.values())
            if total == 0:
                continue
            for curr_e, cnt in counter.items():
                if curr_e < self.num_experts:
                    freq[prev_e, curr_e] = cnt / total
        return freq
