app.py

import sys, os
sys.path.insert(0, os.path.dirname(__file__))

import streamlit as st
import torch
import torch.nn.functional as F
import numpy as np
import pickle
import lmdb
import json
from urllib.parse import quote
from huggingface_hub import snapshot_download

from src.training import models
from src.utils import utils

st.set_page_config(page_title="Draft Visualizer", layout="wide")

MAX_CHOICES = 15
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BASICS = {"Mountain", "Forest", "Swamp", "Island", "Plains"}

HF_MODEL_REPO = os.environ.get("HF_MODEL_REPO", "")
HF_DATA_REPO  = os.environ.get("HF_DATA_REPO",  "")

# ---------------------------------------------------------------------------
# Download HF repos once — returns local paths identical to server layout
# ---------------------------------------------------------------------------

@st.cache_resource(show_spinner="Downloading model files…")
def get_model_dir():
    return snapshot_download(HF_MODEL_REPO)

@st.cache_resource(show_spinner="Downloading draft data…")
def get_data_dir():
    return snapshot_download(HF_DATA_REPO, repo_type="dataset")


def scryfall_image_url(card_name: str) -> str:
    return f"https://api.scryfall.com/cards/named?fuzzy={quote(card_name)}&format=image&version=normal"


# ---------------------------------------------------------------------------
# Model + vocab
# ---------------------------------------------------------------------------

@st.cache_resource(show_spinner="Loading model…")
def load_model_and_vocab(checkpoint_path: str, model_dir: str):
    config = utils.load_config(os.path.join(model_dir, "config.yaml"))
    config["embedding_path"] = os.path.join(model_dir, "card_encodings.pt")

    embedding_dict = utils.get_embedding_dict(config["embedding_path"], add_nontransformed=True)
    all_vecs = np.array(list(embedding_dict.values()))
    mean, std = all_vecs.mean(axis=0), all_vecs.std(axis=0)
    std[std == 0] = 1

    cards = sorted(embedding_dict.keys())
    card_to_idx = {c: i for i, c in enumerate(cards)}
    idx_to_card = {i: c for c, i in card_to_idx.items()}

    embedding_matrix = torch.tensor(
        np.stack([(embedding_dict[c] - mean) / std for c in cards]),
        dtype=torch.float32,
    )

    gih_wr_matrix = torch.full((len(cards),), -1.0)
    gih_folder = os.path.join(model_dir, "gih_wr")
    if os.path.isdir(gih_folder):
        for fname in os.listdir(gih_folder):
            if not fname.endswith("_gih.json"):
                continue
            with open(os.path.join(gih_folder, fname)) as f:
                for entry in json.load(f):
                    wr = entry.get("ever_drawn_win_rate")
                    if wr is None:
                        continue
                    if isinstance(wr, str):
                        wr = float(wr.rstrip("%")) / 100
                    name = utils.normalize_card_name(entry["name"])
                    if name in card_to_idx:
                        gih_wr_matrix[card_to_idx[name]] = float(wr)

    network = models.DecisionDraftTransformer(
        **config, embedding_matrix=embedding_matrix, gih_wr_matrix=gih_wr_matrix
    )
    state = torch.load(checkpoint_path, map_location="cpu")
    network.load_state_dict(state)
    network.to(DEVICE)
    network.eval()

    return network, card_to_idx, idx_to_card, config


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------

@st.cache_data
def load_draft(lmdb_path: str, draft_idx: int):
    env = lmdb.open(lmdb_path, readonly=True, lock=False)
    with env.begin() as txn:
        cur = txn.cursor()
        keys = [bytes(k) for k, _ in cur if k != b"__len__"]
    key = keys[draft_idx % len(keys)]
    with env.begin() as txn:
        data = pickle.loads(txn.get(key))
    env.close()
    sequence, in_maindeck, wins, losses, u_g, u_wr = data
    return sequence, in_maindeck, int(wins), int(losses), int(u_g), float(u_wr), len(keys)


