import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
from scipy.fftpack import dct, idct
import time

def get_2d_dct(weight_matrix):
    """Applies a standard 2D DCT over the rows and columns."""
    d0 = dct(weight_matrix, axis=0, norm='ortho')
    d1 = dct(d0,            axis=1, norm='ortho')
    return d1

def decompress_2d_dct(dct_matrix, original_shape):
    """Applies a standard 2D Inverse DCT to restore the matrix."""
    i1 = idct(dct_matrix, axis=1, norm='ortho')
    i0 = idct(i1,         axis=0, norm='ortho')
    return i0

if __name__ == "__main__":
    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
    print(f"Loading weights from {model_id}...")
    
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    target_layer_name = "model.embed_tokens.weight"
    W = model.state_dict()[target_layer_name].cpu().numpy()
    orig_shape = W.shape  # (49152, 576)
    print(f"Original Shape: {orig_shape}")

    # 1. Transform the 2D matrix to 2D DCT domain once
    print("Transforming weight matrix into 2D DCT Domain...")
    full_dct = get_2d_dct(W)

    # 2. Extract "LITTLE" Params (Keep fewer rows, keep fewer or all cols)
    # Let's keep 4000 rows out of 49152, and 144 columns out of 576
    little_rows, little_cols = 4000, 144
    little_params = full_dct[:little_rows, :little_cols].copy()
    
    # 3. Extract "BIG" Params (Higher quality extension)
    # Let's expand up to 12000 rows and 288 columns
    big_rows, big_cols = 12000, 288
    big_params_delta = full_dct[:big_rows, :big_cols].copy()
    # Mask out the little parameters so the delta layer only stores unique high frequencies
    big_params_delta[:little_rows, :little_cols] = 0.0

    print(f"Little params footprint: {little_params.size} floats")
    print(f"Big delta footprint: {np.count_nonzero(big_params_delta)} floats")

    # --- SIMULATION 1: Decompress 'Little' Model ---
    dct_reconstruct_little = np.zeros(orig_shape)
    dct_reconstruct_little[:little_rows, :little_cols] = little_params
    W_little = decompress_2d_dct(dct_reconstruct_little, orig_shape)
    
    mse_little = np.mean((W - W_little) ** 2)
    print(f"\n[LITTLE MODEL] Reconstruction MSE: {mse_little:.2e}")

    # --- SIMULATION 2: Decompress 'Big' Model (Patching Little + Delta) ---
    dct_reconstruct_big = dct_reconstruct_little.copy()
    dct_reconstruct_big[:big_rows, :big_cols] += big_params_delta
    W_big = decompress_2d_dct(dct_reconstruct_big, orig_shape)
    
    mse_big = np.mean((W - W_big) ** 2)
    print(f"[BIG MODEL] Patched Reconstruction MSE: {mse_big:.2e}")
    print(f"Accuracy improvement factor: {mse_little / (mse_big + 1e-15):.1f}x closer to original")

    # 4. Update LLM with high-accuracy patched weights
    model.state_dict()[target_layer_name].copy_(torch.from_numpy(W_big).float())
    model.to(device)

    # 5. Run inference test
    messages = [{"role": "user", "content": "Explain what a 2D Discrete Cosine Transform does in one sentence."}]
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(device)
    input_ids = inputs["input_ids"] if isinstance(inputs, dict) or hasattr(inputs, "data") else inputs

    outputs = model.generate(input_ids, max_new_tokens=40, temperature=0.1, do_sample=True)
    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
    print("\nModel Output with 2D-DCT Weights:\n", response)