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, 
            bias=False
        )
        # Using a slightly more aggressive diffusion stencil (stronger neighbor weights)
        with torch.no_grad():
            laplacian = torch.tensor([[[[0.15, 0.2, 0.15],
                                        [0.2,  0.0, 0.2],
                                       [0.15, 0.2, 0.15]]]])
            self.diffusion_kernel.weight.copy_(laplacian.repeat(channels, 1, 1, 1))

    def forward(self, current_state, boundaries, boundary_mask, max_iters=1000, tol=1e-3):
        """
        Relaxes the latent grid until it hits mathematical equilibrium.
        Note: Boosted max_iters ceiling and adjusted tol to prevent premature timeouts.
        """
        state = current_state.clone()
        state = torch.where(boundary_mask, boundaries, state)
        
        iterations = 0
        for i in range(max_iters):
            old_state = state.clone()
            
            # Step 1: Neural Diffusion
            diffused = self.diffusion_kernel(state)
            
            # Step 2: Relaxation Update (Over-relaxation factor to accelerate convergence)
            # Using 0.4 / 0.6 split forces faster propagation
            state = 0.4 * state + 0.6 * diffused
            
            # Step 3: Re-enforce Boundary Conditions
            state = torch.where(boundary_mask, boundaries, state)
            
            # Step 4: Check for Equilibrium using Mean Absolute Error (MAE) for stability
            residual = torch.mean(torch.abs(state - old_state))
            iterations += 1
            if residual < tol:
                break
                
        return state, iterations

# ==========================================
# RUNNING THE PROOF OF CONCEPT
# ==========================================
if __name__ == "__main__":
    GRID_SZ = 64
    CHANNELS = 16
    tcr_layer = ThermalContextRelaxation(grid_size=GRID_SZ, channels=CHANNELS)
    
    # Define a Mock Prompt Scenario
    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:20, 10:20] = 1.0  # Slightly larger brush for better gradient
    mock_mask[:, :, 10:20, 10:20] = True
    
    mock_boundaries[:, :, 45:55, 45:55] = -1.0 
    mock_mask[:, :, 45:55, 45: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 

    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 instruction in the middle)
    new_boundaries = mock_boundaries.clone()
    new_mask = mock_mask.clone()
    
    new_boundaries[:, :, 30:35, 30:35] = 0.8 
    new_mask[:, :, 30:35, 30:35] = True
    
    # Seeding it with the pre-converged '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 

    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 based on step mechanics
    speedup = cold_iters / warm_iters
    print(f"\nSpeedup Factor (Compute Efficiency): {speedup:.1f}x")