"""
Next-word prediction with a two-stage reward schedule.

Goal:
- First guess from a word sequence
- If wrong, reveal the first letter of the target word and guess again
- Training should converge toward solving it on the first guess

Training strategy:
- Cross-entropy on both stages
- Exact-match reward shaping
- Hint reward decays over epochs so the model relies less on the hint
"""

import re
import random
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
from torch.utils.data import Dataset, DataLoader


class Config:
    book_folder = "./books"

    # Model
    embed_dim = 128
    hidden_dim = 256
    hint_embed_dim = 16
    num_classes = None

    # Sequence
    seq_len = 8

    # Training
    batch_size = 64
    epochs = 20
    lr = 0.001
    rl_weight = 0.10
    baseline_momentum = 0.9
    hint_levels = [
        (0.00, 0.90, "full_minus_one"),
        (0.15, 0.90, "full_minus_one"),
        (0.25, 0.80, "two_letters"),
        (0.40, 0.70, "one_letter"),
        (0.55, 0.50, "one_letter"),
    ]

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


cfg = Config()


def load_markdown_files(folder_path):
    words = []
    folder = Path(folder_path)

    if not folder.exists():
        print(f"⚠️ Folder '{folder_path}' not found. Using sample text.")
        return sample_words()

    for md_file in folder.glob("*.md"):
        print(f"📖 Loading: {md_file.name}")
        with open(md_file, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()

        extracted = re.findall(r"[a-zA-Z]{2,}", content)
        words.extend([w.lower() for w in extracted])

    if len(words) == 0:
        return sample_words()

    print(f"✅ Extracted {len(words)} words from markdown files.")
    return words


def sample_words():
    sample_text = """
    intelligence artificial machine learning neural network deep
    consciousness awareness thinking reasoning planning algorithm
    """
    return re.findall(r"[a-zA-Z]{2,}", sample_text.lower())


def build_vocabulary(words, max_vocab_size=5000):
    word_counts = defaultdict(int)
    for word in words:
        word_counts[word] += 1

    sorted_words = sorted(word_counts.items(), key=lambda x: -x[1])
    vocab_words = [w for w, _ in sorted_words[:max_vocab_size]]

    word2idx = {word: idx for idx, word in enumerate(vocab_words)}
    idx2word = {idx: word for word, idx in word2idx.items()}

    print(f"📚 Vocabulary size: {len(vocab_words)}")
    return word2idx, idx2word, vocab_words


def char_to_idx(letter):
    if not letter or not letter.isalpha():
        return 26
    return ord(letter.lower()) - ord("a")


class SequenceHintDataset(Dataset):
    """Context sequence -> next word, plus target word for prefix hints."""

    def __init__(self, words, word2idx, seq_len=8, max_samples=50000):
        self.samples = []
        self.seq_len = seq_len
        self.word2idx = word2idx

        vocab_set = set(word2idx.keys())
        valid_words = [w for w in words if w in vocab_set and len(w) > 1]

        if len(valid_words) > max_samples:
            valid_words = valid_words[:max_samples]

        for i in range(len(valid_words) - seq_len):
            context_words = valid_words[i : i + seq_len]
            target_word = valid_words[i + seq_len]

            context_indices = [word2idx[w] for w in context_words]
            target_idx = word2idx[target_word]

            self.samples.append(
                {
                    "context_indices": context_indices,
                    "target_idx": target_idx,
                    "context_words": context_words,
                    "target_word": target_word,
                }
            )

        print(f"🖼️ Sequence-hint dataset: {len(self.samples)} samples (seq_len={seq_len})")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        return (
            torch.LongTensor(sample["context_indices"]),
            torch.LongTensor([sample["target_idx"]]),
            sample["context_words"],
            sample["target_word"],
        )


def collate_seq_hint_batch(batch):
    context_indices, target_idx, context_words, target_words = zip(*batch)
    context_indices = torch.stack(context_indices, dim=0)
    target_idx = torch.stack(target_idx, dim=0)
    return context_indices, target_idx, list(context_words), list(target_words)


def make_hint_prefix(word, mode):
    if len(word) <= 1:
        return word
    if mode == "full_minus_one":
        return word[:-1]
    if mode == "two_letters":
        return word[: min(2, len(word))]
    return word[:1]


def encode_hint_prefix(prefix, max_len=8):
    ids = [char_to_idx(ch) for ch in prefix[:max_len]]
    if not ids:
        ids = [26]
    return torch.LongTensor(ids)


class HintSequencePredictor(nn.Module):
    """Two-stage predictor: first guess, then hint-assisted guess."""

    def __init__(self, num_classes, embed_dim=128, hidden_dim=256, hint_embed_dim=16):
        super().__init__()

        self.word_embeddings = nn.Embedding(num_classes, embed_dim)
        self.hint_embeddings = nn.Embedding(27, hint_embed_dim)

        self.context_encoder = nn.GRU(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.2,
        )

        self.first_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, num_classes),
        )

        self.second_head = nn.Sequential(
            nn.Linear(hidden_dim + hint_embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, num_classes),
        )

    def encode_hint(self, hint_prefix_batch):
        hint_embeds = []
        for prefix_ids in hint_prefix_batch:
            emb = self.hint_embeddings(prefix_ids.to(self.hint_embeddings.weight.device))
            hint_embeds.append(emb.mean(dim=0))
        return torch.stack(hint_embeds, dim=0)

    def forward(self, context_indices, hint_prefix_batch=None):
        context_embed = self.word_embeddings(context_indices)
        _, hidden = self.context_encoder(context_embed)
        state = hidden[-1]

        first_logits = self.first_head(state)

        if hint_prefix_batch is None:
            return first_logits, None, state

        hint_embed = self.encode_hint(hint_prefix_batch)
        second_input = torch.cat([state, hint_embed], dim=1)
        second_logits = self.second_head(second_input)
        return first_logits, second_logits, state


