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

# -------------------------------
# Your compression functions (slightly adapted for torch tensors)
# -------------------------------
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
    for name, param in model.named_parameters():
        if param.requires_grad:
            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_weights(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'])
        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'])
    return state_dict

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

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

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

# -------------------------------
# 3. Decompress into a new model
# -------------------------------
print("\nDecompressing weights...")
state_dict_decomp = decompress_weights(compressed)

# Create a fresh model with the same architecture
model_decomp = GPT2Model(GPT2Config.from_pretrained(model_name))
model_decomp.load_state_dict(state_dict_decomp, strict=True)
model_decomp.eval()

# -------------------------------
# 4. Verify output similarity
# -------------------------------
# Use a sample text
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
input_text = "The future of artificial intelligence is"
inputs = tokenizer(input_text, return_tensors="pt")

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

# Compare last hidden states (or logits if using GPT2LMHeadModel)
# For GPT2Model, the output is a BaseModelOutput with 'last_hidden_state'
hidden_orig = output_orig.last_hidden_state
hidden_decomp = output_decomp.last_hidden_state

# Compute metrics
mse = torch.nn.functional.mse_loss(hidden_orig, hidden_decomp).item()
cosine_sim = torch.nn.functional.cosine_similarity(hidden_orig.flatten(), hidden_decomp.flatten(), dim=0).item()
max_diff = (hidden_orig - hidden_decomp).abs().max().item()

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

# Optional: check that the model still generates reasonable text
# (requires GPT2LMHeadModel, but same principle)