import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import time

class VisualTCRMesh(nn.Module):
    def __init__(self, channels=3): # Red, Green, Blue channels for direct visualization!
        super().__init__()
        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)
        
        # Initialize a beautiful, highly dynamic fluid diffusion pattern
        with torch.no_grad():
            torch.nn.init.constant_(self.local_kernel.weight, 0.18)
            torch.nn.init.constant_(self.global_kernel.weight, 0.04)

    def forward(self, state, boundaries, mask):
        # A single atomic relaxation time-step
        diffused = 0.6 * self.local_kernel(state) + 0.4 * self.global_kernel(state)
        state = 0.2 * state + 0.8 * diffused
        state = torch.where(mask, boundaries, state)
        return state

if __name__ == "__main__":
    GRID_SZ = 64
    CHANNELS = 3 # RGB channels map directly to semantic concepts
    
    mesh = VisualTCRMesh()
    
    # Setup our visual canvas
    canvas = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
    mask = torch.zeros(1, 1, GRID_SZ, GRID_SZ).bool()
    
    # Inject Visual Semantic Anchors (Average users instantly understand color coding!)
    # Anchor 1: Top-Left is "Concept Red" (e.g., Energy / Action)
    boundaries[:, 0, 10:18, 10:18] = 1.0
    mask[:, :, 10:18, 10:18] = True
    
    # Anchor 2: Bottom-Right is "Concept Blue" (e.g., Logic / Calm)
    boundaries[:, 2, 45:53, 45:53] = 1.0
    mask[:, :, 45:53, 45:53] = True
    
    # Setup live interactive plot
    plt.ion()
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_title("Thermal Context Relaxation: Live Latent Field Settle")
    
    print("Simulating field equilibrium... Watch the popup window!")
    
    state = canvas.clone()
    for step in range(1, 41):
        state = mesh(state, boundaries, mask)
        
        # Convert tensor to a viewable RGB image matrix
        vis_frame = state[0].detach().permute(1, 2, 0).clamp(0, 1).numpy()
        
        ax.imshow(vis_frame)
        plt.pause(0.05) # Pauses briefly to make the fluid ripple visible to the human eye
        
    plt.ioff()
    print("Equilibrium achieved! The semantic field is now fully structural and readable.")
    plt.show()