import torch
import torch.nn as nn
import torch.nn.functional as F
import time

# -------------------------------------------------------------
# 1. ARCHITECTURE DEFINITION
# -------------------------------------------------------------
class HierarchicalThermalRelaxation(nn.Module):
    def __init__(self, channels=16):
        super().__init__()
        self.channels = channels
        self.local_kernel = nn.Conv2d(channels, channels, 3, padding=1, groups=channels, bias=False)
        self.global_kernel = nn.Conv2d(channels, channels, 7, padding=3, groups=channels, bias=False)
        
        with torch.no_grad():
            torch.nn.init.constant_(self.local_kernel.weight, 0.1)
            torch.nn.init.constant_(self.global_kernel.weight, 0.02)

    def forward(self, current_state, boundaries, boundary_mask, max_iters=2000, tol=5e-3):
        state = current_state.clone()
        state = torch.where(boundary_mask, boundaries, state)
        
        iterations = 0
        for i in range(max_iters):
            old_state = state.clone()
            diffused = 0.6 * self.local_kernel(state) + 0.4 * self.global_kernel(state)
            state = 0.3 * state + 0.7 * diffused
            state = torch.where(boundary_mask, boundaries, state)
            
            residual = torch.mean(torch.abs(state - old_state))
            iterations += 1
            if residual < tol:
                break
                
        return state, iterations

class SpatialProbingHead(nn.Module):
    def __init__(self, channels=16, vocab_size=8):
        super().__init__()
        # Projects 16 channels at a specific pixel location down to vocabulary logits
        self.projection = nn.Linear(channels, vocab_size)
        
    def forward(self, relaxed_field, trajectory_coords):
        """
        relaxed_field: Tensor of shape (1, channels, H, W)
        trajectory_coords: List of (x, y) tuples representing the reading path
        """
        logits_list = []
        for x, y in trajectory_coords:
            # Extract the 16-channel vector at this specific grid point
            latent_vector = relaxed_field[0, :, y, x] # Shape: (channels,)
            
            # Project to vocabulary space
            logits = self.projection(latent_vector) # Shape: (vocab_size,)
            logits_list.append(logits)
            
        return torch.stack(logits_list) # Shape: (sequence_length, vocab_size)

# -------------------------------------------------------------
# 2. RUNNING THE COMPLETE PIPELINE
# -------------------------------------------------------------
if __name__ == "__main__":
    GRID_SZ = 64
    CHANNELS = 16
    
    # Define our system's vocabulary
    vocab = {0: "<PAD>", 1: "The", 2: "AI", 3: "reaches", 4: "thermal", 5: "equilibrium", 6: "instantly", 7: "<EOS>"}
    VOCAB_SIZE = len(vocab)
    
    # Initialize our modules
    tcr_layer = HierarchicalThermalRelaxation(channels=CHANNELS)
    probing_head = SpatialProbingHead(channels= CHANNELS, vocab_size=VOCAB_SIZE)
    
    # Seed the weights deterministicly for demonstration
    with torch.no_grad():
        torch.nn.init.eye_(probing_head.projection.weight[:CHANNELS, :CHANNELS])

    # Setup the Canvas and Prompt Masks
    mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    mock_mask = torch.zeros(1, 1, GRID_SZ, GRID_SZ).bool()
    
    # Inject a "Hot Prompt" anchor influencing semantic distribution
    mock_boundaries[:, :, 10:20, 10:20] = 1.5 
    mock_mask[:, :, 10:20, 10:20] = True
    
    # Run cold start to settle the field
    cold_canvas = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    relaxed_field, _ = tcr_layer(cold_canvas, mock_boundaries, mock_mask)
    
    # Define a spatial "Reading Trajectory" sweeping across the diffused gradient from (10,10) to (50,50)
    # This replaces sequential token-by-token generation loops!
    reading_path = [(15 + i, 15 + i) for i in range(0, 35, 5)]
    
    # Extract text logits from the static equilibrium canvas in a single step
    output_logits = probing_head(relaxed_field, reading_path)
    predicted_token_ids = torch.argmax(output_logits, dim=-1)
    
    # Decode integers back to text strings
    generated_text = [vocab[token_id.item()] for token_id in predicted_token_ids]
    
    print("\n" + "="*50)
    print("STABILIZED LATENT FIELD TEXT DECODING")
    print("="*50)
    print(f"Sampling Path Coordinates: {reading_path}")
    print(f"Raw Token IDs Sampled:    {predicted_token_ids.tolist()}")
    print(f"Decoded Output Sequence:   {' '.join(generated_text)}")
    print("="*50)