def hint_prob_for_accuracy(first_acc):
    selected = cfg.hint_levels[0][1]
    for threshold, prob, _ in cfg.hint_levels:
        if first_acc >= threshold * 100.0:
            selected = prob
        else:
            break
    return selected


def hint_mode_for_accuracy(first_acc):
    selected = cfg.hint_levels[0][2]
    for threshold, _, mode in cfg.hint_levels:
        if first_acc >= threshold * 100.0:
            selected = mode
        else:
            break
    return selected


def train_model(model, train_loader, val_loader, epochs=20, lr=0.001):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    baseline = 0.0
    hint_prob = cfg.hint_levels[0][1]
    hint_mode = cfg.hint_levels[0][2]
    prev_val_first_acc = 0.0

    model.to(cfg.device)

    for epoch in range(epochs):
        model.train()

        train_loss = 0.0
        train_ce = 0.0
        train_rl = 0.0
        first_correct = 0
        second_correct = 0
        final_correct = 0
        total = 0

        for context_idx, target_idx, context_words, target_words in train_loader:
            context_idx = context_idx.to(cfg.device)
            target_idx = target_idx.squeeze(-1).to(cfg.device)

            # Step 1: first attempt, no hint.
            optimizer.zero_grad()
            first_logits, _, _ = model(context_idx, None)
            first_ce = criterion(first_logits, target_idx)

            # Reward shaping uses exact-match samples from both stages.
            dist1 = Categorical(logits=first_logits)
            action1 = dist1.sample()

            reward1 = (action1 == target_idx).float()
            shaped_reward = reward1
            batch_reward = shaped_reward.mean().item()

            baseline = cfg.baseline_momentum * baseline + (1.0 - cfg.baseline_momentum) * batch_reward
            advantage = shaped_reward - baseline

            policy_loss = -(advantage.detach() * dist1.log_prob(action1)).mean()

            first_loss = first_ce + cfg.rl_weight * policy_loss
            first_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            first_pred = first_logits.argmax(dim=1)
            first_ok = first_pred == target_idx
            second_ok = first_ok

            # Step 2: hint-assisted attempt, only on some batches.
            second_ce = torch.tensor(0.0, device=cfg.device)
            second_pred = first_pred
            use_hint = random.random() < hint_prob
            if use_hint:
                optimizer.zero_grad()
                hint_prefixes = [encode_hint_prefix(make_hint_prefix(word, hint_mode)) for word in target_words]
                _, second_logits, _ = model(context_idx, hint_prefixes)
                second_ce = criterion(second_logits, target_idx)
                second_ce.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                optimizer.step()

                second_pred = second_logits.argmax(dim=1)
                second_ok = second_pred == target_idx

            final_ok = first_ok | second_ok

            train_loss += first_loss.item() + second_ce.item()
            train_ce += first_ce.item() + second_ce.item()
            train_rl += policy_loss.item() if torch.is_tensor(policy_loss) else float(policy_loss)
            first_correct += first_ok.sum().item()
            second_correct += second_ok.sum().item()
            final_correct += final_ok.sum().item()
            total += target_idx.size(0)

            first_acc = 100.0 * first_correct / total
            second_acc = 100.0 * second_correct / total
            final_acc = 100.0 * final_correct / total
        avg_loss = train_loss / max(len(train_loader), 1)
        avg_ce = train_ce / max(len(train_loader), 1)
        avg_rl = train_rl / max(len(train_loader), 1)

        model.eval()
        val_first = 0
        val_second = 0
        val_final = 0
        val_total = 0
        val_loss = 0.0

        with torch.no_grad():
            for context_idx, target_idx, context_words, target_words in val_loader:
                context_idx = context_idx.to(cfg.device)
                target_idx = target_idx.squeeze(-1).to(cfg.device)

                use_hint = hint_prob >= 0.5
                if use_hint:
                    hint_prefixes = [encode_hint_prefix(make_hint_prefix(word, hint_mode)) for word in target_words]
                    first_logits, second_logits, _ = model(context_idx, hint_prefixes)
                else:
                    first_logits, _, _ = model(context_idx, None)
                    second_logits = None
                first_ce = criterion(first_logits, target_idx)
                second_ce = criterion(second_logits, target_idx) if second_logits is not None else torch.tensor(0.0, device=cfg.device)

                first_pred = first_logits.argmax(dim=1)
                second_pred = second_logits.argmax(dim=1) if second_logits is not None else first_pred
                first_ok = first_pred == target_idx
                second_ok = second_pred == target_idx
                final_ok = first_ok | second_ok

                val_loss += (first_ce.item() + second_ce.item())
                val_first += first_ok.sum().item()
                val_second += second_ok.sum().item()
                val_final += final_ok.sum().item()
                val_total += target_idx.size(0)

        val_first_acc = 100.0 * val_first / val_total
        val_second_acc = 100.0 * val_second / val_total
        val_final_acc = 100.0 * val_final / val_total
        avg_val_loss = val_loss / max(len(val_loader), 1)

        prev_val_first_acc = val_first_acc
        hint_prob = hint_prob_for_accuracy(prev_val_first_acc)
        hint_mode = hint_mode_for_accuracy(prev_val_first_acc)

        print(
            f"Epoch {epoch+1}/{epochs} | "
            f"Loss: {avg_loss:.4f} | CE: {avg_ce:.4f} | RL: {avg_rl:.4f} | "
            f"Train First: {first_acc:.2f}% | Train Hint: {second_acc:.2f}% | Train Final: {final_acc:.2f}% | "
            f"HintP(next): {hint_prob:.3f} | HintMode(next): {hint_mode} | Val Gate First: {prev_val_first_acc:.2f}% | "
            f"Val Loss: {avg_val_loss:.4f} | Val First: {val_first_acc:.2f}% | "
            f"Val Hint: {val_second_acc:.2f}% | Val Final: {val_final_acc:.2f}%"
        )

    return model


