import random
from dataclasses import dataclass

import matplotlib

matplotlib.use("Agg")  # headless backend for Spaces
import matplotlib.pyplot as plt

import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F


# ----------------------------
# Simple starter datasets
# ----------------------------
DATASETS = {
    "🪄 PoemBot": """Roses are red and skies are blue.
The moon shines softly over you.
A little bird sings in the tree.
The wind is dancing wild and free.
Morning light begins to glow.
Tiny flowers start to grow.
Clouds are floating, soft and slow.
Kindness is the seed we sow.
Stars are bright in velvet night.
Dreams can fly like paper kites.
Rain can tap a gentle beat.
Puddles sparkle on the street.
The sun comes up, the shadows run.
A day begins with hope and fun.
A quiet river hums a song.
It carries little leaves along.
""",
    "📖 StoryBot": """Once upon a time, a small turtle found a golden key. It opened a tiny door under an old tree. Inside, the turtle discovered a library for animals.
One day, Maya built a robot from cardboard and tape. The robot could only say kind things. Soon, everyone wanted to build one too.
A brave squirrel wanted to touch a cloud. It climbed the tallest pine tree in the park. From the top, the cloud looked like a giant pillow.
Leo lost his red balloon at the fair. A bird carried it across the sky. The next morning, Leo found it tied to his mailbox.
A dragon lived behind the school garden. It was not scary at all. Every Friday, it helped water the tomatoes.
""",
    "🎬 ReviewBot": """Movie: The Lion King
Review: This movie is emotional, exciting, and full of memorable songs.
Movie: Frozen
Review: This movie is magical and funny, with strong characters and great music.
Movie: Toy Story
Review: This movie is creative, warm, and teaches a lesson about friendship.
Movie: Finding Nemo
Review: This movie is colorful, adventurous, and perfect for families.
Movie: Spider-Man
Review: This movie is action packed, funny, and inspiring.
Movie: Inside Out
Review: This movie is smart, creative, and helps explain feelings.
Movie: Moana
Review: This movie is beautiful, brave, and filled with adventure.
Movie: Coco
Review: This movie is touching, musical, and full of family love.
""",
    "🍿 PopcornBot": """Movie: The Lion King
Review: A heartfelt adventure with unforgettable songs and big emotions.
Score: 9/10

Movie: Frozen
Review: Magical, funny, and full of heart, with music you will sing for days.
Score: 8/10

Movie: Toy Story
Review: A warm and clever story about friendship that never gets old.
Score: 10/10

Movie: Finding Nemo
Review: A colorful ocean journey that is exciting and perfect for families.
Score: 9/10

Movie: Spider-Man
Review: Action packed and inspiring, with humor and a big beating heart.
Score: 8/10

Movie: Inside Out
Review: A smart and creative look at feelings that makes you think and smile.
Score: 9/10

Movie: Moana
Review: A brave and beautiful voyage with stunning scenes and great songs.
Score: 9/10

Movie: Coco
Review: A touching musical full of family love and gorgeous colors.
Score: 10/10

Movie: Zootopia
Review: A funny and clever mystery with a kind and important message.
Score: 8/10

Movie: Encanto
Review: A bright and magical family story with catchy songs and big feelings.
Score: 9/10

Movie: Up
Review: A sweet and adventurous tale that is both happy and a little sad.
Score: 9/10

Movie: Ratatouille
Review: A charming story that says anyone can do great things with heart.
Score: 8/10
""",
    "🎥 CameronBot": """Scene 1: A young explorer finds an old map hidden inside a dusty book.
Scene 2: She follows the map into a deep forest and discovers a secret cave.

Scene 1: A robot wakes up alone in an empty factory.
Scene 2: It opens the rusty doors and steps into a bright new world outside.

Scene 1: A small boat drifts on a calm and quiet sea at sunrise.
Scene 2: A giant friendly whale rises beside it and guides the boat home.

Scene 1: A girl plants a single seed in her tiny garden.
Scene 2: By morning a tall tree has grown with glowing golden fruit.

Scene 1: A spaceship lands softly on a strange purple planet.
Scene 2: Curious aliens gather around and offer the crew a warm welcome.

Scene 1: A cat chases a ball of yarn across the living room floor.
Scene 2: The yarn rolls under the door and leads the cat on a wild adventure.

Scene 1: A knight stands before a locked castle gate at midnight.
Scene 2: The gate creaks open and reveals a hall full of sleeping dragons.

Scene 1: A scientist mixes two glowing liquids in her quiet lab.
Scene 2: A cloud of sparkling stars floats up and fills the whole room.

Scene 1: A boy finds a tiny door at the base of an old oak tree.
Scene 2: He crawls through and lands in a village built for mice.

Scene 1: A train stops at a station that is not on any map.
Scene 2: The passengers step out into a city made entirely of glass.

Scene 1: A painter draws a bird on a blank white wall.
Scene 2: The bird flaps its wings, lifts off the wall, and flies away.

Scene 1: A child whispers a wish into an old brass lantern.
Scene 2: The lantern glows and a kind genie floats out with a smile.
""",
}

