# By Forge


import torch


def native_unpack_4x4bits_in_1x16bits_to_4x8bits_in_1x32bits(x):
    x = x.view(torch.uint8).view(x.size(0), -1)
    unpacked = torch.stack([x & 15, x >> 4], dim=-1)
    reshaped = unpacked.view(x.size(0), -1)
    reshaped = reshaped.view(torch.int8) - 8
    return reshaped.view(torch.int32)


def native_unpack_4x4bits_in_1x16bits_to_4x8bits_in_1x32bits_u(x):
    x = x.view(torch.uint8).view(x.size(0), -1)
    unpacked = torch.stack([x & 15, x >> 4], dim=-1)
    reshaped = unpacked.view(x.size(0), -1)
    return reshaped.view(torch.int32)


disable_all_optimizations = False

if not hasattr(torch, "uint16"):
    disable_all_optimizations = True

if disable_all_optimizations:
    print(
        "You are using PyTorch below version 2.3. Some optimizations will be disabled."
    )

if not disable_all_optimizations:
    native_4bits_lookup_table = (
        native_unpack_4x4bits_in_1x16bits_to_4x8bits_in_1x32bits(
            torch.arange(start=0, end=256 * 256, dtype=torch.long).to(torch.uint16)
        )[:, 0]
    )
    native_4bits_lookup_table_u = (
        native_unpack_4x4bits_in_1x16bits_to_4x8bits_in_1x32bits_u(
            torch.arange(start=0, end=256 * 256, dtype=torch.long).to(torch.uint16)
        )[:, 0]
    )


def quick_unpack_4bits(x):
    if disable_all_optimizations:
        return (
            torch.stack([x & 15, x >> 4], dim=-1).view(x.size(0), -1).view(torch.int8)
            - 8
        )

    global native_4bits_lookup_table

    s0 = x.size(0)
    x = x.view(torch.uint16)

    if native_4bits_lookup_table.device != x.device:
        native_4bits_lookup_table = native_4bits_lookup_table.to(device=x.device)

    y = torch.index_select(
        input=native_4bits_lookup_table, dim=0, index=x.to(dtype=torch.int32).flatten()
    )
    y = y.view(torch.int8)
    y = y.view(s0, -1)

    return y


def quick_unpack_4bits_u(x):
    if disable_all_optimizations:
        return torch.stack([x & 15, x >> 4], dim=-1).view(x.size(0), -1)

    global native_4bits_lookup_table_u

    s0 = x.size(0)
    x = x.view(torch.uint16)

    if native_4bits_lookup_table_u.device != x.device:
        native_4bits_lookup_table_u = native_4bits_lookup_table_u.to(device=x.device)

    y = torch.index_select(
        input=native_4bits_lookup_table_u,
        dim=0,
        index=x.to(dtype=torch.int32).flatten(),
    )
    y = y.view(torch.uint8)
    y = y.view(s0, -1)

    return y


def change_4bits_order(x):
    y = torch.stack([x & 15, x >> 4], dim=-2).view(x.size(0), -1)
    z = y[:, ::2] | (y[:, 1::2] << 4)
    return z
