"""
Idea - Big little params (Fetch higher accuracy param method)

If the params in the smaller model is not enough for accuracy 
the model fetches higher accuracy params from disk

"""

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

def compress_4d(weight_matrix, shape_4d, keep_dims=(8, 8, 8, 8)):
    """
    Reshapes a 2D weight matrix into a 4D tensor and applies a 4D-DCT.
    Trims high-frequency coefficients down to `keep_dims`.
    """
    # 1. Reshape 2D weight matrix to the target 4D hypercube
    tensor_4d = weight_matrix.reshape(shape_4d)
    
    # 2. Sequential 4D Forward DCT
    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')
    
    # 3. Keep only the low-frequency parameters (Top-left corners of the hypercube)
    k0, k1, k2, k3 = keep_dims
    compressed_params = d3[:k0, :k1, :k2, :k3]
    return compressed_params

def decompress_4d(compressed_params, target_shape_4d, original_shape_2d):
    """
    Pads the 4D compressed parameters back with zeros and runs a 4D-IDCT.
    """
    k0, k1, k2, k3 = compressed_params.shape
    
    # 1. Pad back to the full 4D hypercube resolution with zeros
    padded_dct = np.zeros(target_shape_4d)
    padded_dct[:k0, :k1, :k2, :k3] = compressed_params
    
    # 2. Sequential 4D Inverse DCT
    i3 = idct(padded_dct, 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')
    
    # 3. Reshape back to the original 2D matrix
    return i0.reshape(original_shape_2d)

if __name__ == "__main__":
    # Load model weights (Using SmolLM2-135M as an ideal baseline open model)
    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
    print(f"Loading weights from {model_id}...")
    
    # We load onto CPU for numpy manipulation
    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"
    model.to(device)
    # Select a heavy internal projection layer to compress
    # For SmolLM2, intermediate_size = 5632, hidden_size = 576
    # Let's target an attention projection layer instead: self_attn.q_proj (576 x 576)
    #target_layer_name = "model.layers.0.self_attn.q_proj.weight"
    target_layer_name = "model.embed_tokens.weight"
    W = model.state_dict()[target_layer_name].numpy()
    #W = list(model.parameters())[0].detach().numpy()
    
    orig_shape = W.shape  # (576, 576)
    print(f"\nTarget Layer: {target_layer_name}")
    print(f"Original Shape: {orig_shape} ({W.size} parameters)")
    
    # Factorize (576, 576) into a 4D Shape: (24, 24, 24, 24)
    # Since 24 * 24 = 576
    #shape_4d = (24, 24, 24, 24)
    shape_4d = (192, 256, 24, 24)
    
    # Choose compression constraints (Hyper-volume truncation)
    # Keeping only 8x8x8x8 parameters out of 24x24x24x24
    keep_dims = (8*8, 8*8, 8, 8) 
    
    print(f"Reshaping to 4D Hypercube: {shape_4d}")
    print(f"Truncating 4D DCT box to: {keep_dims}")
    
    # --- Execute Compression Loop ---
    t0 = time.time()
    
    # Step 1: Base compression pass
    params = compress_4d(W, shape_4d, keep_dims=keep_dims)
    params_W = np.zeros_like(params)
    for _ in range(10):
        W_recon = decompress_4d(params, shape_4d, orig_shape)
        W_recon += decompress_4d(params_W, shape_4d, orig_shape)
        err = (W - W_recon)
        params_W += 1e3 * compress_4d(err, shape_4d, keep_dims=keep_dims)

        #params_W += 0.01 * error_params
        print(np.sum(err**2))

    # Final functionalized weight restoration
    
    #W_final = W_recon + decompress_4d(params, shape_4d, orig_shape)
    W_final = decompress_4d(params, shape_4d, orig_shape)
    W_final += decompress_4d(params_W, shape_4d, orig_shape) / 1e7
    
    duration = time.time() - t0
    
    # --- Evaluate Metric Success ---
    total_compressed_size = params.size# + error_params.size
    compression_ratio = W.size / total_compressed_size
    
    # Calculate Mean Squared Error & Mean Absolute Error
    mse = np.mean((W - W_final) ** 2)
    mae = np.mean(np.abs(W - W_final))
    
    print("\n--- Compression Results ---")
    print(f"Time Taken: {duration:.4f} seconds")
    print(f"Original Parameters: {W.size}")
    print(f"Compressed Parameter Footprint: {total_compressed_size}")
    print(f"Compression Ratio: {compression_ratio:.2f}x smaller")
    print(f"Reconstruction MSE Accuracy: {mse:.2e}")
    print(f"Reconstruction MAE Accuracy: {mae:.2e}")
    
    # Directly inject weights back to test functional model consistency
    W_final_tensor = torch.from_numpy(W_final).float()
    model.state_dict()[target_layer_name].copy_(W_final_tensor)
    print("\nSuccessfully updated LLM tensor with functional 4D-DCT data.")

    messages = [
        {"role": "user", "content": "Explain what a 4D Discrete Cosine Transform does in one sentence."}
    ]

    # --- Your previous tokenization step ---
    inputs = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to(device)

    # --- THE FIX ---
    # Check if inputs is a dictionary/BatchEncoding; grab the underlying tensor
    if isinstance(inputs, dict) or hasattr(inputs, "data"):
        input_ids = inputs["input_ids"]
    else:
        input_ids = inputs

    # Pass the tensor directly to generate
    outputs = model.generate(
        input_ids,  # Use input_ids tensor instead of the dictionary wrapper
        max_new_tokens=60, 
        temperature=0.2, 
        do_sample=True
    )

    # Use input_ids.shape[-1] to strip away the prompt when decoding
    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)

    print("Model Output:", response)