"""TITAN expert execution block, drop-in replacement for mlx-lm's `SwitchGLU`.

The router (gate linear, softmax, top-k, gating weights, weighted sum) stays in the host model's
code so that expert identity and aggregation are exactly those of the reference implementation;
this module only decides *where* the selected experts' weights come from and executes them with
`mx.gather_qmm` on the layer's slot table.
"""
from __future__ import annotations

from typing import Callable, List, Optional

import mlx.core as mx
from mlx_lm.models.activations import swiglu

from .slots import SlotTable


class TitanSwitchGLU:
    """Callable with the SwitchGLU signature: (x[..., 1, D], indices[..., k]) -> y[..., k, D]."""

    def __init__(self, table: SlotTable, group_size: int, bits: int, activation=None,
                 on_route: Optional[Callable[[int, List[int]], None]] = None):
        self.table = table
        self.group_size = group_size
        self.bits = bits
        self.on_route = on_route  # hook(layer, unique_experts) used by predictors/prefetchers (Gate C)
        self.last_indices: Optional[List[int]] = None

    def _qmm(self, x, m, idx, sorted_indices):
        b = self.table.buf
        return mx.gather_qmm(x, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=idx,
                             transpose=True, group_size=self.group_size, bits=self.bits, sorted_indices=sorted_indices)

    def _run_device(self, x, slot_idx):
        x = mx.expand_dims(x, (-2, -3))
        x_up = self._qmm(x, "up", slot_idx, False)
        x_gate = self._qmm(x, "gate", slot_idx, False)
        h = swiglu(x_gate, x_up)
        y = self._qmm(h, "down", slot_idx, False)
        return y.squeeze(-2)

    def _run(self, x, flat, shape):
        uniq = list(dict.fromkeys(flat))
        if self.on_route is not None:
            self.on_route(self.table.layer, uniq)
        slot_of = dict(zip(uniq, self.table.ensure(uniq)))
        slot_idx = mx.array([slot_of[e] for e in flat], dtype=mx.uint32).reshape(shape)
        x = mx.expand_dims(x, (-2, -3))
        x_up = self._qmm(x, "up", slot_idx, False)
        x_gate = self._qmm(x, "gate", slot_idx, False)
        h = swiglu(x_gate, x_up)  # identical compiled activation to mlx-lm SwitchGLU
        y = self._qmm(h, "down", slot_idx, False)
        return y.squeeze(-2)

    def __call__(self, x: mx.array, indices: mx.array) -> mx.array:
        t = self.table
        if t.complete():
            # Fast path: every expert is resident, so the slot map is static and lives on the device: no host sync.
            self.n_chunks = 1
            self.last_indices = None
            t.hits_device += indices.size
            return self._run_device(x, t.slot_map[indices])
        # Materialise the router's choice: this is the unavoidable host sync of any demand-loading design.
        flat = indices.reshape(-1).tolist()
        self.last_indices = flat
        k = indices.shape[-1]
        n_rows = len(flat) // k
        cap = self.table.capacity
        x2 = x.reshape(n_rows, x.shape[-1])
        # Greedy row chunking: a chunk's set of unique experts must fit in the slot table (exact: rows are independent).
        chunks, cur, cur_set = [], [], set()
        for r in range(n_rows):
            es = set(flat[r * k:(r + 1) * k])
            if cur and len(cur_set | es) > cap:
                chunks.append(cur); cur, cur_set = [], set()
            cur.append(r); cur_set |= es
        if cur:
            chunks.append(cur)
        self.n_chunks = len(chunks)
        if len(chunks) == 1:
            return self._run(x, flat, indices.shape)
        outs = []
        for rows in chunks:
            f = [flat[r * k + j] for r in rows for j in range(k)]
            outs.append(self._run(x2[rows[0]:rows[-1] + 1], f, (len(rows), k)))
        return mx.concatenate(outs, axis=0).reshape(indices.shape + (x.shape[-1],))