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

def get_full_dct_hypercube(weight_matrix, shape_4d):
    """Transforms the entire weight matrix into the 4D DCT domain."""
    tensor_4d = weight_matrix.reshape(shape_4d)
    d0 = dct(tensor_4d, axis=0, norm='ortho')
    d1 = dct(d0,        axis=1, norm='ortho')
    d2 = dct(d1,        axis=2, norm='ortho')
    d3 = dct(d2,        axis=3, norm='ortho')
    return d3

def decompress_from_dct(dct_cube, original_shape_2d):
    """Transforms a full 4D DCT domain tensor back to the original weight matrix."""
    i3 = idct(dct_cube, axis=3, norm='ortho')
    i2 = idct(i3,       axis=2, norm='ortho')
    i1 = idct(i2,       axis=1, norm='ortho')
    i0 = idct(i1,       axis=0, norm='ortho')
    return i0.reshape(original_shape_2d)

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  
    
    # SmolLM2 Embedding layer shape: (49152, 576)
    # 49152 factors to: 192 * 256
    # 576 factors to: 24 * 24
    shape_4d = (192, 256, 24, 24)
    print(f"Original Shape: {orig_shape} | 4D Hypercube: {shape_4d}")

    # 1. Map the entire matrix to 4D DCT frequency space once
    print("Transforming weight matrix into 4D DCT Domain...")
    full_dct_cube = get_full_dct_hypercube(W, shape_4d)

    # 2. Extract "LITTLE" Params (Ultra compressed base model)
    # Keeping a compact frequency zone
    little_dims = (16, 16, 6, 6) 
    little_params = full_dct_cube[:little_dims[0], :little_dims[1], :little_dims[2], :little_dims[3]].copy()
    
    # 3. Simulate "BIG" Params (Higher accuracy delta stored on disk)
    # When accuracy is low, we fetch a larger chunk of the spectrum
    big_dims = (48, 48, 12, 12)
    big_params_delta = full_dct_cube[:big_dims[0], :big_dims[1], :big_dims[2], :big_dims[3]].copy()
    # Mask out the little params so it only contains the high-frequency delta
    big_params_delta[:little_dims[0], :little_dims[1], :little_dims[2], :little_dims[3]] = 0.0

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

    # --- SIMULATION 1: Run the 'Little' Model (Low accuracy baseline) ---
    dct_reconstruct_little = np.zeros(shape_4d)
    dct_reconstruct_little[:little_dims[0], :little_dims[1], :little_dims[2], :little_dims[3]] = little_params
    W_little = decompress_from_dct(dct_reconstruct_little, orig_shape)
    
    mse_little = np.mean((W - W_little) ** 2)
    print(f"\n[LITTLE MODEL] Base Reconstruction MSE: {mse_little:.2e}")

    # --- SIMULATION 2: Fetch 'Big' Params from Disk & Patch ---
    print("Fetching higher accuracy parameter chunks from disk...")
    
    # We combine them directly in the frequency domain safely without errors
    dct_reconstruct_big = dct_reconstruct_little.copy()
    dct_reconstruct_big[:big_dims[0], :big_dims[1], :big_dims[2], :big_dims[3]] += big_params_delta
    W_big = decompress_from_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 better!")

    # 4. Inject the high-accuracy version back into the LLM
    model.state_dict()[target_layer_name].copy_(torch.from_numpy(W_big).float())
    model.to(device)

    # 5. Run inference to test readability
    messages = [{"role": "user", "content": "Explain what a 4D 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 Functional Weights:\n", response)