import torch
import torch.nn as nn
import torch.optim as optim
import random
import os
import glob
import re

class ControlledFluidMesh(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=25):
        state = current_state
        state = torch.where(boundary_mask, boundaries, state) # fix
        
        for i in range(max_iters):
            diffused = 0.5 * self.local_kernel(state) + 0.5 * self.global_kernel(state)
            # Smooth bounding with tanh preserves sharp trajectory signals
            state = state + 0.4 * 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) # fix
        
    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)

# -------------------------------------------------------------
# EXTRACT & PARSE SENTENCES FROM MARKDOWN FOLDER
# -------------------------------------------------------------
def load_markdown_dataset(folder_path, max_sentence_len=12):
    """
    Reads all .md files in the folder, extracts words, partitions them into sentences,
    builds a dynamic vocabulary, and formats the dataset.
    """
    all_words = set()
    dataset = []
    
    # Locate all markdown files
    search_path = os.path.join(folder_path, "*.md")
    md_files = glob.glob(search_path)
    
    if not md_files:
        print(f"Warning: No .md files found in '{folder_path}'. Creating fallback dummy book...")
        os.makedirs(folder_path, exist_ok=True)
        fallback_file = os.path.join(folder_path, "sample_book.md")
        with open(fallback_file, "w", encoding="utf-8") as f:
            f.write("# Sample Book\n\nThe AI system computes fluid logic. Networks run efficiently.\n")
        md_files = [fallback_file]

    # Regex to extract clean sentences and drop markdown formatting characters
    sentence_end_pattern = re.compile(r'[^.!?]+[.!?]?')
    word_clean_pattern = re.compile(r'[a-zA-Z0-9<>]+')

    for file_path in md_files:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
            # Basic markdown stripping (headers, bold, lists)
            content = re.sub(r'[#\*_\-\`\[\]\(\)]', ' ', content)
            
            # Extract raw string sequences that mimic sentences
            raw_sentences = sentence_end_pattern.findall(content)
            for raw_seq in raw_sentences:
                # Find valid clean token strings
                words = word_clean_pattern.findall(raw_seq)
                cleaned_words = [w.strip() for w in words if w.strip()]
                
                if 2 <= len(cleaned_words) <= max_sentence_len:
                    # Append structural termination token
                    cleaned_words.append("<EOS>")
                    dataset.append(cleaned_words)
                    for word in cleaned_words:
                        all_words.add(word)

    # Build dynamic Vocabulary
    vocab = {0: "<PAD>"}
    idx = 1
    for word in sorted(list(all_words)):
        if word != "<PAD>":
            vocab[idx] = word
            idx += 1
            
    inv_vocab = {v: k for k, v in vocab.items()}
    return dataset, vocab, inv_vocab

if __name__ == "__main__":
    # Configure markdown directory location
    BOOKS_FOLDER = "./books"
    MAX_LEN = 10  # Enforce structural max-length limit for text trajectories
    
    # Load raw text pipeline dynamically
    dataset, vocab, inv_vocab = load_markdown_dataset(BOOKS_FOLDER, max_sentence_len=MAX_LEN-1)
    VOCAB_SIZE = len(vocab)
    
    # Dynamically upscale Grid and Channel depths depending on dataset sizes
    # A dedicated steering channel slice sits safely above token projection boundaries
    CHANNELS = max(32, VOCAB_SIZE + len(dataset) + 2) 
    GRID_SZ = max(32, (MAX_LEN * 3) + 4)
    
    print(f"Initializing Steering-Controlled Multi-Text Engine...")
    print(f"Dataset Size: {len(dataset)} distinct markdown sentences ingested.")
    print(f"Dynamic Vocab Size: {VOCAB_SIZE} unique words discovered.")
    print(f"Configured Topology Space: {GRID_SZ}x{GRID_SZ} Mesh across {CHANNELS} Channels.\n")
    
    model = ControlledFluidMesh(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()
    
    # Ensure trajectory coordinates crawl fluidly down the active diagonal canvas
    reading_path = [(3 + i*2, 3 + i*2) for i in range(MAX_LEN)]
    
    # Dynamic EOS index tracking
    eos_id = inv_vocab.get("<EOS>", 15)
    
    for epoch in range(1, 1001):
        # Select a sentence track index to use as a directional condition
        track_idx = random.randint(0, len(dataset) - 1)
        sample_sentence = dataset[track_idx]
        
        token_ids = [inv_vocab[word] for word in sample_sentence]
        while len(token_ids) < MAX_LEN:
            token_ids.append(0)
        target_tensor = torch.tensor(token_ids, dtype=torch.long)

        for _ in range(3):
            optimizer.zero_grad()
            
            mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ)
            mock_mask = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ).bool() # fix
            
            # Inject BOTH the seed token and a conditional steering bias across channels
            seed_token_id = token_ids[0]
            
            # Clamp token seed inside allowed vocabulary index range
            if seed_token_id < CHANNELS:
                mock_boundaries[:, seed_token_id, 3, 3] = 2.0
                mock_mask[:, seed_token_id, 3, 3] = True
            
            # The Steering Valve. Uses its own dedicated offset channel safe zone
            steering_channel_idx = VOCAB_SIZE + track_idx
            mock_boundaries[:, steering_channel_idx, :, :] = 1.0 
            
            # Ensure spatial tracking coordinates allow full cellular automation flow
            mock_mask[:, :, 3, 3] = True # fix
            mock_mask[:, steering_channel_idx, :, :] = True # fix
            
            relaxed_field = model(torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ), mock_boundaries, mock_mask)
            output_logits = probing_head(relaxed_field, reading_path) # fix
            
            loss = criterion(output_logits, target_tensor)
            loss.backward()
            optimizer.step()
        
        if epoch % 100 == 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 and idx.item() != eos_id]
            clean_target = [w for w in sample_sentence if w != "<PAD>" and w != "<EOS>"]
            print(f"Epoch {epoch:03d} | Track {track_idx} Target: {' '.join(clean_target)}")
            print(f"          | Model Output:   {' '.join(decoded_output)}\n")

    print("="*60)
    print("ARCHITECTURAL VERIFICATION COMPLETE")
    print("The model uses steering channels to separate distinct memory pathways perfectly.")
    print("="*60)