"""Per-layer slot table: a bounded set of resident experts held in stacked MLX buffers.

Compute always runs `gather_qmm` on the slot buffers with expert indices remapped to slot indices,
so the kernel is the same whether an expert was resident or just loaded. Loading an expert into a
slot is an in-place slice assignment (measured 0.3 ms for 9 tensors / 3.4 MB on M5 Max, see
benchmarks/raw/gateB/design/). Eviction policy is pluggable (Gate C); the default is LRU.
Every miss is accounted as a stall (wall time between the request and the slot being ready).
"""
from __future__ import annotations

import time
from collections import OrderedDict
from typing import Dict, Iterable, List, Optional

import mlx.core as mx

from .store import ExpertStore


class EvictionPolicy:
    """Interface: called on every access; must return a victim expert id when the table is full."""
    name = "lru"

    def __init__(self):
        self.order: "OrderedDict[int, None]" = OrderedDict()

    def touch(self, expert: int):
        self.order.pop(expert, None)
        self.order[expert] = None

    def victim(self, protected: Iterable[int]) -> int:
        prot = set(protected)
        for e in self.order:
            if e not in prot:
                return e
        raise RuntimeError("no evictable expert")

    def remove(self, expert: int):
        self.order.pop(expert, None)


class LFUPolicy(EvictionPolicy):
    name = "lfu"

    def __init__(self):
        super().__init__()
        self.count: Dict[int, int] = {}

    def touch(self, expert: int):
        super().touch(expert)
        self.count[expert] = self.count.get(expert, 0) + 1

    def victim(self, protected: Iterable[int]) -> int:
        prot = set(protected)
        cands = [(self.count.get(e, 0), i, e) for i, e in enumerate(self.order) if e not in prot]
        if not cands:
            raise RuntimeError("no evictable expert")
        return min(cands)[2]

    def remove(self, expert: int):
        super().remove(expert)
        self.count.pop(expert, None)


class SlotTable:
    def __init__(self, layer: int, capacity: int, store: ExpertStore, policy: Optional[EvictionPolicy] = None,
                 matrices=("gate", "up", "down"), fused_upgate: bool = False):
        self.layer = layer
        self.capacity = capacity
        self.store = store
        self.policy = policy or EvictionPolicy()
        self.matrices = matrices
        self.spec = store.spec
        self.slot_of: Dict[int, int] = {}
        self.expert_in: List[Optional[int]] = [None] * capacity
        self.free: List[int] = list(range(capacity - 1, -1, -1))
        self.buf: Dict[str, mx.array] = {}
        # allocate stacked buffers from the first expert's shapes
        es = self.spec.experts[(layer, 0)]
        from .store import MX_DTYPE
        self.fused = fused_upgate
        for name, loc in es.tensors.items():
            if self.fused and name.startswith(("up.", "gate.")):
                continue
            self.buf[name] = mx.zeros((capacity,) + tuple(loc.shape), dtype=MX_DTYPE[loc.dtype])
        if self.fused:  # [C, 2H, ...]: rows 0..H-1 = up, H..2H-1 = gate
            for part in ("weight", "scales", "biases"):
                loc = es.tensors.get(f"up.{part}")
                if loc is not None:
                    self.buf[f"upgate.{part}"] = mx.zeros((capacity, 2 * loc.shape[0]) + tuple(loc.shape[1:]), dtype=MX_DTYPE[loc.dtype])
        mx.eval(*self.buf.values())
        # accounting
        self.hits = 0
        self.misses = 0
        self.evictions = 0
        self.stall_time = 0.0
        self.loads = 0
        self.prefetched_used = 0
        self.hits_device = 0  # accesses served on the fast path (table complete, no host sync)
        self._complete = False
        self.slot_map = None  # device array expert -> slot, rebuilt when the table becomes complete

    # ---- residency management -------------------------------------------------
    def _install(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int]) -> int:
        if self.free:
            slot = self.free.pop()
        else:
            victim = self.policy.victim(protected)
            slot = self.slot_of.pop(victim)
            self.policy.remove(victim)
            self.evictions += 1
        if self.fused:
            H = tensors["up.weight"].shape[0]
            for part in ("weight", "scales", "biases"):
                if f"up.{part}" in tensors:
                    self.buf[f"upgate.{part}"][slot, :H] = tensors[f"up.{part}"]
                    self.buf[f"upgate.{part}"][slot, H:] = tensors[f"gate.{part}"]
            for name, arr in tensors.items():
                if name.startswith("down."):
                    self.buf[name][slot] = arr
        else:
            for name, arr in tensors.items():
                self.buf[name][slot] = arr
        self.slot_of[expert] = slot
        self.expert_in[slot] = expert
        self.policy.touch(expert)  # newly installed experts are MRU (prefetched ones included), so they are always evictable
        self.loads += 1
        self._complete = False
        return slot

    def complete(self) -> bool:
        """True iff every expert of the layer is resident (then no host synchronisation is needed)."""
        if not self._complete and len(self.slot_of) == self.spec.num_experts:
            self.slot_map = mx.array([self.slot_of[e] for e in range(self.spec.num_experts)], dtype=mx.uint32)
            mx.eval(self.slot_map)
            self._complete = True
        return self._complete

    def load(self, expert: int, protected: Iterable[int] = ()) -> int:
        """Load an expert from the store into a slot (demand miss). The slot write is a lazy slice assignment that the
        layer's gather_qmm will consume: no per-expert GPU synchronisation (the stall counted here is the read time)."""
        t = time.perf_counter()
        tensors = self.store.read_expert(self.layer, expert)
        self.stall_time += time.perf_counter() - t
        return self._install(expert, tensors, protected)

    def insert_prefetched(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int] = ()) -> Optional[int]:
        """Install tensors that were fetched asynchronously (no stall accounted)."""
        if expert in self.slot_of:
            return self.slot_of[expert]
        if not self.free and all(e in set(protected) for e in self.slot_of):
            return None
        return self._install(expert, tensors, protected)

    def ensure(self, experts: List[int]) -> List[int]:
        """Make every expert resident (demand-loading misses, read as a batch) and return their slot indices."""
        needed = set(experts)
        missing = [e for e in experts if e not in self.slot_of]
        self.hits += len(experts) - len(missing)
        self.misses += len(missing)
        if missing:
            t = time.perf_counter()
            tensors = self.store.read_experts(self.layer, missing)
            self.stall_time += time.perf_counter() - t
            for e in missing:
                self._install(e, tensors[e], needed)
        for e in experts:
            self.policy.touch(e)
        return [self.slot_of[e] for e in experts]

    def resident(self) -> List[int]:
        return list(self.slot_of)

    def stats(self) -> dict:
        n = self.hits + self.misses
        return {"layer": self.layer, "capacity": self.capacity, "hits": self.hits, "misses": self.misses, "hits_device": self.hits_device,
                "hit_rate": (self.hits / n) if n else None, "evictions": self.evictions, "loads": self.loads,
                "stall_time_s": self.stall_time, "policy": self.policy.name}

    def reset_stats(self):
        self.hits = self.misses = self.evictions = self.loads = self.hits_device = 0
        self.stall_time = 0.0