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

class HierarchicalThermalRelaxation(nn.Module):
    def __init__(self, channels=16):
        super().__init__()
        self.channels = channels
        
        # Multi-scale kernels to capture both local detail and global context in ONE step
        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():
            # Initialize with strong, stable diffusion weights
            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()
            
            # Combine local and wide-range diffusion to move "information" rapidly across the grid
            diffused = 0.6 * self.local_kernel(state) + 0.4 * self.global_kernel(state)
            
            # Succession Over-Relaxation update step
            state = 0.3 * state + 0.7 * diffused
            
            # Enforce immutable user prompts
            state = torch.where(boundary_mask, boundaries, state)
            
            # Use Mean Absolute Error to measure stability
            residual = torch.mean(torch.abs(state - old_state))
            iterations += 1
            
            if residual < tol:
                break
                
        return state, iterations

# ==========================================
# TEST EXECUTION
# ==========================================
if __name__ == "__main__":
    GRID_SZ = 64
    CHANNELS = 16
    tcr_layer = HierarchicalThermalRelaxation(channels=CHANNELS)
    
    # Base boundaries (Prompt A)
    mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    mock_mask = torch.zeros(1, 1, GRID_SZ, GRID_SZ).bool()
    
    mock_boundaries[:, :, 10:20, 10:20] = 1.0  # Hot prompt anchor
    mock_mask[:, :, 10:20, 10:20] = True
    mock_boundaries[:, :, 45:55, 45:55] = -1.0 # Cold prompt anchor
    mock_mask[:, :, 45:55, 45:55] = True

    # ------------------------------------------
    # 1. COLD START (From Zeros)
    # ------------------------------------------
    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")

    # ------------------------------------------
    # 2. BRIDGED STATE (Warm Start from State A)
    # ------------------------------------------
    # Add a small localized secondary update (like an incremental prompt modification)
    new_boundaries = mock_boundaries.clone()
    new_mask = mock_mask.clone()
    new_boundaries[:, :, 30:34, 30:34] = 0.8 
    new_mask[:, :, 30:34, 30:34] = True
    
    start_time = time.time()
    # Passing 'equilibrium_state_A' bridges the state!
    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 True Compute Efficiency Speedup
    iter_speedup = cold_iters / warm_iters
    time_speedup = cold_time / warm_time
    print("\n" + "="*40)
    print(f"Iteration Speedup Factor:  {iter_speedup:.1f}x")
    print(f"Wall-Clock Speedup Factor:  {time_speedup:.1f}x")
    print("="*40)