def test_model(model, test_loader, idx2word):
    model.eval()
    criterion = nn.CrossEntropyLoss()

    first_correct = 0
    second_correct = 0
    final_correct = 0
    total = 0
    test_loss = 0.0
    reward_total = 0.0

    print("\n" + "=" * 60)
    print("📊 TWO-STAGE NEXT-WORD TEST RESULTS")
    print("=" * 60)

    with torch.no_grad():
        for context_idx, target_idx, context_words, target_words in test_loader:
            context_idx = context_idx.to(cfg.device)
            target_idx = target_idx.squeeze(-1).to(cfg.device)

            hint_prefixes = [encode_hint_prefix(make_hint_prefix(word, "one_letter")) for word in target_words]
            first_logits, second_logits, _ = model(context_idx, hint_prefixes)
            first_ce = criterion(first_logits, target_idx)
            second_ce = criterion(second_logits, target_idx)

            first_pred = first_logits.argmax(dim=1)
            second_pred = second_logits.argmax(dim=1)

            first_ok = first_pred == target_idx
            second_ok = second_pred == target_idx
            final_ok = first_ok | second_ok

            first_correct += first_ok.sum().item()
            second_correct += second_ok.sum().item()
            final_correct += final_ok.sum().item()
            total += target_idx.size(0)
            test_loss += (first_ce.item() + second_ce.item())
            reward_total += first_ok.float().sum().item() + 0.25 * second_ok.float().sum().item()

    first_acc = 100.0 * first_correct / total
    second_acc = 100.0 * second_correct / total
    final_acc = 100.0 * final_correct / total
    avg_loss = test_loss / max(len(test_loader), 1)
    avg_reward = reward_total / total

    print(f"Test Loss: {avg_loss:.4f}")
    print(f"First Guess Accuracy: {first_acc:.2f}% ({first_correct}/{total})")
    print(f"Hint-Assisted Accuracy: {second_acc:.2f}% ({second_correct}/{total})")
    print(f"Final Accuracy After Hint: {final_acc:.2f}% ({final_correct}/{total})")
    print(f"Average Reward: {avg_reward:.4f}")

    print("\n🔍 Sample Predictions:")
    shown = 0
    with torch.no_grad():
        for context_idx, target_idx, context_words, target_words in test_loader:
            hint_prefixes = [encode_hint_prefix(make_hint_prefix(word, "one_letter")) for word in target_words]
            first_logits, second_logits, _ = model(context_idx.to(cfg.device), hint_prefixes)
            first_pred = first_logits.argmax(dim=1).cpu()
            second_pred = second_logits.argmax(dim=1).cpu()

            for i in range(min(4, len(target_words))):
                true_word = target_words[i]
                first_word = idx2word.get(first_pred[i].item(), "?")
                second_word = idx2word.get(second_pred[i].item(), "?")
                hint_letter = make_hint_prefix(true_word, "one_letter")

                if first_word == true_word:
                    status = "✅ first guess"
                elif second_word == true_word:
                    status = "🟡 hint recovered"
                else:
                    status = "❌ miss"

                print(
                    f"  {status}: {' '.join(context_words[i][-4:])} -> "
                    f"first='{first_word}', hint='{hint_letter}', second='{second_word}', true='{true_word}'"
                )
                shown += 1
                if shown >= 8:
                    break
            if shown >= 8:
                break

    print("=" * 60)
    return first_acc, final_acc