DEFAULT_PROJECT = "🪄 PoemBot"


# ----------------------------
# Tiny character language model
# ----------------------------
class TinyCharModel(nn.Module):
    def __init__(self, vocab_size, emb_size=48, hidden_size=96):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.rnn = nn.GRU(emb_size, hidden_size, batch_first=True)
        self.head = nn.Linear(hidden_size, vocab_size)

    def forward(self, idx, hidden=None):
        x = self.embedding(idx)
        out, hidden = self.rnn(x, hidden)
        logits = self.head(out)
        return logits, hidden


@dataclass
class TrainState:
    model: object = None
    stoi: object = None
    itos: object = None
    vocab: object = None
    device: str = "cpu"
    trained: bool = False


def build_vocab(text):
    chars = sorted(list(set(text)))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    return chars, stoi, itos


def encode(text, stoi):
    return torch.tensor([stoi[c] for c in text if c in stoi], dtype=torch.long)


def decode(indices, itos):
    return "".join(itos[int(i)] for i in indices)


def sample_batch(data, block_size=64, batch_size=16):
    if len(data) <= block_size + 1:
        block_size = max(4, len(data) - 2)
    ix = torch.randint(0, len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i : i + block_size] for i in ix])
    y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix])
    return x, y


@torch.no_grad()
def generate_text(model, start_text, stoi, itos, length=300, temperature=0.8):
    model.eval()
    device = next(model.parameters()).device

    # keep only characters known to the model
    clean_start = "".join([c for c in start_text if c in stoi])
    if clean_start == "":
        clean_start = random.choice(list(stoi.keys()))

    idx = torch.tensor([[stoi[c] for c in clean_start]], dtype=torch.long, device=device)

    for _ in range(length):
        idx_cond = idx[:, -64:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / max(temperature, 0.1)
        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_id], dim=1)

    return decode(idx[0].tolist(), itos)


def make_loss_plot(losses):
    fig, ax = plt.subplots(figsize=(6, 3.4))
    fig.patch.set_facecolor("#ffffff")
    ax.set_facecolor("#fbfbfd")

    ax.plot(losses, color="#7c3aed", linewidth=2.2)
    ax.fill_between(range(len(losses)), losses, min(losses) if losses else 0,
                    color="#7c3aed", alpha=0.12)

    ax.set_xlabel("Training step", fontsize=11)
    ax.set_ylabel("Loss", fontsize=11)
    ax.set_title("Training loss goes down as the model learns 📉",
                 fontsize=12, fontweight="bold", color="#1f2937")
    ax.grid(True, linestyle="--", alpha=0.35)
    for spine in ["top", "right"]:
        ax.spines[spine].set_visible(False)
    fig.tight_layout()
    return fig


def load_dataset(project):
    return DATASETS[project]


def create_fresh_model(dataset_text):
    vocab, stoi, itos = build_vocab(dataset_text)
    model = TinyCharModel(len(vocab))
    return TrainState(model=model, stoi=stoi, itos=itos, vocab=vocab, trained=False)


def untrained_output(project, dataset_text, start_text, output_length, temperature):
    if len(dataset_text.strip()) < 100:
        return "⚠️ Please add more training text first. Try at least 100 characters."

    state = create_fresh_model(dataset_text)
    text = generate_text(
        state.model,
        start_text=start_text,
        stoi=state.stoi,
        itos=state.itos,
        length=int(output_length),
        temperature=float(temperature),
    )
    return text


def train_model(
    project,
    dataset_text,
    training_steps,
    start_text,
    output_length,
    temperature,
    progress=gr.Progress(),
):
    if len(dataset_text.strip()) < 100:
        return None, "⚠️ Please add more training text first. Try at least 100 characters.", ""

    torch.manual_seed(7)
    random.seed(7)

    state = create_fresh_model(dataset_text)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    state.device = device
    state.model.to(device)

    data = encode(dataset_text, state.stoi).to(device)
    optimizer = torch.optim.AdamW(state.model.parameters(), lr=2e-3)

    losses = []
    steps = int(training_steps)

    state.model.train()
    for step in progress.tqdm(range(steps), desc="Training tiny model"):
        xb, yb = sample_batch(data, block_size=64, batch_size=16)
        xb, yb = xb.to(device), yb.to(device)

        logits, _ = state.model(xb)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), yb.reshape(-1))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        losses.append(float(loss.item()))

    fig = make_loss_plot(losses)
    sample = generate_text(
        state.model,
        start_text=start_text,
        stoi=state.stoi,
        itos=state.itos,
        length=int(output_length),
        temperature=float(temperature),
    )

    state.trained = True
    first_loss = losses[0] if losses else 0.0
    last_loss = losses[-1] if losses else 0.0
    message = (
        f"✅ Done! The tiny model trained on **{len(dataset_text)} characters** "
        f"using **{len(state.vocab)} unique characters**.\n\n"
        f"Loss dropped from **{first_loss:.3f} → {last_loss:.3f}** over {steps} steps. "
        f"Device used: `{device}`."
    )
    return fig, sample, message


