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

# --- 1D DCT Operations ---
def compress_1d(weight_matrix):
    # Transform only across the vocabulary axis (axis=0)
    return dct(weight_matrix, axis=0, norm='ortho')

def decompress_1d(dct_matrix):
    return idct(dct_matrix, axis=0, norm='ortho')

if __name__ == "__main__":
    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
    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)

    # Transform vocabulary space to frequency space
    full_dct = compress_1d(W)

    # LITTLE MODEL: Keep the low-frequency global token trends (e.g., top 8000 frequencies)
    little_rows = 8000
    little_params = full_dct[:little_rows, :].copy()
    
    # BIG MODEL Extension: Fetch a much wider spectrum from disk
    big_rows = 50000
    big_params_delta = full_dct[:big_rows, :].copy()
    big_params_delta[:little_rows, :] = 0.0  # Zero out little to keep delta pure

    # --- Reconstruction ---
    dct_rec_big = np.zeros(orig_shape)
    dct_rec_big[:little_rows, :] = little_params
    dct_rec_big[:big_rows, :] += big_params_delta
    W_big = decompress_1d(dct_rec_big)
    
    # Inject and run inference
    model.state_dict()[target_layer_name].copy_(torch.from_numpy(W_big).float())
    model.to(device)

    messages = [{"role": "user", "content": "Explain what a 1D DCT 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)
    print("\nModel Output (1D DCT):\n", tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True))