import torch

@torch.no_grad()
def unpack_sc_lut_indices(packed: torch.Tensor, n_tiles: int, bits: int) -> torch.Tensor:
    R = packed.shape[0]
    if bits == 10:
        p = packed.view(R, n_tiles // 4, 5).to(torch.int32)
        i0 = p[:, :, 0] | ((p[:, :, 1] & 0x03) << 8)
        i1 = (p[:, :, 1] >> 2) | ((p[:, :, 2] & 0x0F) << 6)
        i2 = (p[:, :, 2] >> 4) | ((p[:, :, 3] & 0x3F) << 4)
        i3 = (p[:, :, 3] >> 6) | (p[:, :, 4] << 2)
        return torch.stack([i0, i1, i2, i3], dim=-1).reshape(R, n_tiles)
    if bits == 9:
        p = packed.view(R, n_tiles // 8, 9).to(torch.int64)
        lo = torch.zeros((R, n_tiles // 8), dtype=torch.int64, device=packed.device)
        for k in range(8):
            lo |= (p[:, :, k] & 0xFF) << (8 * k)
        out = [((lo >> (9 * t)) & 0x1FF).to(torch.int32) for t in range(7)]
        i7 = (((lo >> 63) & 0x01) | ((p[:, :, 8] & 0xFF) << 1)).to(torch.int32)
        out.append(i7)
        return torch.stack(out, dim=-1).reshape(R, n_tiles)
    raise ValueError(f"Unsupported bit width: {bits}")

@torch.no_grad()
def decode_sc_lut_matrix(idx_or_packed, scales, u_svd, v_svd, lut, q_basis, bits=0, out_dtype=torch.bfloat16):
    R, C = u_svd.shape[0], v_svd.shape[1]
    tile = lut.shape[1]
    n_tiles = C // tile
    idx = unpack_sc_lut_indices(idx_or_packed, n_tiles, bits) if bits > 0 else idx_or_packed
    w_rot = lut.float()[idx.long()].reshape(R, scales.shape[1], -1) * scales.float().unsqueeze(-1)
    w_main = (w_rot.reshape(R, C // 64, 64) @ q_basis.float()).reshape(R, C)
    return (w_main + (u_svd.float() @ v_svd.float())).to(out_dtype)