app.py

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st
from PIL import Image
from scipy.ndimage import binary_dilation

# --- minicv (Milestone 1) ---
from minicv.frequency import apply_frequency_filter
from minicv.io import read_image
from minicv.utils import rgb2gray, to_float01
from minicv.transforms import resize as mc_resize
from minicv.filters import (
    gaussian_filter, sobel_filter, threshold_otsu,
    median_filter, threshold_adaptive, erode, dilate
)
from minicv.processing import gamma_correction, equalize_histogram, histogram_matching
from minicv.features import harris_corner_detector, canny_edge_detector
from minicv.transforms import rotate, resize

# --- el_nos_el_tany (Milestone 2) ---
PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT / "el_nos_el_tany"))

# --- Page Config ---
st.set_page_config(page_title="minicv & Milestone 2", layout="wide")

# --- Top-level Mode Selector ---
st.sidebar.title("Navigation")
APP_MODE = st.sidebar.radio(
    "Choose mode",
    ["minicv Studio", "Milestone 2: Predict an Image"],
    index=1,
)
st.sidebar.divider()


# ===========================================================================
# MILESTONE 2 — INTERACTIVE PREDICTOR
# ===========================================================================
M2_ROOT     = PROJECT_ROOT / "el_nos_el_tany"
SPLITS_DIR  = M2_ROOT / "data" / "splits"
FEATS_DIR   = M2_ROOT / "data" / "features"
RUNS_DIR    = M2_ROOT / "data" / "runs"

CLASSES = (
    "cleaning_the_floor", "climbing", "cutting_trees",
    "riding_a_horse", "rowing_a_boat", "watching_tv",
)

FAMILY_COLORS = {
    "color_hist": "#1f77b4", "hsv_hist":   "#9467bd",
    "stats":      "#ff7f0e", "grad_hist":  "#2ca02c",
    "canny_grid": "#8c564b", "hog":        "#d62728",
}


# ----- cached resources --------------------------------------------------
@st.cache_resource(show_spinner="Loading schema, MRMR selection, training stats...")
def load_pipeline_state():
    """Load everything that's expensive to construct."""
    from el_nos_el_tany import features as ftmod, mrmr as mrmrmod
    schema = ftmod.FeatureSchema.load(FEATS_DIR / "schema.json")
    sel    = mrmrmod.load_selection(FEATS_DIR / "mrmr_top800.json")
    idx    = list(sel.indices)

    X_tr_full, y_train = ftmod.load_features(FEATS_DIR / "train.npz")
    X_tr = X_tr_full[:, idx]
    mu = X_tr.mean(axis=0)
    sd = X_tr.std(axis=0)
    sd[sd == 0] = 1.0
    X_train_std = ((X_tr - mu) / sd).astype(np.float64)

    return {
        "schema":      schema,
        "selection":   sel,
        "mrmr_idx":    idx,
        "X_train":     X_tr_full,            # full 1107-dim train (for nearest-neighbor display)
        "X_train_std": X_train_std,          # MRMR + z-scored (for KNN inference)
        "y_train":     y_train,
        "mu":          mu,
        "sd":          sd,
    }


@st.cache_resource(show_spinner="Loading KNN classifier...")
def load_knn(k: int = 3, metric: str = "cosine"):
    from el_nos_el_tany.models import knn as knn_mod
    state = load_pipeline_state()
    return knn_mod.KNNClassifier(k=k, metric=metric).fit(state["X_train_std"], state["y_train"])


@st.cache_resource(show_spinner="Computing MRMR selection at requested K (first time only)...")
def get_mrmr_selection(K: int):
    """Load mrmr_topK.json if it exists; otherwise compute and persist it."""
    from el_nos_el_tany import features as ftmod, mrmr as mrmrmod
    cache_path = FEATS_DIR / f"mrmr_top{K}.json"
    if cache_path.exists():
        return mrmrmod.load_selection(cache_path)
    schema = ftmod.FeatureSchema.load(FEATS_DIR / "schema.json")
    X_train, y_train = ftmod.load_features(FEATS_DIR / "train.npz")
    sel = mrmrmod.select_topk(X_train, y_train, K=K, schema=schema, source="train")
    mrmrmod.save_selection(cache_path, sel, schema=schema)
    return sel


@st.cache_resource(show_spinner="Building KNN for feature subset...")
def build_knn_for_subset(label: str, k: int = 3, metric: str = "cosine"):
    """Build a from-scratch KNN classifier on a chosen feature subset of the
    training set. Used by the MRMR before/after comparison.

    label is one of: 'full', 'mrmr_800', 'mrmr_175'.
    Returns (knn, mu, sd, idx_list_or_None, D).
    """
    from el_nos_el_tany import features as ftmod
    from el_nos_el_tany.models import knn as knn_mod
    X_full, y_train = ftmod.load_features(FEATS_DIR / "train.npz")

    if label == "full":
        idx_list, X_tr = None, X_full
    elif label == "mrmr_800":
        idx_list = list(get_mrmr_selection(800).indices)
        X_tr = X_full[:, idx_list]
    elif label == "mrmr_175":
        idx_list = list(get_mrmr_selection(175).indices)
        X_tr = X_full[:, idx_list]
    else:
        raise ValueError(f"unknown subset label {label!r}")

    mu = X_tr.mean(axis=0)
    sd = X_tr.std(axis=0); sd[sd == 0] = 1.0
    X_tr_z = ((X_tr - mu) / sd).astype(np.float64)
    knn = knn_mod.KNNClassifier(k=k, metric=metric).fit(X_tr_z, y_train)
    return knn, mu, sd, idx_list, X_tr.shape[1]


