"""Quantized weights for the port (MLX affine quantization), with a per-tensor recipe chosen by measurement.

A quantized weight is stored as `<name>.q` (packed uint32), `<name>.scales`, `<name>.biases`; `quant.json` in the
model folder records bits and group size per tensor. The runtime's linear layers call `mm`, which uses
`mx.quantized_matmul` for a quantized weight and a plain matmul otherwise; embedding tables use `rows`.
"""
import json
import os

import mlx.core as mx


class QW:
    """A quantized weight [out, in]: x @ W.T through mx.quantized_matmul."""
    __slots__ = ("q", "scales", "biases", "bits", "group_size", "shape")

    def __init__(self, q, scales, biases, bits, group_size, shape):
        self.q, self.scales, self.biases = q, scales, biases
        self.bits, self.group_size, self.shape = int(bits), int(group_size), tuple(shape)

    @classmethod
    def from_dense(cls, w, bits, group_size):
        q, s, b = mx.quantize(w, group_size=group_size, bits=bits)
        return cls(q, s, b, bits, group_size, w.shape)

    def dense(self, dtype=None):
        w = mx.dequantize(self.q, self.scales, self.biases, group_size=self.group_size, bits=self.bits)
        return w if dtype is None else w.astype(dtype)

    def arrays(self):
        return [self.q, self.scales, self.biases]


def mm(x, w):
    """x @ w.T for a dense [out, in] weight or a QW."""
    if isinstance(w, QW):
        return mx.quantized_matmul(x, w.q, w.scales, w.biases, transpose=True, group_size=w.group_size, bits=w.bits)
    return x @ w.T


def rows(w, idx):
    """Embedding lookup w[idx] for a dense table or a QW (dequantizes only the gathered rows)."""
    if isinstance(w, QW):
        return mx.dequantize(w.q[idx], w.scales[idx], w.biases[idx], group_size=w.group_size, bits=w.bits)
    return w[idx]


def cast(w, dtype):
    """A dense weight is cast; a QW keeps its packed values and casts its scales and biases."""
    if isinstance(w, QW):
        return QW(w.q, w.scales.astype(dtype), w.biases.astype(dtype), w.bits, w.group_size, w.shape)
    return w.astype(dtype)


def arrays(tree):
    """Every mx.array in a nested dict/list/QW structure (for mx.eval)."""
    if isinstance(tree, QW):
        return tree.arrays()
    if isinstance(tree, dict):
        return [a for v in tree.values() for a in arrays(v)]
    if isinstance(tree, (list, tuple)):
        return [a for v in tree for a in arrays(v)]
    return [tree] if isinstance(tree, mx.array) else []


def load_recipe(path):
    p = os.path.join(path, "quant.json")
    return json.load(open(p)) if os.path.exists(p) else None


def take(w, name, recipe):
    """Pop `name` from a loaded file dict as a dense array or a QW, as quant.json says."""
    if recipe and name in recipe["tensors"]:
        r = recipe["tensors"][name]
        return QW(w.pop(name + ".q"), w.pop(name + ".scales"), w.pop(name + ".biases"), r["bits"], r["group_size"],
                  r["shape"])
    return w.pop(name)