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

# -------------------------------------------------------------
# 1. RESIDUAL THERMAL RELAXATION CORE
# -------------------------------------------------------------
class TextFluidMesh(nn.Module):
    def __init__(self, channels=32):
        super().__init__()
        self.channels = channels
        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)
        
        nn.init.kaiming_normal_(self.local_kernel.weight)
        nn.init.kaiming_normal_(self.global_kernel.weight)

    def forward(self, current_state, boundaries, boundary_mask, max_iters=20):
        state = current_state
        state = torch.where(boundary_mask, boundaries, state)
        
        for i in range(max_iters):
            diffused = 0.5 * self.local_kernel(state) + 0.5 * self.global_kernel(state)
            # Residual connection to preserve gradient flow over text sequences
            state = state + 0.3 * torch.tanh(diffused)
            state = torch.where(boundary_mask, boundaries, state)
                
        return state

class SpatialProbingHead(nn.Module):
    def __init__(self, channels=32, vocab_size=16):
        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. TRAINING ENVIRONMENT & TEXT DATASET
# -------------------------------------------------------------
if __name__ == "__main__":
    GRID_SZ = 32
    CHANNELS = 32 # Expanded channel space to hold a diverse vocabulary
    
    # 16-word vocabulary corpus
    vocab = {
        0: "<PAD>", 1: "The", 2: "AI", 3: "reaches", 4: "equilibrium", 
        5: "system", 6: "computes", 7: "logic", 8: "instantly", 9: "beautifully",
        10: "running", 11: "on", 12: "fluid", 13: "networks", 14: "efficiently", 15: "<EOS>"
    }
    inv_vocab = {v: k for k, v in vocab.items()}
    VOCAB_SIZE = len(vocab)
    
    # Dataset containing varied conceptual structures
    dataset = [
        ["The", "AI", "reaches", "equilibrium", "instantly", "<EOS>"],
        ["The", "system", "computes", "logic", "beautifully", "<EOS>"],
        ["AI", "networks", "running", "on", "fluid", "logic", "<EOS>"],
        ["The", "fluid", "system", "computes", "efficiently", "<EOS>"]
    ]
    
    # Instantiate
    model = TextFluidMesh(channels=CHANNELS)
    probing_head = SpatialProbingHead(channels=CHANNELS, vocab_size=VOCAB_SIZE)
    
    optimizer = optim.Adam(list(model.parameters()) + list(probing_head.parameters()), lr=0.01)
    criterion = nn.CrossEntropyLoss()
    
    # Layout a reading trajectory path matching the maximum sentence size
    max_len = 7
    reading_path = [(4 + i*4, 4 + i*4) for i in range(max_len)]
    
    print("Initializing Multi-Text Training Engine...")
    print(f"Dataset Size: {len(dataset)} distinct text paths.\n")
    
    for epoch in range(1, 401):
        # Pick a random sample sentence from our training dataset each iteration
        sample_sentence = random.choice(dataset)
        
        # Convert words to integer tokens and pad out to uniform length
        token_ids = [inv_vocab[word] for word in sample_sentence]
        while len(token_ids) < max_len:
            token_ids.append(0) # Pad
        target_tensor = torch.tensor(token_ids, dtype=torch.long)
        
        optimizer.zero_grad()
        
        # Setup clean boundary canvas
        mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
        mock_mask = torch.zeros(1, 1, GRID_SZ, GRID_SZ).bool()
        
        # Inject the FIRST token as the absolute physical seed anchor point
        seed_token_id = token_ids[0]
        mock_boundaries[:, seed_token_id, 4, 4] = 2.0
        mock_mask[:, :, 4, 4] = True
        
        # Let the field relax based on the starting seed word
        cold_canvas = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
        relaxed_field = model(cold_canvas, mock_boundaries, mock_mask)
        
        # Sample the whole path simultaneously
        output_logits = probing_head(relaxed_field, reading_path)
        
        # Compute loss and optimize the fluid weights
        loss = criterion(output_logits, target_tensor)
        loss.backward()
        optimizer.step()
        
        # Monitor convergence properties over time
        if epoch % 50 == 0 or epoch == 1:
            predicted_ids = torch.argmax(output_logits, dim=-1)
            decoded_output = [vocab[idx.item()] for idx in predicted_ids if idx.item() != 0]
            clean_target = [w for w in sample_sentence if w != "<PAD>"]
            print(f"Epoch {epoch:03d} | Target: {' '.join(clean_target)}")
            print(f"          | Output: {' '.join(decoded_output)}\n")

    print("="*60)
    print("TRAINING ENGINE PERFORMANCE COMPLETE")
    print("The weights have learned how to structure fields for distinct sentences.")
    print("="*60)