def build_tensors(sequence, card_to_idx):
    T = len(sequence)
    history_idx = torch.zeros(1, T, dtype=torch.long)
    pack_idx    = torch.zeros(1, T, MAX_CHOICES, dtype=torch.long)
    pack_mask   = torch.zeros(1, T, MAX_CHOICES, dtype=torch.bool)
    seq_mask    = torch.zeros(1, T, dtype=torch.bool)
    for t, pack_cards in enumerate(sequence):
        history_idx[0, t] = card_to_idx.get(utils.normalize_card_name(pack_cards[0]), 0)
        for j, card in enumerate(pack_cards[:MAX_CHOICES]):
            pack_idx[0, t, j]  = card_to_idx.get(utils.normalize_card_name(card), 0)
            pack_mask[0, t, j] = True
    return (history_idx.to(DEVICE), pack_idx.to(DEVICE),
            pack_mask.to(DEVICE), seq_mask.to(DEVICE))


@torch.no_grad()
def run_model(network, sequence, card_to_idx, skill_target=0.60):
    history_idx, pack_idx, pack_mask, seq_mask = build_tensors(sequence, card_to_idx)
    B         = history_idx.shape[0]
    outcome   = torch.full((B,), skill_target, device=DEVICE)
    player_wr = torch.full((B,), skill_target, device=DEVICE)
    logits, play_logits, pick_play_logits, _, _, _ = network(history_idx, pack_idx, pack_mask, seq_mask, outcome, player_wr)
    bc_probs      = F.softmax(logits[0], dim=-1).cpu()
    play_sig      = torch.sigmoid(play_logits[0]).cpu()
    pick_play_sig = torch.sigmoid(pick_play_logits[0]).cpu()
    return bc_probs, play_sig, pick_play_sig, pack_mask[0].cpu(), pack_idx[0].cpu()


# ---------------------------------------------------------------------------
# Sidebar
# ---------------------------------------------------------------------------

if not HF_MODEL_REPO or not HF_DATA_REPO:
    st.error("Set HF_MODEL_REPO and HF_DATA_REPO as Space secrets.")
    st.stop()

model_dir = get_model_dir()
data_dir  = get_data_dir()

with st.sidebar:
    st.header("Config")

    run_folders = sorted(
        [d for d in os.listdir(model_dir) if os.path.isdir(os.path.join(model_dir, d))
         and any(f.endswith(".pt") for f in os.listdir(os.path.join(model_dir, d)))],
        reverse=True,
    )
    selected_run = st.selectbox("Run", options=run_folders)
    run_dir = os.path.join(model_dir, selected_run) if selected_run else ""
    ckpt_files = sorted(
        [f for f in os.listdir(run_dir) if f.endswith(".pt")]
        if run_dir else [],
        reverse=True,
    )
    selected_ckpt = st.selectbox("Checkpoint", options=ckpt_files)
    checkpoint_path = os.path.join(run_dir, selected_ckpt) if selected_ckpt else ""

    sets_with_lmdb = sorted([
        s for s in os.listdir(data_dir)
        if os.path.exists(os.path.join(data_dir, s, "test.lmdb"))
    ])
    selected_set = st.selectbox("Set", options=sets_with_lmdb)
    lmdb_split   = st.radio("Split", ["test", "train"], horizontal=True)
    lmdb_path    = os.path.join(data_dir, selected_set, f"{lmdb_split}.lmdb") if selected_set else ""

    skill_target = st.slider("Skill level", min_value=0.0, max_value=1.0, value=0.60, step=0.05)
    show_images = st.toggle("Show card images", value=True)

    if not (checkpoint_path and os.path.exists(lmdb_path)):
        st.warning("Select a run, checkpoint, and set.")
        st.stop()

    network, card_to_idx, idx_to_card, config = load_model_and_vocab(checkpoint_path, model_dir)
    st.success(f"Model on {DEVICE}")

    draft_idx = st.number_input("Draft index", min_value=0, value=0, step=1)

# ---------------------------------------------------------------------------
# Load draft + run model
# ---------------------------------------------------------------------------

sequence, in_maindeck, wins, losses, u_g, u_wr, n_drafts = load_draft(lmdb_path, draft_idx)
st.sidebar.caption(f"{n_drafts} drafts available")

T = len(sequence)
bc_probs, play_sig, pick_play_sig, pack_mask, pack_idx = run_model(network, sequence, card_to_idx, skill_target)

# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------

st.title("MTG Draft Visualizer")

c1, c2, c3, c4, c5 = st.columns(5)
c1.metric("Draft result", f"{wins}W – {losses}L")
c2.metric("Player WR", f"{u_wr:.1%}" if u_wr > 0 else "unknown")
c3.metric("Player games", f"{u_g:,}" if u_g > 0 else "unknown")
c4.metric("Total picks", T)
c5.metric("Set", selected_set)

