src/sidebar.py

import os
from pathlib import Path
import streamlit as st

LOCAL = False


DEFAULT_TARGET = "GVHD"
DEFAULT_THRESHOLD = 0.5

SAVED_MODELS_DIR = Path("src/saved_models")


def _init_session_defaults():
    # defaults used by all pages
    st.session_state.setdefault("selected_model", None)
    st.session_state.setdefault("target_choice", DEFAULT_TARGET)   # bulk page uses this
    st.session_state.setdefault("target_col", DEFAULT_TARGET)      # individual/training set this too
    st.session_state.setdefault("threshold", DEFAULT_THRESHOLD)
    st.session_state.setdefault("orig_train_cols", [
        # Raw columns expected from uploaded CSV (NOT engineered features)
        "EPI/ID numbers",
        "Recepient_gender",
        "Recepient_DOB",
        "Recepient_Nationality",
        "Hematological Diagnosis",
        "Date of first diagnosis/BMBx date",
        "Recepient_Blood group before HSCT",
        "Donor_DOB",
        "Donor_gender",
        "D_Blood group",
        "R_HLA_A", "R_HLA_B", "R_HLA_C", "R_HLA_DR", "R_HLA_DQ",
        "D_HLA_A", "D_HLA_B", "D_HLA_C", "D_HLA_DR", "D_HLA_DQ",
        "Number of lines of Rx before HSCT",
        "PreHSCT conditioning regimen+/-ATG+/-TBI",
        "HSCT_date",
        "Source of cells",
        "Donor_relation to recepient",
        "HLA match ratio",
        "First_GVHD prophylaxis",

        # Pro / survival raw columns
        "Last_followup_date",
        "Date_of_death",

        # Pro categorical columns
        "Donor_type",
        "Conditioning_intensity",
        "GVHD_Prophylaxis_Cat",
    ])


def get_model_options():
    models = []

    if LOCAL:
        SAVED_MODELS_DIR.mkdir(exist_ok=True)
        models.extend([p.stem for p in SAVED_MODELS_DIR.glob("*.pkl")])
        return sorted(models, reverse=True)

    from huggingface_hub import HfApi

    repo_id = os.environ.get("HF_REPO_ID", "")
    token = os.environ.get("HF_TOKEN", "")

    if not repo_id or not token:
        return models

    api = HfApi(token=token)
    all_files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")

    parquet_files = [
        Path(f).stem
        for f in all_files
        if f.startswith("models/") and f.endswith(".parquet")
    ]

    # Sort by timestamp prefix (yyMMdd_HHMMSS...)
    def extract_timestamp(name):
        return name.split("_")[0] if "_" in name else ""

    parquet_files.sort(key=extract_timestamp, reverse=True)

    return parquet_files


def sidebar():
    _init_session_defaults()

    st.sidebar.title("GVHD-Intel")
    st.sidebar.subheader("Model")
    options = get_model_options()

    if not options:
        st.sidebar.warning("No trained models found.")
        return
    
    # If no model selected OR selected model not in list → pick latest (first in sorted list)
    if (
        "selected_model" not in st.session_state
        or st.session_state.selected_model not in options
    ):
        st.session_state.selected_model = options[0]  # latest model
    
    st.session_state.selected_model = st.sidebar.selectbox(
        "Model",
        options=options,
        index=options.index(st.session_state.selected_model),
    )
    

    

    # Optional: hide advanced controls behind toggle (recommended)
    show_advanced = st.sidebar.toggle("Show advanced settings", value=False)

    if show_advanced:
        st.sidebar.subheader("Target")
        st.session_state.target_choice = st.sidebar.radio(
            "Choose target:",
            options=["GVHD", "Acute GVHD(<100 days)", "Chronic GVHD>100 days"],
            index=["GVHD", "Acute GVHD(<100 days)", "Chronic GVHD>100 days"].index(
                st.session_state.get("target_choice", DEFAULT_TARGET)
            )
        )

    st.sidebar.subheader("Decision threshold")
    st.session_state.threshold = st.sidebar.slider(
        "Positive threshold",
        min_value=0.05,
        max_value=0.95,
        value=float(st.session_state.get("threshold", DEFAULT_THRESHOLD)),
        step=0.05,
    )