EXPLANATION = """
# 🧠 Tiny Generative AI Trainer

A **kid-friendly mini version** of how generative AI is trained. Students change only three main things:

1. **🎯 Project type** &nbsp;—&nbsp; poems, stories, or movie reviews
2. **⏱️ Training steps** &nbsp;—&nbsp; how long the model learns
3. **✍️ Start text** &nbsp;—&nbsp; the beginning of the text the model completes

The model learns by trying to **predict the next character** again and again.
When the loss curve goes down, the model is getting better at copying the pattern of the examples.
"""

CUSTOM_CSS = """
.gradio-container { max-width: 1100px !important; margin: auto !important; }

#hero {
    background: linear-gradient(135deg, #7c3aed 0%, #db2777 50%, #f59e0b 100%);
    border-radius: 18px;
    padding: 6px 26px;
    color: white;
    box-shadow: 0 10px 30px rgba(124, 58, 237, 0.25);
    margin-bottom: 8px;
}
#hero h1 { color: white !important; font-size: 2.1rem !important; }
#hero p, #hero li, #hero strong { color: rgba(255,255,255,0.95) !important; }
#hero a { color: #fde68a !important; }

.panel-card {
    border-radius: 16px !important;
    padding: 16px !important;
    background: var(--block-background-fill);
    box-shadow: 0 4px 18px rgba(0,0,0,0.06);
    border: 1px solid var(--border-color-primary);
}

#train-btn { font-weight: 700 !important; }

footer { visibility: hidden; }
"""


THEME = gr.themes.Soft(
    primary_hue="purple",
    secondary_hue="pink",
    font=[gr.themes.GoogleFont("Quicksand"), "system-ui", "sans-serif"],
)


with gr.Blocks(title="Tiny Generative AI Trainer") as demo:
    with gr.Group(elem_id="hero"):
        gr.Markdown(EXPLANATION)

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(elem_classes="panel-card"):
                gr.Markdown("### ⚙️ Controls")
                project = gr.Dropdown(
                    choices=list(DATASETS.keys()),
                    value=DEFAULT_PROJECT,
                    label="1. Choose project",
                )
                training_steps = gr.Slider(
                    minimum=50,
                    maximum=1500,
                    value=500,
                    step=50,
                    label="2. Training steps",
                    info="More steps = more learning (but slower).",
                )
                start_text = gr.Textbox(
                    value="Once upon",
                    label="3. Start text / prompt",
                )
                with gr.Row():
                    output_length = gr.Slider(
                        minimum=80,
                        maximum=600,
                        value=250,
                        step=20,
                        label="Output length",
                    )
                    temperature = gr.Slider(
                        minimum=0.3,
                        maximum=1.5,
                        value=0.8,
                        step=0.1,
                        label="Creativity 🎨",
                    )

                with gr.Row():
                    untrained_btn = gr.Button("🎲 Generate (untrained)", variant="secondary")
                    train_btn = gr.Button("🚀 Train & Generate", variant="primary", elem_id="train-btn")

        with gr.Column(scale=1):
            with gr.Group(elem_classes="panel-card"):
                with gr.Accordion("📚 Training examples / dataset", open=True):
                    dataset_text = gr.Textbox(
                        value=DATASETS[DEFAULT_PROJECT],
                        lines=14,
                        label="Students can edit this text or paste their own examples",
                    )

    gr.Markdown("## 🔍 Results")
    with gr.Row():
        with gr.Column():
            untrained_box = gr.Textbox(lines=8, label="🎲 Untrained model output")
        with gr.Column():
            trained_box = gr.Textbox(lines=8, label="✨ Trained model output")

    with gr.Row():
        loss_plot = gr.Plot(label="📉 Training loss plot")
    status = gr.Markdown()

    project.change(load_dataset, inputs=project, outputs=dataset_text)

    untrained_btn.click(
        untrained_output,
        inputs=[project, dataset_text, start_text, output_length, temperature],
        outputs=untrained_box,
    )

    train_btn.click(
        train_model,
        inputs=[project, dataset_text, training_steps, start_text, output_length, temperature],
        outputs=[loss_plot, trained_box, status],
    )

    with gr.Accordion("💬 Classroom discussion questions", open=False):
        gr.Markdown(
            """
- What changed after training?
- Did more training steps make the output better?
- What happens if the dataset is very small?
- What happens if the examples are all poems versus all movie reviews?
- Why does the model sometimes make spelling mistakes?
"""
        )


if __name__ == "__main__":
    demo.launch(theme=THEME, css=CUSTOM_CSS)