Functionality works, UI needs work

src/draft_model.py

@@ -63,7 +63,7 @@ class DraftModel:
     # Public API expected by streamlit_app.py                               #
     # --------------------------------------------------------------------- #
     @torch.no_grad()
-    def predict(self, pack: List[Dict], deck: List[Dict]) -> Dict:
+    def predict(self, pack: List[str], deck: List[str]) -> Dict:
 
         card_t = torch.stack([self._embed(c) for c in pack]).unsqueeze(0).to(self.device)
         if deck is None:
@@ -74,7 +74,6 @@ class DraftModel:
         vals = self.net(deck = deck_t, cards = card_t)
         #scores = torch.softmax(vals, dim=1).squeeze(0).cpu().numpy()
         scores = vals.squeeze(0).cpu().numpy()
-        print(scores)
         return {
             "pick": pack[scores.argmax()],
             "scores": scores.tolist(),

src/streamlit_app.py

@@ -33,6 +33,8 @@ from pathlib import Path
 
 from typing import Dict, List
 import pandas as pd
+import copy
+import numpy as np
 
 import requests
 import streamlit as st
@@ -41,20 +43,20 @@ from draft_model import DraftModel
 
 SUPPORTED_SETS_PATH = Path("src/helper_files/supported_sets.txt")
 
+@st.cache_resource(show_spinner="Loading draft model …")
+def load_model():
+    return DraftModel() 
+
+model = load_model()
 
 @st.cache_data(show_spinner="Reading supported sets …")
 def get_supported_sets(path: Path = SUPPORTED_SETS_PATH) -> List[str]:
     """Return a list of legal set codes read from *supported_sets.txt*."""
-
     if path.is_file():
         return [ln.strip() for ln in path.read_text().splitlines() if ln.strip()]
     return []
 
 
-@st.cache_resource(show_spinner="Loading draft model …")
-def load_model():
-    return DraftModel() 
-
 @st.cache_data(show_spinner="Calculating P1P1 ...")
 def p1p1_ranking(set_code: str):
     names, scores = model.get_p1p1(set_code)
@@ -65,36 +67,11 @@ def p1p1_ranking(set_code: str):
 
 
 @st.cache_data(show_spinner="Calculating card rankings …")
-def rank_cards(set_code: str) -> List[Dict]:
-    """Return a stubbed ranking list for *set_code*.
-
-    Replace with your real evaluation logic. For now we just pull 30 random
-    commons from the set and assign a dummy score.
-    """
-
-    url = f"https://api.scryfall.com/cards/search?q=set%3A{set_code}+unique%3Aprints+is%3Acommon"
-    cards: List[Dict] = []
-    while url and len(cards) < 60:  # cap network use
-        r = requests.get(url)
-        r.raise_for_status()
-        payload = r.json()
-        cards += payload["data"]
-        url = payload.get("next_page") if payload.get("has_more") else None
-
-    sample = random.sample(cards, k=min(30, len(cards))) if cards else []
-    ranked = [
-        {"name": c["name"], "score": round(random.random(), 2)} for c in sample
-    ]
-    ranked.sort(key=lambda x: x["score"], reverse=True)
-    return ranked
-
-
-model = load_model()
-
-def suggest_pick(pack: List[Dict], picks: List[Dict]) -> Dict:
-    if model is None:
-        return random.choice(pack)
-    return model.predict(pack=pack, picks=picks)  # type: ignore[attr-defined]
+def rank_cards(deck: List[str], pack: List[str]) -> List[Dict]:
+    out = model.predict(pack, deck)
+    pick = out["pick"]
+    scores = {pack[i]: score for i, score in enumerate(out["scores"])}
+    return pick, scores
 
 
 def fetch_card_image(card_name: str) -> str:
@@ -125,61 +102,167 @@ with st.sidebar:
     supported_sets = get_supported_sets()
 
     # Hide control in an expander (collapsed by default)
-    with st.expander("Set selection", expanded=False):
-        if supported_sets:
-            # UPDATED: dropdown instead of radio
-            set_code = st.selectbox(
-                "Choose a set",
-                supported_sets,
-                index=0,
-                key="set_code",
-            )
-        else:
-            st.warning(
-                "*supported_sets.txt* not found or empty. Using free-text input instead.",
-                icon="⚠️",
-            )
-            set_code = st.text_input("Set code", value="EOE", key="set_code")
+    if supported_sets:
+        # UPDATED: dropdown instead of radio
+        set_code = st.selectbox(
+            "Choose a set",
+            supported_sets,
+            index=0,
+            key="set_code",
+        )
+    else:
+        st.warning(
+            "*supported_sets.txt* not found or empty. Using free-text input instead.",
+            icon="⚠️",
+        )
+        set_code = st.text_input("Set code", value="EOE", key="set_code")
 
-        if st.button("Start new draft", type="primary"):
-            st.session_state["pack"] = generate_booster(set_code)
-            st.session_state["picks"] = []
 
 # -------- Session state ------------------------------------------------------
 
 st.session_state.setdefault("pack", [])
 st.session_state.setdefault("picks", [])
+st.session_state.setdefault("undo_stack", [])
 
 # -------- Main content organised in tabs ------------------------------------
 
 tabs = st.tabs(["Draft", "P1P1 Rankings"])
+def ensure_state():
+    """Make sure session_state has the pack and picks lists."""
+    if "pack" not in st.session_state:
+        st.session_state["pack"] = []
+    if "picks" not in st.session_state:
+        st.session_state["picks"] = []
+def add_card(target: str, card: str):
+    """target is 'pack' or 'picks'."""
+    ensure_state()
+    st.session_state[target].append(card)
+
+def remove_card(target: str, key: str):
+    ensure_state()
+    lst = st.session_state[target]
+    idx = next((i for i, c in enumerate(lst) if c == key), None)
+    if idx is not None:
+        lst.pop(idx)
+
+def push_undo():
+    """Save a snapshot of pack + picks so we can undo one step."""
+    ensure_state()
+    st.session_state["undo_stack"].append({
+        "pack": copy.deepcopy(st.session_state["pack"]),
+        "picks": copy.deepcopy(st.session_state["picks"]),
+    })
+    # (optional) cap history
+    if len(st.session_state["undo_stack"]) > 20:
+        st.session_state["undo_stack"].pop(0)
+
+def undo_last():
+    if st.session_state.get("undo_stack"):
+        snap = st.session_state["undo_stack"].pop()
+        st.session_state["pack"]  = snap["pack"]
+        st.session_state["picks"] = snap["picks"]
+        st.rerun()
+
+# --- callbacks ---
+def _add_selected_to_deck():
+    val = st.session_state.get("deck_selectbox")
+    if val:
+        add_card("picks", val)
+        st.session_state["deck_selectbox"] = None  # clear selection
+        st.toast(f"Added to deck: {val}")
+        st.rerun()
+
+def _add_selected_to_pack():
+    val = st.session_state.get("pack_selectbox")
+    if val:
+        add_card("pack", val)
+        st.session_state["pack_selectbox"] = None  # clear selection
+        st.toast(f"Added to pack: {val}")
+        st.rerun()
+
 
 # --- Tab 1: Draft ------------------------------------------------------------
 
 with tabs[0]:
-    if not st.session_state["pack"]:
-        st.info("Choose **Start new draft** in the sidebar to open pack 1.")
+    ensure_state()
+
+    
+    if st.session_state["undo_stack"]:
+        st.button("↩️ Undo last action", on_click=undo_last)
+
+    scores = {}
+    pick = None
+
+    if st.session_state["pack"]:
+        pack = st.session_state["pack"]
+        deck = st.session_state["picks"]
+
+        try:
+            pick, scores = rank_cards(deck, pack)
+        except Exception as e:
+            st.error(f"Error calculating card rankings: {e}")
+    if not set_code:
+        st.info("Pick a set to draft")
     else:
-        pack: List[Dict] = st.session_state["pack"]
-        picks: List[Dict] = st.session_state["picks"]
-
-        st.subheader(f"Pack {len(picks) // 15 + 1} — Pick {len(picks) % 15 + 1}")
-        suggested = suggest_pick(pack, picks)
-        st.success(f"**Model suggests:** {suggested['name']}")
-
-        cols = st.columns(5)
-        for idx, card in enumerate(pack):
-            col = cols[idx % 5]
-            col.image(fetch_card_image(card["name"]), use_column_width=True)
-            if col.button(f"Pick {card['name']}", key=f"pick-{idx}"):
-                picks.append(card)
-                pack.remove(card)
-                if not pack:  # end of pack ⇒ open a fresh booster
-                    st.session_state["pack"] = generate_booster(set_code)
-                st.experimental_rerun()
-
-        with st.expander("Current picks", expanded=False):
-            st.write("\n".join([c["name"] for c in picks]))
+        options = list(model.cards[set_code.lower()].keys())
+        c1, c2 = st.columns(2)
+        with c1:
+            st.subheader("Add to Deck")
+            deck_sel = st.selectbox(
+                "Search card (deck)",
+                options,
+                index=None,
+                placeholder="Type to search…",
+                key="deck_selectbox",
+                on_change=_add_selected_to_deck,   # <- auto-add
+            )
+
+            # Show current deck with remove buttons
+            if st.session_state["picks"]:
+                for i, card in enumerate(st.session_state["picks"]):
+                    cols = st.columns([6, 2])
+                    cols[0].write(card)
+                    if cols[1].button("Remove", key=f"rm-deck-{i}"):
+                        remove_card("picks", card)
+                        st.rerun()
+            else:
+                st.caption("Deck is empty.")
+
+        with c2:
+            st.subheader("Add to pack")
+            pack_sel = st.selectbox(
+                "Search card (pack)",
+                options,
+                index=None,
+                placeholder="Type to search…",
+                key="pack_selectbox",
+                on_change=_add_selected_to_pack,   # <- auto-add
+            )
+                
+            if st.session_state["pack"]:
+                pack_list = st.session_state["pack"]
+                vals = [scores.get(c) if scores and c in scores else np.nan for c in pack_list]
+                df_scores = pd.DataFrame({"Card": pack_list, "Score": vals})
+                tbl_col, btn_col = st.columns([4, 1], gap="small")
+                with tbl_col:
+                    st.dataframe(
+                        df_scores,
+                        use_container_width=True,
+                        column_config={
+                            "Score": st.column_config.NumberColumn(format="%.4f")
+                        },
+                        hide_index=True,
+                    )
+                with btn_col:
+                    for i, card in enumerate(pack_list):
+                        cols = st.columns([8, 4])
+                        if cols[1].button("Add", key=f"add_clear_{i}"):
+                            push_undo()
+                            st.session_state["picks"].append(card)
+                            st.session_state["pack"] = []  # or generate_booster(set_code)
+                            st.rerun()
+            else:
+                st.caption("Pack is empty.")
 
 # --- Tab 2: Card rankings ----------------------------------------------------