# ╔═══════════════════════════════════════════════════════════════════╗
# ║ SETUP — run me first! (identical in every Lab notebook)           ║
# ╚═══════════════════════════════════════════════════════════════════╝
import os, sys, pathlib

IN_COLAB = "google.colab" in sys.modules
SMOKE = os.environ.get("SMOKE_TEST", "0") == "1"   # tiny fast run for automated tests

REPO_URL = "https://github.com/A-Halimi/ai-studio-internship.git"  # instructor: set once

if IN_COLAB:
    if not pathlib.Path("/content/ai-studio-internship").exists():
        !git clone -q {REPO_URL} /content/ai-studio-internship
    %cd /content/ai-studio-internship/notebooks/day07
    %pip -q install gradio==6.19.0

HERE = pathlib.Path.cwd()
REPO = HERE.parents[1]                       # .../notebooks/dayNN -> repo root
DATA_DIR = pathlib.Path(os.environ.get("COURSE_DATA_DIR", REPO / "data"))
SAMPLES_DIR = REPO / "data_samples"          # small datasets shipped with the repo
MODELS_DIR = REPO / "ai_studio" / "models"   # where your Studio modules unlock
for p in (REPO / "data", MODELS_DIR):
    p.mkdir(parents=True, exist_ok=True)

import random
import numpy as np
SEED = 42
random.seed(SEED); np.random.seed(SEED)
print(f"✅ Setup done | Colab: {IN_COLAB} | Smoke test: {SMOKE}")
print(f"   data: {DATA_DIR}")

import torch
torch.manual_seed(SEED)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if DEVICE.type == "cuda":
    print(f"⚡ Running on GPU: {torch.cuda.get_device_name(0)}")
else:
    print("🐢 No GPU found — running on CPU. Everything still works, just slower.")

# ── CONFIG — the Lab's control panel. Tweak me later! ──────────────
EPOCHS        = 2 if SMOKE else 25   # passes over the tiny photo set (needs a few to converge)
BATCH_SIZE    = 16                   # photos per gradient step (dataset is tiny)
LEARNING_RATE = 1e-3                 # how big a step the new head takes

from torchvision import transforms

# ImageNet statistics — the giant was trained on pixels normalized this way.
MEAN = [0.485, 0.456, 0.406]
STD  = [0.229, 0.224, 0.225]

# The TRAINING pipeline (given): resize, crop, two random augmentations, normalize.
train_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.RandomHorizontalFlip(),      # random augmentation
    transforms.RandomRotation(10),          # random augmentation
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),
])
print("Training pipeline:", [type(t).__name__ for t in train_tf.transforms])

val_tf = transforms.Compose([
    # ==================== YOUR CODE HERE ====================
    ### HINT: the same first two steps as train_tf...
    ### HINT: ...then ToTensor and Normalize(MEAN, STD) — but NO flip/rotate
    ...
])
assert len(val_tf.transforms) == 4, "val_tf should have exactly 4 steps (no random augmentation)."
print("Validation pipeline:", [type(t).__name__ for t in val_tf.transforms])

from torchvision import datasets
from sklearn.model_selection import train_test_split

STARTER   = SAMPLES_DIR / "rps_starter"
COLLECTED = REPO / "data" / "rps"          # your own photos land here (optional, next section)

def all_targets(ds):
    """Every sample's class index — works for ImageFolder AND ConcatDataset."""
    if hasattr(ds, "targets"):
        return list(ds.targets)
    out = []
    for d in ds.datasets:
        out += list(d.targets)
    return out

def keep_jpg(path):
    """ImageFolder scans recursively — skip Jupyter's .ipynb_checkpoints duplicates."""
    return path.lower().endswith(".jpg") and ".ipynb_checkpoints" not in path

def build(tf):
    """One dataset with transform `tf`: starter photos (+ your collected ones if any)."""
    base = datasets.ImageFolder(STARTER, transform=tf)
    has_extra = COLLECTED.exists() and any(
        (COLLECTED / c).is_dir() and any((COLLECTED / c).glob("*.jpg"))
        for c in base.classes)
    if has_extra:
        extra = datasets.ImageFolder(COLLECTED, transform=tf, is_valid_file=keep_jpg)
        assert extra.class_to_idx == base.class_to_idx, (
            "Name your collected folders exactly rock/paper/scissors so the labels line up.")
        return torch.utils.data.ConcatDataset([base, extra]), base.classes
    return base, base.classes

