import torch
import json
import os
import struct
import numpy as np
from transformers import AutoModelForCausalLM, AutoConfig
from collections import OrderedDict
from tqdm import tqdm


def quantize_tensor_symmetric(tensor, bits_per_param):
    """
    Symmetric linear quantization with zero-point preservation.

    Args:
        tensor: Input float tensor
        bits_per_param: Bit width for quantization (2,3,4,5,6,8,12,16,32)

    Returns:
        quantized_signed: Signed integer tensor
        scale: Quantization scale factor
        level: Maximum absolute level L = 2^(bits-1)-1
        dtype_name: String descriptor
    """
    if bits_per_param == 32:
        return tensor.to(torch.float32), None, None, "float32"
    elif bits_per_param == 16:
        return tensor.to(torch.float16), None, None, "float16"

    max_abs = tensor.abs().max().item()
    if max_abs == 0:
        max_abs = 1e-8

    L = 2 ** (bits_per_param - 1) - 1
    scale = max_abs / L
    quantized = (tensor / scale).round().clamp(-L, L)

    dtype_name = f"int{bits_per_param}_sym"
    storage_dtype = torch.int8 if bits_per_param <= 8 else torch.int16
    return quantized.to(storage_dtype), scale, L, dtype_name


def pack_2bit(values):
    """Pack 4 unsigned 2-bit values into 1 byte."""
    shape = values.shape
    flat = values.flatten().to(torch.uint8)
    pad = (4 - len(flat) % 4) % 4
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)])
    flat = flat.view(-1, 4)
    packed = ((flat[:, 0] & 0x03) |
              ((flat[:, 1] & 0x03) << 2) |
              ((flat[:, 2] & 0x03) << 4) |
              ((flat[:, 3] & 0x03) << 6))
    return packed, shape, pad


def pack_3bit(values):
    """Pack 8 unsigned 3-bit values into 3 bytes."""
    shape = values.shape
    flat = values.flatten().to(torch.uint8)
    pad = (8 - len(flat) % 8) % 8
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)])
    flat = flat.view(-1, 8)
    packed = torch.zeros(len(flat) * 3, dtype=torch.uint8)
    packed[0::3] = ((flat[:, 0] & 0x07) |
                    ((flat[:, 1] & 0x07) << 3) |
                    ((flat[:, 2] & 0x03) << 6))
    packed[1::3] = (((flat[:, 2] >> 2) & 0x01) |
                    ((flat[:, 3] & 0x07) << 1) |
                    ((flat[:, 4] & 0x07) << 4) |
                    ((flat[:, 5] & 0x01) << 7))
    packed[2::3] = (((flat[:, 5] >> 1) & 0x03) |
                    ((flat[:, 6] & 0x07) << 2) |
                    ((flat[:, 7] & 0x07) << 5))
    return packed, shape, pad


def pack_4bit(values):
    """Pack 2 unsigned 4-bit values into 1 byte."""
    shape = values.shape
    flat = values.flatten().to(torch.uint8)
    pad = (2 - len(flat) % 2) % 2
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)])
    flat = flat.view(-1, 2)
    packed = ((flat[:, 0] & 0x0F) |
              ((flat[:, 1] & 0x0F) << 4))
    return packed, shape, pad


def pack_5bit(values):
    """Pack 8 unsigned 5-bit values into 5 bytes."""
    shape = values.shape
    flat = values.flatten().to(torch.uint8)
    pad = (8 - len(flat) % 8) % 8
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)])
    flat = flat.view(-1, 8)
    packed = torch.zeros(len(flat) * 5, dtype=torch.uint8)
    packed[0::5] = ((flat[:, 0] & 0x1F) |
                    ((flat[:, 1] & 0x07) << 5))
    packed[1::5] = (((flat[:, 1] >> 3) & 0x03) |
                    ((flat[:, 2] & 0x1F) << 2) |
                    ((flat[:, 3] & 0x01) << 7))
    packed[2::5] = (((flat[:, 3] >> 1) & 0x0F) |
                    ((flat[:, 4] & 0x0F) << 4))
    packed[3::5] = (((flat[:, 4] >> 4) & 0x01) |
                    ((flat[:, 5] & 0x1F) << 1) |
                    ((flat[:, 6] & 0x03) << 6))
    packed[4::5] = (((flat[:, 6] >> 2) & 0x07) |
                    ((flat[:, 7] & 0x1F) << 3))
    return packed, shape, pad


def pack_6bit(values):
    """Pack 4 unsigned 6-bit values into 3 bytes."""
    shape = values.shape
    flat = values.flatten().to(torch.uint8)
    pad = (4 - len(flat) % 4) % 4
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8)])
    flat = flat.view(-1, 4)
    packed = torch.zeros(len(flat) * 3, dtype=torch.uint8)
    packed[0::3] = ((flat[:, 0] & 0x3F) |
                    ((flat[:, 1] & 0x03) << 6))
    packed[1::3] = (((flat[:, 1] >> 2) & 0x0F) |
                    ((flat[:, 2] & 0x0F) << 4))
    packed[2::3] = (((flat[:, 2] >> 4) & 0x03) |
                    ((flat[:, 3] & 0x3F) << 2))
    return packed, shape, pad