@st.cache_resource(show_spinner="Loading Softmax checkpoint...")
def load_softmax():
    from el_nos_el_tany.models import softmax as sm_mod
    state = np.load(RUNS_DIR / "sec8_softmax_adam" / "best.npz")
    sm = sm_mod.SoftmaxRegression(num_features=800, num_classes=6)
    sm.load_state_dict({"W": state["model.W"], "b": state["model.b"]})
    return sm


@st.cache_resource(show_spinner="Loading CNN checkpoint...")
def load_cnn():
    from el_nos_el_tany.models import cnn as cnn_mod
    state = np.load(RUNS_DIR / "sec8_cnn_adam" / "best.npz")
    clean = {k.replace("model.", ""): v for k, v in state.items() if k.startswith("model.")}
    m = cnn_mod.SimpleCNN(num_classes=6, in_channels=3, input_size=64, seed=42)
    m.load_state_dict(clean)
    return m


@st.cache_data(show_spinner=False)
def load_train_index():
    """A small DataFrame we can sample 'use a training image' demos from."""
    return pd.read_csv(SPLITS_DIR / "train.csv")


@st.cache_data(show_spinner=False)
def load_split_index(split: str) -> pd.DataFrame:
    """Load any split CSV (train / val / test)."""
    return pd.read_csv(SPLITS_DIR / f"{split}.csv")


# ----- helpers -----------------------------------------------------------
def standardize_features(v: np.ndarray) -> np.ndarray:
    """Apply MRMR-select + z-score using cached training stats; returns (1, 800)."""
    s = load_pipeline_state()
    v_sel = v[s["mrmr_idx"]]
    return ((v_sel - s["mu"]) / s["sd"]).astype(np.float64).reshape(1, -1)


def predict_all(img96: np.ndarray, img64: np.ndarray):
    """Run all three models on the image; return a dict with per-model results."""
    from el_nos_el_tany import features as ftmod
    from el_nos_el_tany.models.softmax import stable_softmax

    v_pool = ftmod.extract_pool(img96)
    v_std  = standardize_features(v_pool)

    knn_clf = load_knn()
    knn_pred = int(knn_clf.predict(v_std)[0])
    knn_proba = knn_clf.predict_proba(v_std)[0]
    nn_dist, nn_idx = knn_clf.kneighbors(v_std)

    sm_clf = load_softmax()
    sm_pred = int(sm_clf.predict(v_std)[0])
    sm_proba = sm_clf.predict_proba(v_std)[0]

    cnn_clf = load_cnn()
    img_chw = img64.transpose(2, 0, 1)[None].astype(np.float32)
    logits = cnn_clf.forward(img_chw)
    cnn_proba = stable_softmax(logits)[0]
    cnn_pred = int(np.argmax(cnn_proba))

    return {
        "feature_vector": v_pool,
        "feature_std": v_std[0],
        "KNN":     {"pred": knn_pred, "proba": knn_proba,
                    "neighbors": (nn_dist[0], nn_idx[0])},
        "Softmax": {"pred": sm_pred, "proba": sm_proba},
        "CNN":     {"pred": cnn_pred, "proba": cnn_proba},
    }


def proba_table(probas: np.ndarray, true_class: str | None = None) -> pd.DataFrame:
    """Build a sortable per-class probability table."""
    df = pd.DataFrame({
        "class": CLASSES,
        "probability": [float(p) for p in probas],
    }).sort_values("probability", ascending=False).reset_index(drop=True)
    df["bar"] = df["probability"]
    if true_class is not None:
        df["match"] = df["class"].apply(lambda c: "✅" if c == true_class else "")
    return df


def _show_proba(probas: np.ndarray, true_class: str | None = None):
    df = proba_table(probas, true_class=true_class)
    st.dataframe(
        df.style.format({"probability": "{:.3f}", "bar": "{:.3f}"})
                .background_gradient(subset=["bar"], cmap="Greens"),
        use_container_width=True, hide_index=True,
    )


def _safe_float01(img):
    """Safely forces an image to [0, 1] float32 to exactly match Notebook 01."""
    res = img.astype(np.float32)
    if res.max() > 1.0:
        res /= 255.0
    return res