# Two views of the SAME images: one augmented (train), one clean (val).
full_train, class_names = build(train_tf)
full_val,   _           = build(val_tf)
print("class_to_idx:", {c: i for i, c in enumerate(class_names)})   # alphabetical!
print("classes (training order):", class_names)

# Split the SAME indices 80/20, stratified so every class appears in both sets.
targets = all_targets(full_train)
train_idx, val_idx = train_test_split(
    list(range(len(full_train))), test_size=0.2, random_state=SEED, stratify=targets)
train_dataset = torch.utils.data.Subset(full_train, train_idx)
val_dataset   = torch.utils.data.Subset(full_val,   val_idx)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True,  num_workers=0)
val_loader   = torch.utils.data.DataLoader(val_dataset,   batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
print(f"{len(train_dataset)} training photos, {len(val_dataset)} validation photos.")

import matplotlib.pyplot as plt
from PIL import Image

fig, axes = plt.subplots(len(class_names), 4, figsize=(9, 7))
for row, cls in enumerate(class_names):
    files = sorted((STARTER / cls).glob("*.jpg"))[:4]
    for col, f in enumerate(files):
        axes[row, col].imshow(Image.open(f))
        axes[row, col].set_ylabel(cls if col == 0 else "", fontsize=12)
        axes[row, col].set_xticks([]); axes[row, col].set_yticks([])
fig.suptitle("Your entire training set — 12 photos per class (4 shown)")
plt.tight_layout(); plt.show()

# ⭐ OPTIONAL cell — skip it entirely if you're happy with the 36 starter photos.
import gradio as gr
from PIL import Image

RPS_DIR = REPO / "data" / "rps"
for c in ["rock", "paper", "scissors"]:
    (RPS_DIR / c).mkdir(parents=True, exist_ok=True)

def save_photo(frame, label):
    if frame is None:
        return "📷 Nothing to save yet — show a hand to the webcam (give it a second), or pick a file on the Upload tab."
    try:
        folder = RPS_DIR / label
        n = len(list(folder.glob("*.jpg")))
        path = folder / f"{label}_{n:03d}.jpg"
        Image.fromarray(frame).convert("RGB").save(path, quality=90)
    except Exception as e:
        return f"⚠️ Could not save that frame: {e}"
    total = sum(len(list((RPS_DIR / c).glob("*.jpg"))) for c in ["rock", "paper", "scissors"])
    return f"✅ Saved {label}/{path.name}. You now have {total} collected photo(s) total."

with gr.Blocks() as collector:
    gr.Markdown(
        "### 📸 Lab Camera — grow your dataset\n"
        "Pick the **Webcam** or **Upload** tab, get a hand photo, set the true move, hit **Save**."
    )
    wlatest = gr.State(None)          # latest live webcam frame (Save reads THIS, not the Image -> no 422)
    with gr.Row():
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("📷 Webcam"):
                    webcam = gr.Image(sources=["webcam"], streaming=True, type="numpy",
                                      height=360, label="Show your hand")
                    save_w = gr.Button("Save webcam shot 📥", variant="primary")
                with gr.Tab("📁 Upload"):
                    upload = gr.Image(sources=["upload"], type="numpy",
                                      height=360, label="Pick a photo")
                    save_u = gr.Button("Save uploaded photo 📥", variant="primary")
        with gr.Column(scale=1):
            lab = gr.Radio(["rock", "paper", "scissors"], value="rock", label="What move is this?")
            out = gr.Markdown()
    webcam.stream(lambda f, cur: f if f is not None else cur,
                  inputs=[webcam, wlatest], outputs=wlatest, stream_every=0.25)
    save_w.click(save_photo, [wlatest, lab], out)
    save_u.click(save_photo, [upload, lab], out)   # upload payloads are valid -> read the component directly

if not SMOKE:
    collector.launch(share=IN_COLAB)
    if not IN_COLAB:
        print("👉 For the full-size camera, open http://127.0.0.1:7860 in a browser tab.")

# 🛑 Run this to STOP the Lab Camera and free port 7860 (re-run the cell above to restart it).
if not SMOKE:
    collector.close()
    print("Lab Camera stopped — port 7860 is free.")

from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights

weights = MobileNet_V3_Small_Weights.DEFAULT     # downloads once (~10 MB), then cached
model = mobilenet_v3_small(weights=weights).to(DEVICE)

# The head is a small Sequential stack. Note the LAST layer outputs 1000 classes.
print(model.classifier)

# ==================== YOUR CODE HERE ====================
### HINT: loop over model.parameters() and set p.requires_grad = False
...

total     = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total params:     {total:,}")
print(f"Trainable params: {trainable:,}   <- frozen solid for now")

import torch.nn as nn

# For reference, model.classifier is:
#   Sequential(Linear(576,1024), Hardswish, Dropout, Linear(1024,1000))
#                                                     ^^^^^^^^ index 3 = the last layer
# ==================== YOUR CODE HERE ====================
### HINT: model.classifier[3] = nn.Linear(model.classifier[3].in_features, 3)
...
model = model.to(DEVICE)                 # the new layer needs to live on DEVICE too

total     = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable params now: {trainable:,}   ({100*trainable/total:.2f}% of the brain)")
assert trainable > 0, "The new head should have trainable parameters — did the transplant run?"

import torch.nn as nn
from tqdm.auto import tqdm

criterion = nn.CrossEntropyLoss()

# ==================== YOUR CODE HERE ====================
### HINT: Adam over filter(lambda p: p.requires_grad, model.parameters()), lr=LEARNING_RATE
...

def run_epoch(net, loader, train, optim=None):
    """One pass over `loader`. Returns (avg_loss, accuracy)."""
    net.train() if train else net.eval()
    torch.set_grad_enabled(train)
    total_loss, correct, n = 0.0, 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(DEVICE), yb.to(DEVICE)
        logits = net(xb)
        loss = criterion(logits, yb)
        if train:
            optim.zero_grad(); loss.backward(); optim.step()
        total_loss += loss.item() * len(xb)
        correct    += (logits.argmax(1) == yb).sum().item()
        n          += len(xb)
    torch.set_grad_enabled(True)
    return total_loss / n, correct / n

val_accs = []
for epoch in range(1, EPOCHS + 1):
    tr_loss, tr_acc = run_epoch(model, train_loader, True, optimizer)
    va_loss, va_acc = run_epoch(model, val_loader,   False)
    val_accs.append(va_acc)
    print(f"Epoch {epoch:2d}/{EPOCHS} | train loss {tr_loss:.3f} | val acc {va_acc:.0%}")

pretrained_acc = max(val_accs)
print(f"\n🏆 Best validation accuracy: {pretrained_acc:.0%}")
if not SMOKE:
    assert pretrained_acc >= 0.7, (
            f"Best val accuracy {pretrained_acc:.0%} on only {len(val_dataset)} validation "
            f"photos - re-run, or add a few of your own photos (section 3) for a steadier score.")

%matplotlib inline
model.eval()
mean_t = torch.tensor(MEAN).view(3, 1, 1)
std_t  = torch.tensor(STD).view(3, 1, 1)

wrong = []
with torch.no_grad():
    for i in range(len(val_dataset)):
        x, y = val_dataset[i]
        pred = model(x.unsqueeze(0).to(DEVICE)).argmax(1).item()
        if pred != y:
            wrong.append((x, y, pred))

print(f"{len(wrong)} of {len(val_dataset)} validation photos misclassified.")
if wrong:
    fig, axes = plt.subplots(1, min(len(wrong), 4), figsize=(9, 3), squeeze=False)
    for ax, (x, y, p) in zip(axes[0], wrong[:4]):
        img = (x * std_t + mean_t).clamp(0, 1).permute(1, 2, 0).numpy()   # undo Normalize
        ax.imshow(img); ax.axis("off")
        ax.set_title(f"true: {class_names[y]}\npred: {class_names[p]}", fontsize=10)
    plt.tight_layout(); plt.show()
else:
    print("🎉 Zero mistakes on the validation set!")

# ==================== YOUR CODE HERE ====================
### HINT: mobilenet_v3_small(weights=None)  — no borrowed knowledge
...
scratch.classifier[3] = nn.Linear(scratch.classifier[3].in_features, 3)
scratch = scratch.to(DEVICE)

# Nothing is frozen, so we train ALL of it (and Adam sees every parameter):
scratch_opt = torch.optim.Adam(scratch.parameters(), lr=LEARNING_RATE)

scratch_accs = []
for epoch in range(1, EPOCHS + 1):
    tr_loss, _      = run_epoch(scratch, train_loader, True, scratch_opt)
    _,       va_acc = run_epoch(scratch, val_loader,   False)
    scratch_accs.append(va_acc)
    print(f"Epoch {epoch:2d}/{EPOCHS} | train loss {tr_loss:.3f} | val acc {va_acc:.0%}")

scratch_acc = max(scratch_accs)
print(f"\nFrom-scratch best validation accuracy: {scratch_acc:.0%}")

%matplotlib inline
# ^ the Lab Camera (Gradio) silently switches matplotlib to the headless "Agg"
#   backend, which makes plt.show() a no-op — this line switches it back.
labels_bar = ["Pretrained\n(borrowed giant)", "From scratch\n(random giant)"]
accs = [pretrained_acc, scratch_acc]

plt.figure(figsize=(6, 4))
bars = plt.bar(labels_bar, accs, color=["#2a9d8f", "#e76f51"])
plt.axhline(1/3, ls="--", color="gray")
plt.text(1.0, 1/3 + 0.02, "random guessing (33%)", color="gray", ha="center")
for b, a in zip(bars, accs):
    plt.text(b.get_x() + b.get_width() / 2, a + 0.02, f"{a:.0%}", ha="center", fontsize=12)
plt.ylim(0, 1.1); plt.ylabel("Best validation accuracy")
plt.title("1.2 million photos of experience = the difference")
plt.tight_layout(); plt.show()

%matplotlib inline
def predict(pil_img):
    """Return (label, confidence) for one PIL image, using the clean val pipeline."""
    x = val_tf(pil_img.convert("RGB")).unsqueeze(0).to(DEVICE)   # (1, 3, 224, 224)
    # ==================== YOUR CODE HERE ====================
    ### HINT: probs = torch.softmax(model(x)[0], dim=0); idx = int(probs.argmax())
    ...
    return class_names[idx], float(probs[idx])

# Quick test: one starter image per class.
model.eval()
fig, axes = plt.subplots(1, len(class_names), figsize=(9, 3))
for ax, cls in zip(axes, class_names):
    img = Image.open(sorted((STARTER / cls).glob("*.jpg"))[0]).convert("RGB")
    label, conf = predict(img)
    ax.imshow(img); ax.axis("off")
    mark = "✅" if label == cls else "❌"
    ax.set_title(f"{mark} pred: {label}\n({conf:.0%} sure)", fontsize=10)
plt.tight_layout(); plt.show()

import json

model.cpu().eval()                                   # trace/save on CPU (the Studio loads on CPU)
example = torch.zeros(1, 3, 224, 224)                # the input contract: 224x224 RGB
traced = torch.jit.trace(model, example)
traced.save(str(MODELS_DIR / "rps_arena.pt"))

# class_names came straight from ImageFolder.classes — the exact training order.
(MODELS_DIR / "rps_labels.json").write_text(json.dumps(class_names), encoding="utf-8")

REQUIRED_FILES = ["rps_arena.pt", "rps_labels.json"]
missing = [f for f in REQUIRED_FILES if not (MODELS_DIR / f).exists()]
assert not missing, f"Something didn't save: {missing}"

model.to(DEVICE)                                     # move back in case you keep experimenting
print("Saved labels:", class_names)
print("🔓 UNLOCKED: RPS Arena! Run your Studio and beat your own creation:")
print("   python ai_studio/app.py   (from the repo root)   ->   🪨📄✂️ RPS Arena")