def main():
    print("🧠 Two-Stage Next-Word Reward Model")
    print("=" * 60)

    print("\n📂 Step 1: Loading markdown files...")
    words = load_markdown_files(cfg.book_folder)

    print("\n📚 Step 2: Building vocabulary...")
    word2idx, idx2word, vocab_words = build_vocabulary(words, max_vocab_size=2000)
    cfg.num_classes = len(vocab_words)

    print(f"\n🖼️ Step 3: Creating sequence-hint dataset (seq_len={cfg.seq_len})...")
    dataset = SequenceHintDataset(words, word2idx, seq_len=cfg.seq_len, max_samples=40000)

    print("\n📊 Step 4: Splitting data...")
    total = len(dataset)
    train_size = int(0.7 * total)
    val_size = int(0.15 * total)
    test_size = total - train_size - val_size

    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
        dataset, [train_size, val_size, test_size]
    )

    train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, collate_fn=collate_seq_hint_batch)
    val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, collate_fn=collate_seq_hint_batch)
    test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, collate_fn=collate_seq_hint_batch)

    print(f"  Train: {len(train_dataset)} | Val: {len(val_dataset)} | Test: {len(test_dataset)}")

    print("\n🏗️ Step 5: Building model...")
    model = HintSequencePredictor(
        num_classes=cfg.num_classes,
        embed_dim=cfg.embed_dim,
        hidden_dim=cfg.hidden_dim,
        hint_embed_dim=cfg.hint_embed_dim,
    )
    print(model)

    print("\n🚀 Step 6: Training two-stage reward model...")
    model = train_model(model, train_loader, val_loader, epochs=cfg.epochs, lr=cfg.lr)

    print("\n🧪 Step 7: Testing...")
    first_acc, final_acc = test_model(model, test_loader, idx2word)

    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "word2idx": word2idx,
            "idx2word": idx2word,
            "config": {
                "embed_dim": cfg.embed_dim,
                "hidden_dim": cfg.hidden_dim,
                "hint_embed_dim": cfg.hint_embed_dim,
                "num_classes": cfg.num_classes,
                "seq_len": cfg.seq_len,
            },
        },
        "two_stage_reward_sequence.pt",
    )

    print("\n💾 Model saved to 'two_stage_reward_sequence.pt'")
    print(f"First Guess Accuracy: {first_acc:.2f}%")
    print(f"Final Accuracy After Hint: {final_acc:.2f}%")


if __name__ == "__main__":
    main()