def _process_upload(uploaded) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Convert an uploaded file to (raw_rgb, img96, img64) all RGB float32 [0,1]."""
    pil = Image.open(uploaded).convert("RGB")
    raw = np.asarray(pil, dtype=np.uint8)
    
    # Resize, then safely scale to [0, 1]
    img96 = _safe_float01(mc_resize(raw, (96, 96), method="bilinear"))
    img64 = _safe_float01(mc_resize(raw, (64, 64), method="bilinear"))
    return raw, img96, img64

def _process_path(path):
    # 1. Normalize slashes for cross-platform compatibility
    safe_path = str(path).replace("\\", "/")
    
    # 2. Dynamically rebuild the path using your existing PROJECT_ROOT
    if "el_nos_el_tany/" in safe_path:
        relative_tail = safe_path.split("el_nos_el_tany/")[-1]
        safe_path = str(PROJECT_ROOT / "el_nos_el_tany" / relative_tail)
        
    # 3. Load the image 
    raw = read_image(safe_path)
    
    # 4. Generate the resized versions
    img96 = mc_resize(raw, (96, 96), method='bilinear')
    img64 = mc_resize(raw, (64, 64), method='bilinear')
    
    # 5. CRITICAL FIX: Force to [0, 1] float32!
    img96 = _safe_float01(img96)
    img64 = _safe_float01(img64)
    
    return raw, img96, img64

# ----- main page ---------------------------------------------------------
def m2_predictor():
    st.title("Milestone 2 — Image Classifier")
    st.caption(
        "Upload an image (or pick one from the training set) and step through "
        "the full pipeline: preprocess → augment → 1,107-dim feature pool → "
        "MRMR top-800 → KNN / Softmax / CNN prediction."
    )

    # --- Sidebar: image input ---
    st.sidebar.header("Image input")
    src_mode = st.sidebar.radio(
        "Source",
        ["Pick a TEST sample (unseen)", "Upload", "Pick a train / val sample"],
        index=0,
        help="Test images are the ones the model has never seen — those are the "
             "honest demo. Training images can be cherry-picked, so use them only "
             "for the augmentation tab or sanity checks.",
    )

    raw = img96 = img64 = true_class = None

    if src_mode == "Upload":
        up = st.sidebar.file_uploader("Image file", type=["jpg", "jpeg", "png"])
        if up is not None:
            raw, img96, img64 = _process_upload(up)
            true_class = st.sidebar.selectbox(
                "True class (optional, for scoring)",
                ["(unknown)"] + list(CLASSES),
            )
            if true_class == "(unknown)":
                true_class = None
    else:
        if src_mode.startswith("Pick a TEST"):
            split = "test"
            st.sidebar.caption("📊 Held-out test set — the model has **never** seen these.")
        else:
            split = st.sidebar.selectbox("Split", ["train", "val"], index=0)
            if split == "train":
                st.sidebar.caption("⚠️ Training set — model already saw these images, so high accuracy here is not impressive.")
            else:
                st.sidebar.caption("Validation set — used to pick best K / k / hyperparameters.")
        df_split = load_split_index(split)
        cls = st.sidebar.selectbox("Class", list(CLASSES))
        sub = df_split[df_split["class_name"] == cls].reset_index(drop=True)
        if len(sub) == 0:
            st.sidebar.warning(f"No `{cls}` images in {split} split.")
            return
        idx = st.sidebar.slider(f"Index (0..{len(sub)-1})", 0, max(0, len(sub) - 1), 0)
        path = sub.iloc[idx]["image_path"]
        st.sidebar.caption(f"{split}/{Path(path).name}")
        raw, img96, img64 = _process_path(path)
        true_class = cls

    st.sidebar.divider()

    # --- Sidebar: step selector ---
    st.sidebar.header("Pipeline step")
    step = st.sidebar.radio(
        "View",
        [
            "Overview (image + final predictions)",
            "1. Preprocessing",
            "2. Augmentation",
            "3. Feature extraction",
            "4. MRMR feature selection (with before/after)",
            "5. KNN prediction",
            "6. Softmax prediction",
            "7. CNN prediction",
            "8. All predictions side-by-side",
        ],
        index=0,
    )

    if raw is None:
        st.info("Upload an image or pick a training sample from the sidebar to start.")
        st.subheader("The 6 classes the model knows")
        st.write(", ".join(f"`{c}`" for c in CLASSES))
        return

    # ========== STEP DISPATCH ==========
    if step.startswith("Overview"):
        c1, c2 = st.columns([1, 2])
        with c1:
            st.subheader("Input image")
            st.image(raw, use_container_width=True)
            st.caption(f"Raw shape: {raw.shape[1]}×{raw.shape[0]}")
            if true_class:
                st.markdown(f"**True class:** `{true_class}`")
        with c2:
            st.subheader("Predictions from all 3 models")
            with st.spinner("Running KNN + Softmax + CNN..."):
                results = predict_all(img96, img64)
            for name in ("KNN", "Softmax", "CNN"):
                pred_class = CLASSES[results[name]["pred"]]
                ok = "✅" if (true_class and pred_class == true_class) else ""
                st.markdown(
                    f"**{name}** → `{pred_class}` {ok}  "
                    f"(confidence {results[name]['proba'][results[name]['pred']]:.1%})"
                )
        st.divider()
        st.markdown("Use the radio in the sidebar to see what each pipeline step does to this image.")

    elif step.startswith("1."):
        st.subheader("1. Preprocessing — `minicv` pipeline")
        st.markdown(
            "`read_image` → `gray2rgb` (if needed) → `to_uint8` → "
            "`resize(bilinear)` → `to_float01`. Two target sizes are produced — "
            "**96×96** for the handcrafted feature pool (KNN/Softmax) and "
            "**64×64** for the CNN."
        )
        c1, c2, c3 = st.columns(3)
        with c1:
            st.image(raw, caption=f"Raw {raw.shape[1]}×{raw.shape[0]}", use_container_width=True)
        with c2:
            st.image(img96, caption="96×96 (features path)", use_container_width=True)
        with c3:
            st.image(img64, caption="64×64 (CNN path)", use_container_width=True)
        norm_path = SPLITS_DIR / "norm_stats.json"
        if norm_path.exists():
            stats = json.loads(norm_path.read_text())
            st.markdown("**Train-set per-channel mean / std** (computed once, never touches val/test):")
            st.dataframe(pd.DataFrame({"channel": ["R", "G", "B"],
                                       "mean": stats["mean"],
                                       "std":  stats["std"]}).round(4),
                         use_container_width=True, hide_index=True)

    elif step.startswith("2."):
        st.subheader("2. Augmentation (training-only)")
        st.markdown(
            "Six stochastic transforms wrapped around `minicv` functions. "
            "At training time each fires with probability ≈ 0.5 — every epoch "
            "the model sees a slightly different version of the same image."
        )
        from el_nos_el_tany import augment
        n_grid = st.slider("How many random samples", 1, 12, 6)
        seed_base = st.number_input("Base seed", value=42, step=1)
        if st.button("Re-roll"):
            seed_base = int(np.random.randint(0, 1_000_000))

        cols_per_row = 4
        rows = (n_grid + cols_per_row - 1) // cols_per_row
        for r in range(rows):
            cs = st.columns(cols_per_row)
            for c in range(cols_per_row):
                k = r * cols_per_row + c
                if k >= n_grid: break
                aug_fn = augment.default_augmenter(seed=int(seed_base + k))
                out = aug_fn(img96.copy())
                cs[c].image(np.clip(out, 0, 1), caption=f"sample {k}", use_container_width=True)

        st.caption("Compare these to the original 96×96 — augmentation perturbs lighting, "
                   "geometry, and crop while keeping the action recognizable.")

    elif step.startswith("3."):
        st.subheader("3. Feature extraction — 1,107-dim pool")
        st.markdown(
            "Six families concatenated per image; every family is computed by a "
            "`minicv` descriptor from Milestone 1 — no descriptor logic is reimplemented."
        )
        from el_nos_el_tany import features as ftmod
        v = ftmod.extract_pool(img96)

        # Family table
        s = load_pipeline_state()
        rows = []
        for fam, sl in s["schema"].family_slices.items():
            block = v[sl]
            rows.append({
                "family":  fam,
                "dim":     sl.stop - sl.start,
                "min":     float(block.min()),
                "mean":    float(block.mean()),
                "max":     float(block.max()),
                "fraction of vector": (sl.stop - sl.start) / s["schema"].dim,
            })
        st.dataframe(pd.DataFrame(rows).style.format({
            "min": "{:.3f}", "mean": "{:.3f}", "max": "{:.3f}",
            "fraction of vector": "{:.1%}"
        }), use_container_width=True, hide_index=True)

        # Bar plot color-coded by family
        st.markdown("**Per-dim values** (color = family)")
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(1, 1, figsize=(11, 2.6))
        for fam, sl in s["schema"].family_slices.items():
            ax.bar(np.arange(sl.start, sl.stop), v[sl],
                   color=FAMILY_COLORS[fam], width=1.0,
                   label=f"{fam} ({sl.stop - sl.start})")
        ax.set_xlim(0, s["schema"].dim)
        ax.set_xlabel("feature index"); ax.set_ylabel("value")
        ax.legend(loc="upper right", fontsize=8)
        st.pyplot(fig, use_container_width=True)

    elif step.startswith("4."):
        st.subheader("4. MRMR feature selection — before / after")
        st.markdown(
            "MRMR (Min-Redundancy Max-Relevance) compresses the 1,107-dim pool into a smaller subset. "
            "This view shows whether that compression preserved discriminative power."
        )

        # --- Test-set accuracy across feature subsets (from §4.2.2) ---
        st.markdown("**Test-set accuracy across feature subsets** (precomputed from the notebook 02 K-sweep)")
        headline = pd.DataFrame([
            {"feature subset": "full pool",       "D": 1107, "compression": "1.0x",
             "LR test": 0.710, "SVM test": 0.729},
            {"feature subset": "MRMR top-800",    "D":  800, "compression": "1.4x",
             "LR test": 0.710, "SVM test": 0.752},
            {"feature subset": "MRMR top-175 ★",  "D":  175, "compression": "6.3x",
             "LR test": 0.664, "SVM test": 0.724},
        ])
        st.dataframe(
            headline.style.format({"LR test": "{:.3f}", "SVM test": "{:.3f}"})
                          .background_gradient(subset=["SVM test"], cmap="Greens"),
            use_container_width=True, hide_index=True,
        )

        # --- Live KNN comparison for this image ---
        st.markdown("**Live KNN (k=3, cosine) on this image — full pool vs MRMR-800 vs MRMR-175**")
        from el_nos_el_tany import features as ftmod
        v_pool = ftmod.extract_pool(img96)

        rows = []
        for label_key, label_disp in [("full",     "full pool"),
                                      ("mrmr_800", "MRMR 800"),
                                      ("mrmr_175", "MRMR 175 ★")]:
            knn, mu, sd, idx_list, D = build_knn_for_subset(label_key)
            v = v_pool if idx_list is None else v_pool[idx_list]
            v_z = ((v - mu) / sd).astype(np.float64).reshape(1, -1)
            pred  = int(knn.predict(v_z)[0])
            proba = knn.predict_proba(v_z)[0]
            ok = ("✅" if true_class and CLASSES[pred] == true_class else
                  "❌" if true_class else "—")
            rows.append({
                "feature subset":  label_disp,
                "D":               D,
                "predicted class": CLASSES[pred],
                "confidence":      float(proba[pred]),
                "match":           ok,
            })
        st.dataframe(
            pd.DataFrame(rows).style.format({"confidence": "{:.1%}"}),
            use_container_width=True, hide_index=True,
        )
        if true_class:
            st.caption(f"True class: `{true_class}` — if the three rows agree, MRMR's compression is preserving the signal for this image.")

        # --- Family breakdown of the K=800 selection (what MRMR kept) ---
        st.markdown("**What MRMR top-800 kept, by family**")
        s = load_pipeline_state()
        from collections import Counter
        fam_counts = Counter(s["schema"].families[i] for i in s["mrmr_idx"])
        st.dataframe(
            pd.DataFrame([{"family": f,
                           "in pool": sl.stop - sl.start,
                           "kept by MRMR": fam_counts.get(f, 0)}
                          for f, sl in s["schema"].family_slices.items()]),
            use_container_width=True, hide_index=True,
        )

    elif step.startswith("5."):
        st.subheader("5. KNN prediction (from-scratch, MRMR features)")
        st.markdown(
            "KNN is from-scratch — distance metrics implemented in "
            "`el_nos_el_tany.models.knn` (`pairwise_cosine` / `pairwise_l2`). "
            "Soft-vote probability is `count_c / k` over the k nearest training neighbors."
        )

        # ---- interactive controls ----
        ctrl_l, ctrl_r = st.columns(2)
        with ctrl_l:
            metric = st.radio(
                "Distance metric",
                ["cosine", "l2"],
                index=0, horizontal=True,
                help="`cosine` = 1 − cos(θ); scale-invariant. `l2` = Euclidean.",
                key="knn_metric",
            )
        with ctrl_r:
            k = st.slider("k (number of neighbors)", min_value=1, max_value=51, value=3, step=2,
                          help="Odd values avoid 2-2 ties on 6-class problems.",
                          key="knn_k")

        knn_clf = load_knn(k=int(k), metric=metric)
        from el_nos_el_tany import features as ftmod
        v = ftmod.extract_pool(img96)
        v_std = standardize_features(v)
        pred = int(knn_clf.predict(v_std)[0])
        proba = knn_clf.predict_proba(v_std)[0]
        nn_dist, nn_idx = knn_clf.kneighbors(v_std)

        c1, c2 = st.columns([1, 1])
        with c1:
            st.image(raw, caption="Query image", use_container_width=True)
            st.markdown(f"### Predicted: `{CLASSES[pred]}`")
            st.caption(f"using **k={k}**, metric=**{metric}**")
            if true_class:
                ok = "✅ correct" if CLASSES[pred] == true_class else f"❌ true: `{true_class}`"
                st.markdown(f"**{ok}**")
        with c2:
            st.markdown("**Per-class probability** (votes / k)")
            _show_proba(proba, true_class=true_class)

        # ---- nearest neighbors (cap display at 12 to keep layout sane) ----
        n_show = min(int(k), 12)
        st.subheader(f"{n_show} nearest training neighbors  ·  metric = {metric}")
        if n_show < int(k):
            st.caption(f"(showing {n_show} of {k} — vote uses all {k})")
        df_train = load_train_index()
        cols_per_row = 4 if n_show > 3 else n_show
        s = load_pipeline_state()
        y_train = s["y_train"]
        n_correct = int(np.sum(y_train[nn_idx[0][:int(k)]] == CLASSES.index(true_class))) if true_class else None

        rows = (n_show + cols_per_row - 1) // cols_per_row
        flat_iter = iter(zip(nn_dist[0][:n_show], nn_idx[0][:n_show]))
        idx_counter = 0
        for r in range(rows):
            cs = st.columns(cols_per_row)
            for c in range(cols_per_row):
                if idx_counter >= n_show: break
                dist, ti = next(flat_iter)
                row = df_train.iloc[int(ti)]
                tick = ""
                if true_class:
                    tick = "✅" if row["class_name"] == true_class else "❌"
                try:
                    cs[c].image(
                        row["image_path"],
                        caption=f"#{idx_counter+1} {tick} {row['class_name']}  d={dist:.3f}",
                        use_container_width=True,
                    )
                except Exception:
                    cs[c].warning(f"image not found: {row['image_path']}")
                idx_counter += 1

        if true_class is not None:
            st.info(
                f"**{n_correct} / {k}** of the chosen neighbors are the correct class "
                f"(`{true_class}`) — that's the soft-vote probability "
                f"`{n_correct/k:.1%}`."
            )

    elif step.startswith("6."):
        st.subheader("6. Softmax regression (from-scratch, MRMR features)")
        st.markdown(
            "Multi-class softmax: `p_c = exp(z_c) / Σ exp(z_j)` with the `-max(z)` "
            "stability shift. Trained from scratch with Adam + cosine LR."
        )
        sm_clf = load_softmax()
        from el_nos_el_tany import features as ftmod
        v = ftmod.extract_pool(img96)
        v_std = standardize_features(v)
        pred = int(sm_clf.predict(v_std)[0])
        proba = sm_clf.predict_proba(v_std)[0]

        c1, c2 = st.columns([1, 1])
        with c1:
            st.image(raw, caption="Query image", use_container_width=True)
            st.markdown(f"### Predicted: `{CLASSES[pred]}`")
            if true_class:
                ok = "✅ correct" if CLASSES[pred] == true_class else f"❌ true: `{true_class}`"
                st.markdown(f"**{ok}**")
            st.metric("Confidence", f"{proba[pred]:.1%}")
        with c2:
            st.markdown("**Per-class probability**")
            _show_proba(proba, true_class=true_class)

        # Show raw logits too
        with st.expander("Raw logits & weight stats"):
            logits = sm_clf._logits(v_std)[0]
            st.dataframe(pd.DataFrame({"class": CLASSES, "logit": logits}).round(3),
                         use_container_width=True, hide_index=True)
            st.write(f"W shape: {sm_clf.W.shape}, b shape: {sm_clf.b.shape}")

    elif step.startswith("7."):
        st.subheader("7. CNN-from-scratch (raw 64×64 RGB)")
        st.markdown(
            "Three Conv→ReLU→MaxPool blocks → Flatten → FC(128) → FC(6). "
            "Implemented from scratch with custom `im2col` / `col2im`. "
            "Skips the handcrafted features entirely — sees raw pixels."
        )
        cnn_clf = load_cnn()
        from el_nos_el_tany.models.softmax import stable_softmax

        img_chw = img64.transpose(2, 0, 1)[None].astype(np.float32)
        logits = cnn_clf.forward(img_chw)
        proba = stable_softmax(logits)[0]
        pred = int(np.argmax(proba))

        c1, c2 = st.columns([1, 1])
        with c1:
            st.image(img64, caption="Input to CNN (64×64 RGB)", use_container_width=True)
            st.markdown(f"### Predicted: `{CLASSES[pred]}`")
            if true_class:
                ok = "✅ correct" if CLASSES[pred] == true_class else f"❌ true: `{true_class}`"
                st.markdown(f"**{ok}**")
            st.metric("Confidence", f"{proba[pred]:.1%}")
        with c2:
            st.markdown("**Per-class probability**")
            _show_proba(proba, true_class=true_class)

        with st.expander("Raw logits"):
            st.dataframe(pd.DataFrame({"class": CLASSES, "logit": logits[0]}).round(3),
                         use_container_width=True, hide_index=True)

    elif step.startswith("8."):
        st.subheader("8. All predictions side-by-side")
        with st.spinner("Running KNN + Softmax + CNN..."):
            r = predict_all(img96, img64)

        # Build a side-by-side comparison
        rows = []
        for name in ("KNN", "Softmax", "CNN"):
            pred_class = CLASSES[r[name]["pred"]]
            rows.append({
                "model":      name,
                "predicted":  pred_class,
                "confidence": r[name]["proba"][r[name]["pred"]],
                "correct":    "✅" if (true_class and pred_class == true_class) else (
                              "❌" if true_class else "?"),
            })
        df_summary = pd.DataFrame(rows)
        st.dataframe(df_summary.style.format({"confidence": "{:.1%}"}),
                     use_container_width=True, hide_index=True)

        st.markdown("**Per-class probability — all three models**")
        proba_df = pd.DataFrame({
            "class":   CLASSES,
            "KNN":     [float(p) for p in r["KNN"]["proba"]],
            "Softmax": [float(p) for p in r["Softmax"]["proba"]],
            "CNN":     [float(p) for p in r["CNN"]["proba"]],
        })
        if true_class:
            proba_df.insert(1, "match", proba_df["class"].apply(
                lambda c: "✅" if c == true_class else ""))
        st.dataframe(
            proba_df.style.format({"KNN": "{:.3f}", "Softmax": "{:.3f}", "CNN": "{:.3f}"})
                          .background_gradient(subset=["KNN", "Softmax", "CNN"], cmap="Greens"),
            use_container_width=True, hide_index=True,
        )

        st.markdown("**Bar chart per class**")
        chart_df = pd.DataFrame({
            "KNN":     r["KNN"]["proba"],
            "Softmax": r["Softmax"]["proba"],
            "CNN":     r["CNN"]["proba"],
        }, index=CLASSES)
        st.bar_chart(chart_df)


# ===========================================================================
# MINICV STUDIO  (Milestone 1 — original interactive playground)
# ===========================================================================
def minicv_studio():
    st.title("🎨 minicv: Interactive Computer Vision Studio")
    st.markdown("Upload an image and apply the custom library functions in real-time.")

    st.sidebar.header("📂 Image Setup")
    uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

    if uploaded_file is None:
        st.info("👋 Upload an image to start your CV experiments!")
        return

    pil_img = Image.open(uploaded_file)
    img = np.array(pil_img).astype(np.float32)
    if img.max() > 1.0:
        img /= 255.0
    if img.shape[-1] == 4:
        img = img[..., :3]

    st.sidebar.divider()
    st.sidebar.header("🛠 Processing Tools")

    tool_category = st.sidebar.selectbox(
        "Select Category",
        ["Enhancement", "Filters", "Features", "Geometric", "Frequency Domain", "Morphology"]
    )

    processed = img.copy()

    try:
        if tool_category == "Enhancement":
            st.sidebar.subheader("Adjust Intensity")
            gamma = st.sidebar.slider("Gamma Correction", 0.1, 3.0, 1.0)
            processed = gamma_correction(img, gamma)

            hist_action = st.sidebar.radio(
                "Histogram Operations",
                ["None", "Equalize Histogram", "Histogram Matching"]
            )

            if hist_action == "Equalize Histogram":
                try:
                    gray_img = rgb2gray(processed)
                    img_255 = (gray_img * 255).astype(np.uint8)
                    eq_img = equalize_histogram(img_255)
                    processed = eq_img.astype(np.float32) / 255.0
                except Exception as e:
                    st.error(f"Equalization Error: {e}")

            elif hist_action == "Histogram Matching":
                st.sidebar.markdown("---")
                st.sidebar.subheader("Reference Image")
                ref_file = st.sidebar.file_uploader("Upload reference...", type=["jpg", "png", "jpeg"], key="ref_uploader")
                if ref_file is not None:
                    ref_pil = Image.open(ref_file)
                    ref_img = np.array(ref_pil).astype(np.float32)
                    if ref_img.max() > 1.0:
                        ref_img /= 255.0
                    if ref_img.shape[-1] == 4:
                        ref_img = ref_img[..., :3]
                    st.sidebar.image(ref_img, caption="Reference", use_container_width=True)
                    try:
                        src = processed
                        ref = ref_img
                        if src.ndim == 3 and ref.ndim == 2:
                            src = rgb2gray(src)
                        elif src.ndim == 2 and ref.ndim == 3:
                            ref = rgb2gray(ref)
                        matched_img = histogram_matching(src, ref)
                        processed = matched_img.astype(np.float32) / 255.0
                    except Exception as e:
                        st.error(f"Matching Error: {e}")
                else:
                    st.sidebar.info("Upload a reference image to apply matching.")

            st.sidebar.markdown("---")
            if st.sidebar.checkbox("Show Histogram"):
                counts, _ = np.histogram(processed.flatten(), bins=256, range=(0, 1))
                st.sidebar.bar_chart(counts)

        elif tool_category == "Filters":
            st.sidebar.subheader("Spatial Filtering")
            filter_type = st.sidebar.radio(
                "Type",
                ["Gaussian Blur", "Median Blur", "Sobel Edges", "Otsu Threshold", "Adaptive Threshold"]
            )
            if filter_type == "Gaussian Blur":
                k_size = st.sidebar.slider("Kernel Size", 3, 15, 5, step=2)
                sigma = st.sidebar.slider("Sigma", 0.1, 10.0, 1.0)
                processed = gaussian_filter(img, k_size, sigma)
            elif filter_type == "Median Blur":
                k_size = st.sidebar.slider("Kernel Size", 3, 13, 3, step=2)
                processed = median_filter(rgb2gray(img), k_size)
            elif filter_type == "Sobel Edges":
                mag, _ = sobel_filter(rgb2gray(img))
                processed = mag
            elif filter_type == "Otsu Threshold":
                gray = rgb2gray(img)
                thresh = threshold_otsu(gray)
                processed = (gray > thresh).astype(np.float32)
            elif filter_type == "Adaptive Threshold":
                st.sidebar.info("Best for uneven lighting/shadows")
                block_size = st.sidebar.slider("Block Size", 3, 31, 11, step=2)
                offset_val = st.sidebar.slider("Offset (C)", -20, 20, 2)
                t_offset = offset_val / 255.0
                res = threshold_adaptive(rgb2gray(img), block_size, t_offset)
                processed = res.astype(np.float32)

        elif tool_category == "Features":
            st.sidebar.subheader("Feature Detection")
            feat_type = st.sidebar.radio("Algorithm", ["Harris Corners", "Canny Edges"])
            if feat_type == "Harris Corners":
                k_sens = st.sidebar.slider("Sensitivity (k)", 0.01, 0.2, 0.04)
                thresh_ratio = st.sidebar.select_slider(
                    "Threshold Sensitivity",
                    options=[0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1],
                    value=0.01
                )
                edge_grid = harris_corner_detector(img, k=k_sens, threshold_ratio=thresh_ratio)
                processed = np.zeros_like(img)
                visible_mask = binary_dilation(edge_grid > 0, iterations=2)
                processed[visible_mask] = [0.0, 1.0, 0.0]
            elif feat_type == "Canny Edges":
                low_val = st.sidebar.slider("Low Threshold", 0, 100, 10)
                high_val = st.sidebar.slider("High Threshold", 0, 200, 30)
                t_low, t_high = low_val / 255.0, high_val / 255.0
                edges = canny_edge_detector(rgb2gray(img), t_low, t_high)
                if isinstance(edges, tuple): edges = edges[0]
                processed = np.zeros_like(img)
                processed[edges > 0] = [0.0, 1.0, 0.0]

        elif tool_category == "Geometric":
            st.sidebar.subheader("Spatial Transforms")
            angle = st.sidebar.slider("Rotate (Degrees)", -180, 180, 0)
            scale = st.sidebar.slider("Scale Factor", 0.1, 2.0, 1.0)
            processed = rotate(img, angle)
            if scale != 1.0:
                new_h = int(processed.shape[0] * scale)
                new_w = int(processed.shape[1] * scale)
                processed = resize(processed, (new_h, new_w), method='bilinear')

        elif tool_category == "Frequency Domain":
            st.sidebar.subheader("FFT Filtering")
            filt_type = st.sidebar.radio("Filter Type", ["Low-pass (Blur)", "High-pass (Edges)"])
            cutoff = st.sidebar.slider("Cutoff Frequency", 1, 100, 30)
            gray = rgb2gray(img)
            f_type = 'lowpass' if filt_type == "Low-pass (Blur)" else 'highpass'
            processed = apply_frequency_filter(gray, cutoff=cutoff, filter_type=f_type)
            if st.sidebar.checkbox("Show Spectrum"):
                f_transform = np.fft.fft2(gray)
                f_shift = np.fft.fftshift(f_transform)
                mag_spec = 20 * np.log(np.abs(f_shift) + 1)
                spec_min, spec_max = mag_spec.min(), mag_spec.max()
                if spec_max > spec_min:
                    normalized_spec = (mag_spec - spec_min) / (spec_max - spec_min)
                else:
                    normalized_spec = mag_spec
                st.divider()
                st.subheader("📊 Fourier Frequency Spectrum")
                st.image(normalized_spec, caption="Centred Magnitude Spectrum", use_container_width=True)

        elif tool_category == "Morphology":
            st.sidebar.subheader("Binary Operations")
            morph_type = st.sidebar.radio("Type", ["Erosion (Shrink)", "Dilation (Expand)"])
            k_size = st.sidebar.slider("Kernel Size", 3, 15, 3, step=2)
            gray = rgb2gray(img)
            thresh = threshold_otsu(gray)
            binary = (gray > thresh).astype(np.uint8)
            if morph_type == "Erosion (Shrink)":
                res = erode(binary, k_size)
            else:
                res = dilate(binary, k_size)
            processed = res.astype(np.float32)

    except Exception as e:
        st.error(f"⚠️ Error: {e}")
        processed = img

    col1, col2 = st.columns(2)
    with col1:
        st.subheader("Original Image")
        st.image(img, use_container_width=True, clamp=True)
        st.caption(f"Original Dimensions: {img.shape[1]}x{img.shape[0]}")
    with col2:
        st.subheader("Processed Output")
        st.image(processed, use_container_width=False, clamp=True)
        st.write(f"📐 **Processed Dimensions:** {processed.shape[1]}px x {processed.shape[0]}px")


# ===========================================================================
# DISPATCH
# ===========================================================================
if APP_MODE == "minicv Studio":
    minicv_studio()
else:
    m2_predictor()