def pack_12bit(values):
    """Pack 2 unsigned 12-bit values into 3 bytes."""
    shape = values.shape
    flat = values.flatten().to(torch.uint16)
    pad = (2 - len(flat) % 2) % 2
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint16)])
    flat = flat.view(-1, 2)
    a = flat[:, 0].numpy().astype(np.uint16)
    b = flat[:, 1].numpy().astype(np.uint16)
    byte0 = a & 0xFF
    byte1 = ((a >> 8) & 0x0F) | ((b & 0x0F) << 4)
    byte2 = (b >> 4) & 0xFF
    packed = np.stack([byte0, byte1, byte2], axis=1).flatten().astype(np.uint8)
    return torch.from_numpy(packed), shape, pad


def convert_to_dq(model_path, output_path, bits_per_param=8):
    """
    Convert HuggingFace model to .dq quantized format.

    Args:
        model_path: Path to local HuggingFace model directory
        output_path: Output path for .dq file
        bits_per_param: Target bit width (2, 3, 4, 5, 6, 8, 12, 16, 32)

    Format specification:
        Magic: 'DQMD' (4 bytes)
        Version: uint32 (4 bytes)
        Global metadata: length (uint32) + JSON bytes
        Tensor headers: length (uint32) + JSON bytes
        Tensor data: raw bytes per tensor in header order
    """
    config = AutoConfig.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True
    )

    state_dict = model.state_dict()

    dq_data = {
        'config': config.to_dict(),
        'bits_per_param': bits_per_param,
        'tensors': OrderedDict(),
        'metadata': {
            'model_path': model_path,
            'architecture': config.architectures[0] if hasattr(config, 'architectures') else 'unknown'
        }
    }

    pack_functions = {
        2: pack_2bit,
        3: pack_3bit,
        4: pack_4bit,
        5: pack_5bit,
        6: pack_6bit,
        12: pack_12bit,
    }

    total_params = 0

    with tqdm(total=len(state_dict), desc="Quantizing") as pbar:
        for name, tensor in state_dict.items():
            total_params += tensor.numel()

            quantized_signed, scale, level, dtype_name = quantize_tensor_symmetric(
                tensor, bits_per_param
            )

            tensor_info = {
                'shape': list(tensor.shape),
                'dtype': dtype_name,
                'scale': scale,
                'level': level,
            }

            if bits_per_param in pack_functions:
                unsigned = quantized_signed.to(torch.int32) + level
                if bits_per_param <= 8:
                    unsigned = unsigned.to(torch.uint8)
                else:
                    unsigned = unsigned.to(torch.uint16)

                packed, original_shape, pad_len = pack_functions[bits_per_param](unsigned)
                tensor_info['data'] = packed.numpy().tobytes()
                tensor_info['packed'] = True
                tensor_info['packed_shape'] = list(original_shape)
                tensor_info['pad'] = pad_len
                del unsigned, packed
            elif bits_per_param == 8:
                unsigned = (quantized_signed.to(torch.int32) + level).to(torch.uint8)
                tensor_info['data'] = unsigned.numpy().tobytes()
                tensor_info['packed'] = False
                del unsigned
            elif bits_per_param in (16, 32):
                tensor_info['data'] = quantized_signed.numpy().tobytes()
                tensor_info['packed'] = False

            dq_data['tensors'][name] = tensor_info
            del tensor, quantized_signed
            pbar.update(1)

    serializable_tensors = OrderedDict()
    for name, info in dq_data['tensors'].items():
        serializable_tensors[name] = {k: v for k, v in info.items() if k != 'data'}
        serializable_tensors[name]['data_size'] = len(info['data'])

    with open(output_path, 'wb') as f:
        f.write(b'DQMD')
        f.write(struct.pack('<I', 1))

        global_meta = {
            'config': dq_data['config'],
            'bits_per_param': dq_data['bits_per_param'],
            'metadata': dq_data['metadata']
        }
        global_meta_json = json.dumps(global_meta).encode('utf-8')
        f.write(struct.pack('<I', len(global_meta_json)))
        f.write(global_meta_json)

        tensor_headers_json = json.dumps(serializable_tensors).encode('utf-8')
        f.write(struct.pack('<I', len(tensor_headers_json)))
        f.write(tensor_headers_json)

        for info in dq_data['tensors'].values():
            f.write(info['data'])

    original_size_mb = total_params * 4 / 1024 ** 2
    compressed_size_mb = os.path.getsize(output_path) / 1024 ** 2
    ratio = (total_params * 4) / os.path.getsize(output_path)

    return {
        'total_params': total_params,
        'bits_per_param': bits_per_param,
        'original_size_mb': original_size_mb,
        'compressed_size_mb': compressed_size_mb,
        'compression_ratio': ratio,
        'output_path': output_path
    }


if __name__ == "__main__":
    result = convert_to_dq("modelHF", "modelDQ.dq", bits_per_param=8)
    print(f"Conversion complete: {result['compression_ratio']:.2f}x compression "
          f"({result['original_size_mb']:.1f}MB -> {result['compressed_size_mb']:.1f}MB)")