import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np

class ProgressiveWeightLoader:
    def __init__(self, original_weight, base_rank=16, delta_rank=16, num_deltas=4):
        self.orig_shape = original_weight.shape
        
        # Compute SVD once to extract the principal components
        U, S, Vt = np.linalg.svd(original_weight, full_matrices=False)
        
        # 1. Extract Initial Core Component (C)
        self.C_U = U[:, :base_rank]
        self.C_S = S[:base_rank]
        self.C_Vt = Vt[:base_rank, :]
        
        # 2. Extract 'n' distinct delta sub-blocks (K_i)
        self.deltas = []
        for i in range(num_deltas):
            start_r = base_rank + (i * delta_rank)
            end_r = base_rank + ((i + 1) * delta_rank)
            
            # Store each K_i as a compact low-rank tuple to simulate disk chunks
            k_u = U[:, start_r:end_r]
            k_s = S[start_r:end_r]
            k_vt = Vt[start_r:end_r, :]
            
            self.deltas.append((k_u, k_s, k_vt))
            
    def get_base_core(self):
        """Returns the initial core matrix C."""
        return np.dot(self.C_U * self.C_S, self.C_Vt)
    
    def fetch_delta_patch(self, index):
        """Simulates loading a specific K_i sub-block from disk."""
        k_u, k_s, k_vt = self.deltas[index]
        return np.dot(k_u * k_s, k_vt)

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_orig = model.state_dict()[target_layer_name].cpu().numpy()
    
    # Initialize our system: Base rank 16, with 4 cascading updates of rank 16 each
    num_components = 4
    decomposer = ProgressiveWeightLoader(W_orig, base_rank=16, delta_rank=16, num_deltas=num_components)
    
    # --- STAGE 0: Initializing the Model with ONLY Base Core C ---
    W_current = decomposer.get_base_core()
    print(f"\n[Stage 0] Core C Loaded. Footprint: {decomposer.C_U.size + decomposer.C_Vt.size} floats")
    print(f"Base MSE: {np.mean((W_orig - W_current)**2):.2e}")
    
    # --- STAGE 1 to N: Progressively Stream Deltas from "Disk" ---
    for stage in range(num_components):
        # Fetch K_i and add it to our running weight matrix
        K_i = decomposer.fetch_delta_patch(stage)
        W_current += K_i
        
        current_mse = np.mean((W_orig - W_current)**2)
        print(f"[Stage {stage+1}] Fetched K_{stage+1}. Combined MSE: {current_mse:.2e}")
        
    # Inject the accumulated highly-accurate weights back into the LLM
    model.state_dict()[target_layer_name].copy_(torch.from_numpy(W_current).float())
    model.to(device)
    
    # Run a text generation test with the progressive weights
    messages = [{"role": "user", "content": "Write a three word poem."}]
    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=20, temperature=0.1, do_sample=True)
    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
    print("\nModel Output after loading C + all K updates:\n", response.strip())