st.divider()

# ---------------------------------------------------------------------------
# Pick slider
# ---------------------------------------------------------------------------

if "pick_slider" not in st.session_state:
    st.session_state.pick_slider = 0
st.session_state.pick_slider = max(0, min(T - 1, st.session_state.pick_slider))

col_prev, col_slider, col_next = st.columns([1, 10, 1])
with col_prev:
    if st.button("◀", use_container_width=True):
        st.session_state.pick_slider = max(0, st.session_state.pick_slider - 1)
with col_next:
    if st.button("▶", use_container_width=True):
        st.session_state.pick_slider = min(T - 1, st.session_state.pick_slider + 1)
with col_slider:
    pick = st.slider("Pick", min_value=0, max_value=T - 1, key="pick_slider")

pack_num     = pick // 15 + 1
pick_in_pack = pick % 15 + 1
pack_cards   = sequence[pick]
human_name   = utils.normalize_card_name(pack_cards[0])
n_cards      = int(pack_mask[pick].sum().item())

bc_top   = bc_probs[pick].masked_fill(~pack_mask[pick], float('-inf')).argmax().item()
play_top = play_sig[pick].masked_fill(~pack_mask[pick], float('-inf')).argmax().item()
bc_name   = idx_to_card.get(pack_idx[pick, bc_top].item(), "?")
play_name = idx_to_card.get(pack_idx[pick, play_top].item(), "?")

st.subheader(f"Pack {pack_num}, Pick {pick_in_pack}  (step {pick})  —  skill = {skill_target:.2f}")

m1, m2, m3 = st.columns(3)
m1.metric("👤 Human pick",    human_name)
m2.metric("🤖 BC pick",       bc_name,   delta="✓" if bc_name   == human_name else "✗")
m3.metric("▶ Top play rate", play_name, delta="✓" if play_name == human_name else "✗")

# ---------------------------------------------------------------------------
# Card image grid
# ---------------------------------------------------------------------------

if show_images:
    valid_slots = [j for j in range(MAX_CHOICES) if pack_mask[pick, j]]

    def border_color(name):
        is_human = name == human_name
        is_bc    = name == bc_name
        is_play  = name == play_name
        if is_human and is_bc and is_play: return "#00ff00"
        if is_human and is_bc:             return "#4488ff"
        if is_human and is_play:           return "#ff8800"
        if is_human:                       return "#aaccff"
        if is_play:                        return "#ff4400"
        if is_bc:                          return "#aa66ff"
        return "#333333"

    cols = st.columns(5)
    for idx_in_row, j in enumerate(valid_slots):
        cidx  = pack_idx[pick, j].item()
        name  = idx_to_card.get(cidx, "?")
        bc_p  = bc_probs[pick, j].item()
        play  = play_sig[pick, j].item()
        color = border_color(name)
        col   = cols[idx_in_row % 5]
        with col:
            url = scryfall_image_url(name)
            st.markdown(
                f'<img src="{url}" style="width:200px;border:3px solid {color};border-radius:8px"/>',
                unsafe_allow_html=True,
            )
            st.caption(f"bc {bc_p:.3f} | play {play:.2f}")

    st.markdown(
        "<small>🟢 all agree &nbsp; 🔵 human+BC &nbsp; 🟠 human+play &nbsp; "
        "🔴 play only &nbsp; 🟣 BC only &nbsp; 🩵 human only</small>",
        unsafe_allow_html=True,
    )
    st.divider()

# ---------------------------------------------------------------------------
# Deck so far
# ---------------------------------------------------------------------------

with st.expander(f"Deck so far ({pick} cards)", expanded=False):
    if pick == 0:
        st.caption("No picks yet.")
    else:
        deck = [utils.normalize_card_name(sequence[t][0]) for t in range(pick)]
        cols = st.columns(5)
        for i, name in enumerate(deck):
            pred_play = pick_play_sig[pick - 1, i].item() if pick > 0 else 0.0
            gt_play   = in_maindeck[i] if in_maindeck is not None else None
            gt_str    = (" ✓" if gt_play == 1.0 else " ✗") if gt_play is not None else ""
            with cols[i % 5]:
                url = scryfall_image_url(name)
                st.markdown(
                    f'<img src="{url}" style="width:200px;border-radius:8px"/>',
                    unsafe_allow_html=True,
                )
                st.caption(f"#{i+1} play {pred_play:.2f}{gt_str}")