import torch
import torch.nn as nn
import torch.optim as optim

# -------------------------------------------------------------
# 1. THE LEARNABLE TCR PARADIGM
# -------------------------------------------------------------
class LearnableThermalRelaxation(nn.Module):
    def __init__(self, channels=16):
        super().__init__()
        self.channels = channels
        
        # Convolutions with a slightly larger local neighborhood to speed up training propagation
        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 weights with slightly higher variance to kickstart gradient flow
        nn.init.uniform_(self.local_kernel.weight, -0.2, 0.2)
        nn.init.uniform_(self.global_kernel.weight, -0.1, 0.1)

    def forward(self, current_state, boundaries, boundary_mask, max_iters=25):
        state = current_state
        state = torch.where(boundary_mask, boundaries, state)
        
        for i in range(max_iters):
            # Combined neural diffusion step
            diffused = 0.5 * self.local_kernel(state) + 0.5 * self.global_kernel(state)
            state = 0.4 * state + 0.6 * diffused
            
            # Re-enforce immutable prompt constraints
            state = torch.where(boundary_mask, boundaries, state)
                
        return state

class SpatialProbingHead(nn.Module):
    def __init__(self, channels=16, vocab_size=8):
        super().__init__()
        self.projection = nn.Linear(channels, vocab_size)
        
    def forward(self, relaxed_field, trajectory_coords):
        logits_list = []
        for x, y in trajectory_coords:
            latent_vector = relaxed_field[0, :, y, x]
            logits = self.projection(latent_vector)
            logits_list.append(logits)
        return torch.stack(logits_list)

# -------------------------------------------------------------
# 2. RUNNING THE OPTIMIZED TRAINING LOOP
# -------------------------------------------------------------
if __name__ == "__main__":
    GRID_SZ = 32 
    CHANNELS = 16
    VOCAB_SIZE = 8
    
    vocab = {0: "<PAD>", 1: "The", 2: "AI", 3: "reaches", 4: "thermal", 5: "equilibrium", 6: "instantly", 7: "<EOS>"}
    
    model = LearnableThermalRelaxation(channels=CHANNELS)
    probing_head = SpatialProbingHead(channels=CHANNELS, vocab_size=VOCAB_SIZE)
    
    # FIX 1: Boost learning rate to 0.03 to help gradients push past the stall point
    optimizer = optim.Adam(list(model.parameters()) + list(probing_head.parameters()), lr=0.03)
    criterion = nn.CrossEntropyLoss()
    
    # Path coordinates
    reading_path = [(5 + i, 5 + i) for i in range(0, 21, 3)] 
    target_tokens = torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.long)
    
    print("Beginning Optimized Optimization Loop...")
    print("Goal: Force gradients completely through the diffusion trajectory.\n")
    
    # FIX 2: Extend to 200 epochs to give the spatial waves time to settle out completely
    for epoch in range(1, 201):
        optimizer.zero_grad()
        
        mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
        mock_mask = torch.zeros(1, 1, GRID_SZ, GRID_SZ).bool()
        
        # Seed the source point with a broad activation signature across the initial channels
        mock_boundaries[:, :8, 5, 5] = 2.0 
        mock_mask[:, :, 5, 5] = True
        
        relaxed_field = model(torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ), mock_boundaries, mock_mask)
        output_logits = probing_head(relaxed_field, reading_path)
        
        loss = criterion(output_logits, target_tokens)
        loss.backward()
        optimizer.step()
        
        if epoch % 20 == 0 or epoch == 1:
            predicted_ids = torch.argmax(output_logits, dim=-1)
            decoded_text = [vocab[idx.item()] for idx in predicted_ids]
            print(f"Epoch {epoch:03d} | Loss: {loss.item():.4f} | Output: {' '.join(decoded_text)}")

    print("\n" + "="*60)
    print("TRAINING COMPLETE: CONVERGED FIELD PERFORMANCE")
    print("="*60)
    print(f"Final Spatial Sequence: {' '.join(decoded_text)}")
    print("="*60)