import torch
import numpy as np
import zlib
import pickle
from transformers import GPT2LMHeadModel, GPT2Config, AutoTokenizer

# -------------------------------
# Compression/decompression functions
# -------------------------------
def quantize_tensor(tensor, num_levels=256):
    min_val = tensor.min().item()
    max_val = tensor.max().item()
    if min_val == max_val:
        q = torch.zeros_like(tensor, dtype=torch.uint8)
        scale = 0.0
        return q, min_val, max_val, scale
    scale = (max_val - min_val) / (num_levels - 1)
    q = torch.round((tensor - min_val) / scale).clamp(0, num_levels-1).to(torch.uint8)
    return q, min_val, max_val, scale

def compress_weights(model, num_levels=256):
    compressed_data = {}
    total_original_bytes = 0
    total_compressed_bytes = 0
    seen_tensors = set()
    for name, param in model.named_parameters():
        if param.requires_grad and id(param) not in seen_tensors:
            seen_tensors.add(id(param))
            q, min_val, max_val, scale = quantize_tensor(param.detach().cpu(), num_levels)
            q_np = q.numpy().tobytes()
            comp = zlib.compress(q_np, level=9)
            compressed_data[name] = {
                'compressed': comp,
                'shape': param.shape,
                'min_val': min_val,
                'max_val': max_val,
                'num_levels': num_levels,
                'dtype': param.dtype
            }
            total_original_bytes += param.numel() * param.element_size()
            total_compressed_bytes += len(comp)
    ratio = total_compressed_bytes / total_original_bytes
    print(f"Compression ratio: {ratio:.4f}")
    return compressed_data

def decompress_state_dict(compressed_data):
    state_dict = {}
    for name, info in compressed_data.items():
        q_bytes = zlib.decompress(info['compressed'])
        q_np = np.frombuffer(q_bytes, dtype=np.uint8).reshape(info['shape']).copy()
        q = torch.from_numpy(q_np).float()
        min_val = info['min_val']
        max_val = info['max_val']
        num_levels = info['num_levels']
        if num_levels > 1:
            scale = (max_val - min_val) / (num_levels - 1)
            dequantized = min_val + scale * q
        else:
            dequantized = torch.full_like(q, min_val)
        state_dict[name] = dequantized.to(info['dtype'])

    # Handle tied weights: lm_head shares weights with wte in GPT-2
    if 'transformer.wte.weight' in state_dict:
        state_dict['lm_head.weight'] = state_dict['transformer.wte.weight'].clone()

    return state_dict

# -------------------------------
# 1. Load original GPT-2 LM model
# -------------------------------
model_name = "distilgpt2"   # or "gpt2" for the larger one
print(f"Loading {model_name}...")
model_orig = GPT2LMHeadModel.from_pretrained(model_name)
model_orig.eval()

# -------------------------------
# 2. Compress weights
# -------------------------------
print("\nCompressing weights...")
compressed = compress_weights(model_orig, num_levels=256)

# Save compressed dict to disk
compressed_file = f"{model_name}_compressed.pkl"
with open(compressed_file, "wb") as f:
    pickle.dump(compressed, f)
print(f"Compressed model saved to {compressed_file}")

# -------------------------------
# 3. Load compressed weights from disk and decompress
# -------------------------------
print("\nLoading compressed weights from disk...")
with open(compressed_file, "rb") as f:
    compressed_loaded = pickle.load(f)

print("Decompressing weights...")
decompressed_state_dict = decompress_state_dict(compressed_loaded)

# Create a fresh LM model with the same architecture
config = GPT2Config.from_pretrained(model_name)
model_decomp = GPT2LMHeadModel(config)
model_decomp.load_state_dict(decompressed_state_dict, strict=True)
model_decomp.eval()

# -------------------------------
# 4. Verify output similarity
# -------------------------------
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
input_text = "The future of artificial intelligence is"
input_ids = tokenizer.encode(input_text, return_tensors="pt")

# Forward pass through both models
with torch.no_grad():
    output_orig = model_orig(input_ids)
    output_decomp = model_decomp(input_ids)

logits_orig = output_orig.logits
logits_decomp = output_decomp.logits

mse = torch.nn.functional.mse_loss(logits_orig, logits_decomp).item()
cosine_sim = torch.nn.functional.cosine_similarity(logits_orig.flatten(), logits_decomp.flatten(), dim=0).item()
max_diff = (logits_orig - logits_decomp).abs().max().item()

print("\n--- Verification Results ---")
print(f"MSE between logits: {mse:.6e}")
print(f"Cosine similarity: {cosine_sim:.8f}")
print(f"Maximum absolute difference: {max_diff:.6f}")

# -------------------------------
# 5. Generate text from both models
# -------------------------------
print("\n--- Text Generation ---")
with torch.no_grad():
    gen_orig = model_orig.generate(input_ids, max_length=50, do_sample=False)
    gen_decomp = model_decomp.generate(input_ids, max_length=50, do_sample=False)

print(f"Original output:     {tokenizer.decode(gen_orig[0], skip_special_tokens=True)}")
print(f"Decompressed output: {tokenizer.decode(gen_decomp[0], skip_special_tokens=True)}")