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, coords_y, coords_x):
        # Vectorized gather: pull every trajectory latent at once instead of
        # looping in Python. Shape: (num_coords, channels) -> single matmul.
        latents = relaxed_field[0, :, coords_y, coords_x].t()
        return self.projection(latents)

# -------------------------------------------------------------
# EXTRACT & PARSE SENTENCES FROM MARKDOWN FOLDER
# -------------------------------------------------------------
def load_markdown_dataset(folder_path, fixed_len=8):
    """
    Reads all .md files in the folder, extracts words, keeps only sentences with
    EXACTLY `fixed_len` words, 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 across the wide mesh today.\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()]

                # Keep ONLY sentences of exactly `fixed_len` words
                if len(cleaned_words) == fixed_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"

    # ---- Constant sentence length ----
    FIXED_WORDS = 5                 # every sentence is exactly this many words
    MAX_LEN = FIXED_WORDS + 1       # 8 words + <EOS> -> one trajectory slot per token

    # Run on GPU if one is available (no-op on CPU-only machines).
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load raw text pipeline dynamically
    dataset, vocab, inv_vocab = load_markdown_dataset(BOOKS_FOLDER, fixed_len=FIXED_WORDS)
    if len(dataset) == 0:
        raise SystemExit(
            f"No {FIXED_WORDS}-word sentences found in '{BOOKS_FOLDER}'. "
            f"Add some {FIXED_WORDS}-word sentences or change FIXED_WORDS."
        )
    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
    # NOTE: channel count grows with len(dataset). This is the single biggest cost
    # driver (every conv is linear in channels).
    CHANNELS = max(32, VOCAB_SIZE + len(dataset) + 2)
    GRID_SZ = max(32, (MAX_LEN * 3) + 4)

    print(f"Initializing Steering-Controlled Multi-Text Engine on {device}...")
    print(f"Sentence Length: fixed at {FIXED_WORDS} words (+<EOS> = {MAX_LEN} tokens).")
    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).to(device)
    probing_head = SpatialProbingHead(channels=CHANNELS, vocab_size=VOCAB_SIZE).to(device)

    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.
    # Precompute the index tensors once (they never change).
    reading_path = [(3 + i*2, 3 + i*2) for i in range(MAX_LEN)]
    coords_x = torch.tensor([x for x, y in reading_path], dtype=torch.long, device=device)
    coords_y = torch.tensor([y for x, y in reading_path], dtype=torch.long, device=device)

    # Dynamic EOS index tracking
    eos_id = inv_vocab.get("<EOS>", 15)

    # Preallocate the persistent buffers once and reuse them every step.
    init_state = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ, device=device)
    mock_boundaries = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ, device=device)
    mock_mask = torch.zeros(1, CHANNELS, GRID_SZ, GRID_SZ, dtype=torch.bool, device=device)

    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, device=device)

        # ---- Build boundaries/mask ONCE per epoch ----
        mock_boundaries.zero_()
        mock_mask.zero_()

        # 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

        output_logits = None
        for _ in range(random.randint(1,30)):
            optimizer.zero_grad(set_to_none=True)

            relaxed_field = model(init_state, mock_boundaries, mock_mask)
            output_logits = probing_head(relaxed_field, coords_y, coords_x)  # fix

            loss = criterion(output_logits, target_tensor)
            loss.backward()
            optimizer.step()

        if epoch % 10 == 0 or epoch == 1:
            with torch.no_grad():
                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)