import torch
import torch.nn as nn
import time

class ThermalContextRelaxation(nn.Module):
    def __init__(self, grid_size=64, channels=16):
        super().__init__()
        self.grid_size = grid_size
        self.channels = channels
        
        # Neural Diffusion Kernel: Learns how "logic" propagates to neighboring cells
        self.diffusion_kernel = nn.Conv2d(
            in_channels=channels, 
            out_channels=channels, 
            kernel_size=3, 
            padding=1, 
            groups=channels, # Depthwise convolution simulating localized diffusion
            bias=False
        )
        # Initialize kernel to act like a discrete Laplace operator (averaging neighbors)
        with torch.no_grad():
            laplacian = torch.tensor([[[[0.1, 0.2, 0.1],
                                        [0.2, 0.0, 0.2],
                                       [0.1, 0.2, 0.1]]]])
            self.diffusion_kernel.weight.copy_(laplacian.repeat(channels, 1, 1, 1))

    def forward(self, current_state, boundaries, boundary_mask, max_iters=200, tol=1e-4):
        """
        Relaxes the latent grid until it hits mathematical equilibrium.
        
        current_state: Tensor (1, C, H, W) -> The initial canvas (zeros for Cold Start, old state for Warm Start)
        boundaries: Tensor (1, C, H, W) -> The prompt values injected into the system
        boundary_mask: Tensor (1, 1, H, W) -> Binary mask indicating where the prompt is locked
        """
        state = current_state.clone()
        
        # Enforce boundary conditions initially
        state = torch.where(boundary_mask, boundaries, state)
        
        iterations = 0
        for i in range(max_iters):
            old_state = state.clone()
            
            # Step 1: Neural Diffusion (Heat propagation step)
            diffused = self.diffusion_kernel(state)
            
            # Step 2: Relaxation Update (Mixing old state with diffused state)
            state = 0.5 * state + 0.5 * diffused
            
            # Step 3: Re-enforce Boundary Conditions (The user's prompt cannot be overwritten)
            state = torch.where(boundary_mask, boundaries, state)
            
            # Step 4: Check for Equilibrium (Residual Error -> 0)
            residual = torch.norm(state - old_state)
            iterations += 1
            if residual < tol:
                break
                
        return state, iterations

# ==========================================
# RUNNING THE PROOF OF CONCEPT
# ==========================================
if __name__ == "__main__":
    # Setup hyperparameters
    GRID_SZ = 64
    CHANNELS = 16
    tcr_layer = ThermalContextRelaxation(grid_size=GRID_SZ, channels=CHANNELS)
    
    # Define a Mock Prompt Scenario
    # Let's inject a "Heat Source" at the top-left and a "Sink" at the bottom-right
    mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    mock_mask = torch.zeros(1, 1, GRID_SZ, GRID_SZ).bool()
    
    # Prompt A: Source at (10,10), Sink at (50,50)
    mock_boundaries[:, :, 10:15, 10:15] = 1.0  # Hot Prompt token
    mock_mask[:, :, 10:15, 10:15] = True
    
    mock_boundaries[:, :, 50:55, 50:55] = -1.0 # Cold Prompt token
    mock_mask[:, :, 50:55, 50:55] = True

    # ------------------------------------------
    # CASE 1: SOLVE FROM SCRATCH (COLD START)
    # ------------------------------------------
    cold_canvas = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    
    start_time = time.time()
    equilibrium_state_A, cold_iters = tcr_layer(cold_canvas, mock_boundaries, mock_mask)
    cold_time = (time.time() - start_time) * 1000 # ms

    print("--- 1. SOLVE FROM SCRATCH (COLD START) ---")
    print(f"Iterations to Equilibrium: {cold_iters}")
    print(f"Execution Time:            {cold_time:.2f} ms")

    # ------------------------------------------
    # CASE 2: SOLVE VIA BRIDGED STATE (WARM START)
    # ------------------------------------------
    # The user modifies the prompt slightly (adds a new localized prompt block near the center)
    new_boundaries = mock_boundaries.clone()
    new_mask = mock_mask.clone()
    
    new_boundaries[:, :, 30:35, 30:35] = 0.8 # New injection
    new_mask[:, :, 30:35, 30:35] = True
    
    # Crucial step: Instead of passing zeros, we seed it with 'equilibrium_state_A'
    start_time = time.time()
    equilibrium_state_B, warm_iters = tcr_layer(equilibrium_state_A, new_boundaries, new_mask)
    warm_time = (time.time() - start_time) * 1000 # ms

    print("\n--- 2. SOLVE VIA BRIDGED STATE (WARM/STATEFUL) ---")
    print(f"Iterations to Equilibrium: {warm_iters}")
    print(f"Execution Time:            {warm_time:.2f} ms")
    
    # Calculate Speedup
    speedup = cold_iters / warm_iters
    print(f"\nSpeedup Factor:            {speedup:.1f}x")