app.py

import os
import re
import json
import random
import hashlib
import hmac
import uuid
import shutil
import tempfile
import numpy as np
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Any, Optional, Tuple, Union

import gradio as gr
from gradio_modal import Modal
import pandas as pd
from huggingface_hub import InferenceClient
from huggingface_hub import CommitScheduler,HfApi,hf_hub_download
from datetime import datetime
#from google import genai
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold

# LangChain imports for retrieval
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

# =========================
# Config via env
# =========================
HF_ENDPOINT_URL = os.environ.get("HF_ENDPOINT_URL", "").strip()
HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "google/gemma-2-2b-it").strip()
HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", os.environ.get("GEMINI_API_KEY", "")).strip()
HF_DATASET_ID = os.environ.get("HF_DATASET_ID", "").strip()
ANON_ID_SECRET = os.environ.get("ANON_ID_SECRET", "").strip()
DEMO_MODE = os.environ.get("DEMO_MODE", "1").strip().lower() not in {"0", "false", "no"}
CONSENT_VERSION = os.environ.get("CONSENT_VERSION", "EPSIC-DEMO-CONSENT-v1").strip()
DEMO_RECORD_PREFIX = os.environ.get("DEMO_RECORD_PREFIX", "demo_records").strip().strip("/") or "demo_records"
if GOOGLE_API_KEY:
    genai.configure(api_key=GOOGLE_API_KEY)
else:
    print("[WARN] GOOGLE_API_KEY/GEMINI_API_KEY is not set. Add it as a Space secret to enable simulated-patient responses.")
if not HF_API_TOKEN:
    print("[WARN] HF_API_TOKEN is not set. Add it as a Space secret to enable demo feedback uploads.")
if DEMO_MODE:
    print("[INFO] Running in DEMO_MODE. Login is not required and feedback saving is consent-gated.")

HERE = os.path.dirname(os.path.abspath(__file__))
DISTRIBUTIONS_PATH = os.path.join(HERE, "distributions.json")
EPSIC_PATH = os.path.join(HERE, "EPSIC.json")
RULES_PATH = os.path.join(HERE, "rules.json")
DSM5_MAPPING_PATH = os.path.join(HERE, "dsm5_mapping.json")

# --- Retrieval Configuration ---
FAISS_INDEX_BASE_PATH = os.path.join(HERE, "faiss_index")
EMBEDDING_MODEL = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
KNOWLEDGE_BASE_MODES = {
    "None": None,
    "DSM-5": "dsm5_only",
    "ICD-11": "icd11_only",
    "DSM-5 + ICD-11": "dsm5_icd11",
}
# --- End Retrieval Configuration ---

# Dataset folder + main file
DATASET_DIR = os.path.join(HERE, "data")
os.makedirs(DATASET_DIR, exist_ok=True)
MAIN_DATASET_PATH = os.path.join(DATASET_DIR, "epsic_demo_feedback.csv" if DEMO_MODE else "epsic_dataset_main.csv")
MAIN_JSON_PATH = os.path.join(DATASET_DIR, "epsic_demo_conversations.json" if DEMO_MODE else "epsic_conversations_main.json")
DEFAULT_MISSING = "Information missing/screen question not asked"

# =========================
# Load EPSIC & Rules
# =========================
with open(DISTRIBUTIONS_PATH, "r", encoding="utf-8") as fh:
    DISTRIBUTIONS = json.load(fh)
with open(EPSIC_PATH, "r", encoding="utf-8") as fh:
    EPSIC = json.load(fh)

with open(RULES_PATH, "r", encoding="utf-8") as fh:
    DIAGNOSTIC_RULES = json.load(fh)

with open(DSM5_MAPPING_PATH, "r", encoding="utf-8") as fh:
    DSM5_MAPPING = json.load(fh)

ALLOWED_MODULES = [
    "SEPARATION ANXIETY DISORDER",
    "SOCIAL ANXIETY DISORDER",
    "PANIC DISORDER",
    "AGORAPHOBIA",
    "GENERALISED ANXIETY DISORDER",
    "OBSESSIVE–COMPULSIVE DISORDER",
    "PTSD/COMPLEX PTSD",
    "MAJOR DEPRESSIVE EPISODE",
    "SCREEN FOR PSYCHOTIC DISORDERS",
    "ADHD",
    "AUTISM SPECTRUM DISORDER"
]
GENDER_DISTRIBUTION = {"female": 0.686, "male": 0.314}
AGE_DISTRIBUTION = {"mean":14.86, "sd":1.385, "min":13, "max":18}
SEVERITY_LEVELS = ["none","moderate", "severe"]
PPROF_DISORDERS = [
    "SEPARATION ANXIETY DISORDER", "SOCIAL ANXIETY DISORDER", "PANIC DISORDER",
    "AGORAPHOBIA", "GENERALISED ANXIETY DISORDER", "OBSESSIVE–COMPULSIVE DISORDER",
    "PTSD/COMPLEX PTSD", "MAJOR DEPRESSIVE EPISODE", "SCREEN FOR PSYCHOTIC DISORDERS",
    "ADHD", "ADD", "AUTISM SPECTRUM DISORDER"
]
PREGEN_PROFILES_PATH = os.path.join(HERE, "simulated_patient_profiles.npy")

print("Loading pre-generated patient profiles...")
PREGEN_PROFILES = np.load(PREGEN_PROFILES_PATH)
print(f"Successfully loaded {len(PREGEN_PROFILES)} profiles.")

# =========================
# Data structures
# =========================
@dataclass
class SelectionRecord:
    item_number: Union[int, float, str, None]
    choice: Optional[str] = None
    choices: Optional[List[str]] = None
    other_choices: Optional[List[str]] = None
    # --- NEW FIELD ---
    realism_rating: Optional[int] = None

@dataclass
class EPSICItem:
    item_number: Union[int, float, str, None]
    assessment: str
    probing_questions: List[str]
    choices: Optional[Dict[str, Any]] = None
    other_choices: Optional[Dict[str, Any]] = None
    is_screening: bool = False


@dataclass
class SessionState:
    module_name: str
    severity: str
    age: int
    sex: str
    items: List[EPSICItem] = field(default_factory=list)
    idx: int = 0
    agent: Optional["HFInferencePatient"] = None
    selections: Dict[str, SelectionRecord] = field(default_factory=dict)
    transcript: List[Tuple[str, str]] = field(default_factory=list)
    asked_probes: set = field(default_factory=set)
    module_comments: str = ""
    module_rating: int = 0  # <-- NEW: Stores the overall module realism score


@dataclass
class DatasetState:
    columns: List[str]
    df: pd.DataFrame
    # --- MODIFIED: Added User State Fields ---
    clinician_id: str = "Unknown"
    profession: str = "Not specified"
    yrs_diag_exp: str = "Not specified"
    yrs_prof_exp: str = "Not specified"
    consent_given: bool = False
    consent_timestamp_utc: str = ""
    consent_version: str = CONSENT_VERSION
    feedback_saving_enabled: bool = False
    demo_session_id: str = ""
    # -----------------------------------------
    current_patient_data: Dict[str, Any] = field(default_factory=dict)
    patient_id: Optional[str] = None
    patient_age: Optional[int] = None
    patient_sex: Optional[str] = None
    patient_profile: Dict[str, str] = field(default_factory=dict)
    patient_adhd_subtype: str = "none" 
    completed_sessions: Dict[str, SessionState] = field(default_factory=dict)
    current_module_index: int = -1
    diag_yes_modules: set = field(default_factory=set)
    last_saved_path: Optional[str] = None

# --- HF Hub persistence helpers ---
def upload_to_hub(file_path: str, path_in_repo: Optional[str] = None):
    """Uploads a file to the configured HF Dataset repository."""
    if not HF_API_TOKEN or not HF_DATASET_ID:
        print("[WARN] Cannot upload to Hub: Token or Dataset ID missing.")
        return

    api = HfApi(token=HF_API_TOKEN)
    repo_path = path_in_repo or os.path.basename(file_path)
    print(f"Uploading {file_path} to {HF_DATASET_ID}:{repo_path}...")

    try:
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=repo_path,
            repo_id=HF_DATASET_ID,
            repo_type="dataset"
        )
        print(f"Successfully uploaded to {repo_path}")
    except Exception as e:
        print(f"[ERROR] Upload failed for {repo_path}: {e}")


def upload_json_to_hub(data: Dict[str, Any], path_in_repo: str):
    if not HF_API_TOKEN or not HF_DATASET_ID:
        print("[WARN] Cannot upload JSON to Hub: Token or Dataset ID missing.")
        return
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as tmp:
            json.dump(data, tmp, ensure_ascii=False, indent=2)
            tmp_path = tmp.name
        upload_to_hub(tmp_path, path_in_repo=path_in_repo)
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
            except OSError:
                pass


def download_from_hub(path_in_repo: str, local_path: Optional[str] = None, force_download: bool = True) -> Optional[str]:
    if not HF_API_TOKEN or not HF_DATASET_ID:
        return None
    try:
        downloaded = hf_hub_download(
            repo_id=HF_DATASET_ID,
            repo_type="dataset",
            filename=path_in_repo,
            token=HF_API_TOKEN,
            force_download=force_download,
        )
        if local_path:
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            shutil.copyfile(downloaded, local_path)
            return local_path
        return downloaded
    except Exception as e:
        print(f"[INFO] Could not download {path_in_repo} from Hub: {e}")
        return None


def load_json_from_hub(path_in_repo: str) -> Optional[Dict[str, Any]]:
    local_path = download_from_hub(path_in_repo, force_download=True)
    if not local_path or not os.path.exists(local_path):
        return None
    try:
        with open(local_path, "r", encoding="utf-8") as fh:
            return json.load(fh)
    except Exception as e:
        print(f"[WARN] Failed to load JSON from Hub path {path_in_repo}: {e}")
        return None


def sync_main_artifacts_from_hub():
    for repo_name, local_path in [
        (os.path.basename(MAIN_DATASET_PATH), MAIN_DATASET_PATH),
        (os.path.basename(MAIN_JSON_PATH), MAIN_JSON_PATH),
        (os.path.basename(SUS_DATASET_PATH), SUS_DATASET_PATH),
    ]:
        download_from_hub(repo_name, local_path=local_path, force_download=True)


def resolve_space_username(profile: Optional[gr.OAuthProfile] = None, request: Optional[gr.Request] = None) -> Optional[str]:
    username = None
    if profile is not None:
        username = getattr(profile, "username", None) or getattr(profile, "preferred_username", None)
    if not username and request is not None:
        username = getattr(request, "username", None)
    if not username:
        return None
    return str(username).strip().lower()


def make_anonymous_clinician_id(
    profile: Optional[gr.OAuthProfile] = None,
    request: Optional[gr.Request] = None,
    fallback_identifier: Optional[str] = None,
) -> Optional[str]:
    base_identifier = resolve_space_username(profile, request) or (fallback_identifier.strip().lower() if fallback_identifier else None)
    if not base_identifier:
        return None
    secret = ANON_ID_SECRET or f"insecure::{HF_DATASET_ID or 'local-space'}"
    digest = hmac.new(secret.encode("utf-8"), base_identifier.encode("utf-8"), hashlib.sha256).hexdigest()
    return f"CLN-{digest[:12].upper()}"




def make_demo_session_id() -> str:
    """Create a non-login, session-scoped identifier for symposium demo feedback."""
    timestamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
    suffix = uuid.uuid4().hex[:8].upper()
    return f"DEMO-{timestamp}-{suffix}"


def normalize_optional_background(value: Any, default: str = "Not specified") -> str:
    if value is None:
        return default
    text = str(value).strip()
    if text in {"", "None", "nan"}:
        return default
    try:
        if isinstance(value, float) and np.isnan(value):
            return default
    except Exception:
        pass
    return text


def optional_number_value(value: Any) -> Optional[int]:
    try:
        text = str(value).strip()
        if text in {"", "None", "Not specified", "nan"}:
            return None
        return int(float(text))
    except Exception:
        return None


def apply_demo_feedback_choice(ds: DatasetState, feedback_choice: Optional[str]):
    """Set consent/saving flags from the attendee's choice on the landing page."""
    wants_to_save = str(feedback_choice or "").startswith("I consent")
    ds.feedback_saving_enabled = bool(wants_to_save)
    ds.consent_given = bool(wants_to_save)
    ds.consent_version = CONSENT_VERSION
    if wants_to_save and not ds.consent_timestamp_utc:
        ds.consent_timestamp_utc = datetime.utcnow().isoformat()
    if not wants_to_save:
        ds.consent_timestamp_utc = ""
    if not ds.demo_session_id:
        ds.demo_session_id = make_demo_session_id()
    ds.clinician_id = ds.demo_session_id
    return ds


def feedback_not_saved_message() -> str:
    return "No feedback was saved because this session was started without consent to contribute demo feedback."

def checkpoint_repo_path(anon_id: str, patient_id: str) -> str:
    return f"{DEMO_RECORD_PREFIX}/active_sessions/{anon_id}/{patient_id}.json" if DEMO_MODE else f"active_sessions/{anon_id}/{patient_id}.json"


def active_pointer_repo_path(anon_id: str) -> str:
    return f"{DEMO_RECORD_PREFIX}/active_sessions/{anon_id}/active.json" if DEMO_MODE else f"active_sessions/{anon_id}/active.json"


def patient_csv_repo_path(patient_id: str) -> str:
    return f"{DEMO_RECORD_PREFIX}/records/{patient_id}.csv" if DEMO_MODE else f"records/{patient_id}.csv"


def patient_json_repo_path(patient_id: str) -> str:
    return f"{DEMO_RECORD_PREFIX}/records/{patient_id}.json" if DEMO_MODE else f"records/{patient_id}.json"

# =========================
# Retrieval Logic
# =========================
def load_retriever(mode: str) -> Optional[Any]:
    """Load the appropriate FAISS vector store retriever for the given mode."""
    index_name = KNOWLEDGE_BASE_MODES.get(mode)
    if not index_name:
        print("No knowledge base selected.")
        return None

    index_path = os.path.join(FAISS_INDEX_BASE_PATH, index_name)
    if not os.path.exists(index_path):
        print(f"[WARN] Knowledge base not found at path: {index_path}")
        return None
    
    try:
        faiss_store = FAISS.load_local(
            index_path, 
            EMBEDDING_MODEL, 
            allow_dangerous_deserialization=True
        )
        print(f"Successfully loaded knowledge base: {mode}")
        return faiss_store.as_retriever(search_kwargs={'k': 2})
    except Exception as e:
        print(f"[ERROR] Could not load knowledge base '{mode}': {e}")
        return None

# =========================
# EPSIC utilities
# =========================
def _collect_items(module_obj: Dict[str, Any]) -> List[EPSICItem]:
    items: List[EPSICItem] = []
    scr = module_obj.get("screening")
    if scr:
        items.append(EPSICItem(
            item_number=scr.get("item_number", "S1"),
            assessment=scr.get("assessment", "") or "",
            probing_questions=scr.get("probing questions", []) or [],
            choices=scr.get("choices"),
            other_choices=scr.get("Other_choices"),
            is_screening=True,
        ))
    assessments = module_obj.get("assessments", []) or []
    for blk in assessments:
        if isinstance(blk, dict) and "section" in blk and "items" in blk:
            for it in blk.get("items", []) or []:
                items.append(EPSICItem(
                    item_number=it.get("item_number"),
                    assessment=it.get("assessment", "") or "",
                    probing_questions=it.get("probing questions", []) or [],
                    choices=it.get("choices"),
                    other_choices=None,
                    is_screening=False,
                ))
        else:
            items.append(EPSICItem(
                item_number=blk.get("item_number"),
                assessment=blk.get("assessment", "") or "",
                probing_questions=blk.get("probing questions", []) or [],
                choices=blk.get("choices"),
                other_choices=None,
                is_screening=False,
            ))
    return items

def get_module_data(module_name: str) -> Dict[str, Any]:
    for m in EPSIC:
        if m["module"] == module_name:
            return m
    raise ValueError(f"Module not found: {module_name}")

def get_items_for_module(module_name: str) -> List[EPSICItem]:
    return _collect_items(get_module_data(module_name))

def module_prefix(module_name: str) -> str:
    words = [w for w in re.split(r'[\s\–/]+', module_name) if w.strip()]
    if words and words[-1].upper().startswith("DISORDER"):
        words = words[:-1]
    segs = [w[:3].upper() for w in words]
    return "_".join(segs)

def item_code(mod_prefix: str, item: EPSICItem) -> str:
    num = str(item.item_number) if item.item_number is not None else "NA"
    num = num.replace(".", "_").replace("-", "_").upper()
    code = f"{mod_prefix}_{num}"
    if item.is_screening:
        code = f"{code}_SCREEN"
    return code

# Precompute schema
MODULE_PREFIX = {m: module_prefix(m) for m in ALLOWED_MODULES}
MODULE_ITEMS: Dict[str, List[EPSICItem]] = {m: get_items_for_module(m) for m in ALLOWED_MODULES}
MODULE_ITEM_CODES: Dict[str, List[str]] = {m: [item_code(MODULE_PREFIX[m], it) for it in MODULE_ITEMS[m]] for m in ALLOWED_MODULES}
ALL_ITEM_COLUMNS: List[str] = [code for m in ALLOWED_MODULES for code in MODULE_ITEM_CODES[m]]
RATING_COLUMNS: List[str] = [f"RAT_{code}" for code in ALL_ITEM_COLUMNS]
MODULE_RATING_COLUMNS: List[str] = [f"MOD_RAT_{MODULE_PREFIX[m]}" for m in ALLOWED_MODULES]
COMMENT_COLUMNS: List[str] = [f"COM_{MODULE_PREFIX[m]}" for m in ALLOWED_MODULES]
DIAG_COLUMNS: List[str] = [f"DIA_{MODULE_PREFIX[m]}" for m in ALLOWED_MODULES]
DEMO_COLUMNS = [
    "DemoSessionID", "ClinicianID", "ConsentGiven", "ConsentTimestampUTC", "ConsentVersion",
    "Profession", "YrsDiagExp", "YrsProfExp"
]
MASTER_COLUMNS: List[str] = DEMO_COLUMNS + ["PatientID", "Sex", "Age"] + ALL_ITEM_COLUMNS + RATING_COLUMNS + MODULE_RATING_COLUMNS + COMMENT_COLUMNS + DIAG_COLUMNS
# Path for SUS Validation Data
SUS_DATASET_PATH = os.path.join(DATASET_DIR, "epsic_demo_sus_results.csv" if DEMO_MODE else "validation_sus_results.csv")

#severity display mapping
SEVERITY_DISPLAY_MAP = {
    "none": "No Symptoms",
    "moderate": "Subclinical Symptoms",
    "severe": "Clinical Symptoms"
}
# =========================
# Dataset I/O + ID helpers
# =========================
def load_or_init_main_dataset() -> pd.DataFrame:
    sync_main_artifacts_from_hub()
    if os.path.exists(MAIN_DATASET_PATH):
        try:
            df = pd.read_csv(MAIN_DATASET_PATH)
            for col in MASTER_COLUMNS:
                if col not in df.columns:
                    if col.startswith("DIA_"):
                        df[col] = "No"
                    else:
                        df[col] = ""
            df = df[MASTER_COLUMNS]
        except Exception as e:
            print(f"Error loading main dataset: {e}. Re-initializing.")
            df = pd.DataFrame(columns=MASTER_COLUMNS)
    else:
        df = pd.DataFrame(columns=MASTER_COLUMNS)
    return df

def scan_all_patient_ids() -> set:
    ids = set()
    for fn in os.listdir(DATASET_DIR):
        if fn.lower().endswith(".csv"):
            try:
                path = os.path.join(DATASET_DIR, fn)
                df = pd.read_csv(path, usecols=["PatientID"])
                ids.update(set(str(x) for x in df["PatientID"].dropna().astype(str)))
            except Exception:
                continue
    return ids



def sanitize_filename(name: str) -> str:
    s = re.sub(r"[^A-Za-z0-9_\-]", "_", name.strip())
    return s or "dataset"

def unique_dataset_path(preferred_name: Optional[str] = None) -> str:
    if preferred_name:
        base = sanitize_filename(preferred_name)
    else:
        base = f"epsic_dataset_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}"
    path = os.path.join(DATASET_DIR, f"{base}.csv")
    if not os.path.exists(path):
        return path
    k = 1
    while True:
        path_try = os.path.join(DATASET_DIR, f"{base}_{k}.csv")
        if not os.path.exists(path_try):
            return path_try
        k += 1

# =========================
# HF Inference patient (chat only)
# =========================
class HFInferencePatient:
    def __init__(self, module_data: Dict[str, Any], severity: str, age: int, sex: str, 
                 retriever: Optional[Any] = None, diagnostic_rule: Optional[Dict] = None, 
                 patient_profile: Dict[str, str] = {}, 
                 profile_adhd_subtype: Optional[str] = None,    # <-- This is the profile's subtype
                 current_module_adhd_subtype: Optional[str] = None): # <-- This is for the current module
        
        self.module_name = module_data["module"]
        self.info = module_data.get("information", "")
        self.severity = severity
        self.age = age
        self.sex = sex
        self.retriever = retriever
        self.diagnostic_rule = diagnostic_rule
        self.patient_profile = patient_profile
        self.profile_adhd_subtype = profile_adhd_subtype           # <-- Store profile subtype
        self.current_module_adhd_subtype = current_module_adhd_subtype # <-- Store current module subtype
        self.persona = self._generate_persona(module_data, severity)

        # --- 1. DEFINE SAFETY SETTINGS ---
        safety_settings = {
            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
        }
        
        # --- 2. DEFINE GENERATION CONFIG ---
        generation_config = {
            "temperature": 0.9,
            "top_p": 0.95,
            "top_k": 40,
        }

        # --- 3. INITIALIZE MODEL & CHAT ---
        try:
            # --- MODIFICATION ---
            # generation_config is passed HERE, in the constructor
            self.model = genai.GenerativeModel(
                "gemini-2.5-flash",
                safety_settings=safety_settings,
                generation_config=generation_config # <-- MOVED TO HERE
            )
            # --- END MODIFICATION ---

        except Exception as e:
            print(f"[ERROR] Could not initialize GenerativeModel: {e}")
            self.model = None

        static_system_prompt = self._get_static_system_prompt()
        
        if self.model:
            # --- MODIFICATION ---
            # generation_config is REMOVED from here
            self.chat_session = self.model.start_chat(
                history=[
                    {'role': 'user', 'parts': [static_system_prompt]},
                    {'role': 'model', 'parts': ["Okay, I understand. I am ready to answer as the patient. I will be consistent with my *full* diagnostic profile, but I will focus my answers on the module I am currently being assessed for. I will wait for your first question."]}
                ]
            )
            # --- END MODIFICATION ---
        else:
            self.chat_session = None

    def _generate_persona(self, module_data: Dict[str, Any], severity: str) -> str:
        # (This method is unchanged)
        persona_parts = []
        if severity == "none":
            return "This patient is generally healthy and does not present with any symptoms for this module."

        if self.module_name == "SPECIFIC PHOBIA" and severity != "none":
            try:
                phobia_choices = module_data["screening"]["Other_choices"]["choices"]["names"]
                chosen_phobia = random.choice(phobia_choices)
                persona_parts.append(f"The patient has a specific phobia of: {chosen_phobia}.")
            except (KeyError, IndexError):
                pass
        
        response_style = random.choice([
            "tends to be very brief and to the point (1-2 sentences, sometimes just one word)", 
            "is hesitant and thoughtful in their responses (might be brief or pause)",
            "is a bit more talkative and will share some examples (but keeps answers to 3-4 sentences at most)"
        ])
        persona_parts.append(f"Their response style is that they {response_style}.")

        return " ".join(persona_parts)

    def _retrieve_context(self) -> str:
        # (This method is unchanged)
        if not self.retriever:
            return ""
        try:
            query = f"{self.module_name} symptoms and diagnostic criteria for a {self.age}-year-old {self.sex}"
            docs = self.retriever.invoke(query)
            if not docs:
                return ""
            context = "\n".join([doc.page_content for doc in docs])
            return f"Retrieved Knowledge Base Context:\n{context}\n\n"
        except Exception as e:
            print(f"[WARN] Could not retrieve context from knowledge base: {e}")
            return ""

    def _get_static_system_prompt(self) -> str:
        """Generates the static system prompt, fixed for the entire module session."""
        
        sev_style = {
            "moderate": (
                "Your symptoms are moderate. You have noticeable symptoms, but they are *not* severe or frequent enough to meet the full diagnostic criteria. You must present as a 'subclinical' case."
                " You MUST provide examples of symptoms when asked."
            ),
            "severe": "Your symptoms are severe, causing significant distress and impairment. They substantially interfere with your daily functioning and quality of life. You meet the full diagnostic criteria.",
            "none": "You have NO symptoms of this disorder. You are generally healthy and functioning well in this area."
        }.get(self.severity)

        # --- This part uses current_module_adhd_subtype ---
        special_instructions = ""
        if self.severity == "moderate":
            special_instructions = (
                "IMPORTANT: As a 'moderate' case, you must *not* meet the full diagnostic criteria. "
                "When asked about symptoms, you should confirm some, but deny or minimize others, especially those related to duration, frequency, or severe impairment."
            )

        # Use the subtype specific to the *current module*
        if self.module_name == "ADHD" and self.current_module_adhd_subtype and self.current_module_adhd_subtype != "none":
            subtype_text = self.current_module_adhd_subtype.replace("_", " ").title()
            special_instructions += f"\nFor this ADHD module, you are specifically simulating the **{subtype_text} Presentation**."

        rule_text = ""
        if self.diagnostic_rule:
            rule_text = f"The overall diagnostic rule for {self.module_name} is: \"{self.diagnostic_rule.get('rule', 'Not specified.')}\"\n"

        # --- This part uses profile_adhd_subtype ---
        profile_summary_parts = [
            "--- Your Full Diagnostic Profile (Background Context) ---",
            "This is your complete psychological profile. Your answers should be consistent with this full picture. For example, if you have severe depression, it should color your answers even when being asked about anxiety."
        ]
        
        if not self.patient_profile:
             profile_summary_parts.append("- No profile provided.")
        else:
            for mod, sev in self.patient_profile.items():
                sev_text = sev.capitalize()
                if mod == self.module_name:
                    sev_text = f"**{sev_text} (This is the module you are being assessed for *now*)**"
                
                # Use the subtype from the *overall profile*
                if mod == "ADHD" and self.patient_profile.get("ADHD", "none") != "none" and self.profile_adhd_subtype and self.profile_adhd_subtype != "none":
                     sev_text += f" ({self.profile_adhd_subtype.replace('_', ' ')} type)"

                profile_summary_parts.append(f"- {mod}: {sev_text}")
        
        full_profile_text = "\n".join(profile_summary_parts)

        # --- Assemble the final prompt ---
        system_prompt = (
            f"You are simulating a {self.age}-year-old {self.sex} patient during a clinician-led EPSIC interview.\n"
            f"Your response style is that you {self.persona}\n\n"
            f"{full_profile_text}\n\n"
            f"--- Instructions for the CURRENT Module ({self.module_name}) ---\n"
            f"You are now being assessed for **{self.module_name}**.\n"
            f"Your simulated severity level for *this* module is: **{self.severity.capitalize()}**.\n"
            f"Severity Definition: {sev_style}\n"
            f"{special_instructions}\n"
            f"{rule_text}"
            f"You MUST adhere to your full profile and *especially* the instructions for the current module. Answer ONLY as the patient, in the first person. Keep your response natural and conversational."
            f"Your memory of this conversation is preserved. Be consistent with your previous answers."
        )
        return system_prompt

    def _build_user_turn(self, question: str, item: EPSICItem) -> str:
        """Builds the content for the user's turn, including dynamic context."""
        
        dynamic_instructions = []
        
        if self.severity == "moderate" and item.is_screening:
            dynamic_instructions.append(
                "Note for this item: This is a screening question. You MUST answer affirmatively (e.g., start with 'Yes, sometimes I do...', 'I guess, ...') to proceed with the assessment. You should give the impression that the symptoms exist, but not to a sufficient extent for diagnosis."
            )

        dsm_code = DSM5_MAPPING.get(self.module_name, {}).get(str(item.item_number))
        if dsm_code:
            dynamic_instructions.append(f"This question relates to DSM-5 Criterion {dsm_code} for this disorder.")

        instruction_text = ""
        if dynamic_instructions:
            instruction_text = "--- Context for this specific item ---\n" + "\n".join(dynamic_instructions) + "\n\n"

        epsic_context = f"Clinical context on {self.module_name}:\n{self.info}\n\n" if self.info else ""
        retrieved_context = self._retrieve_context()
        
        item_context = (
            f"Background for Clinician (do not repeat):\n{epsic_context}{retrieved_context}"
            f"Current EPSIC Item Being Assessed:\n{item.assessment}\n\n"
        )

        user_turn_content = (
            f"{instruction_text}"
            f"{item_context}"
            f"The clinician now asks you: \"{question}\""
        )
        
        print(f"[DEBUG] User turn content:\n{user_turn_content}")
        
        return user_turn_content

    def respond(self, question: str, item: EPSICItem) -> str:
        user_turn_content = self._build_user_turn(question, item)

        if not GOOGLE_API_KEY:
            return "[Agent Response Disabled: GOOGLE_API_KEY/GEMINI_API_KEY not set]"
        
        if not self.chat_session:
             return "[Agent Error: Chat session not initialized]"
            
        try:
            response = self.chat_session.send_message(user_turn_content)
            response_text = response.text.strip()
            return response_text
        
        except Exception as e:
            print(f"[ERROR] LLM generation failed: {e}")
            try:
                print("[DEBUG] Chat History on Error:")
                for msg in self.chat_session.history:
                    role = msg.role
                    text = msg.parts[0].text if msg.parts else "[No Text]"
                    print(f"  [{role}] {text[:150]}...")
            except Exception as de:
                print(f"[ERROR] Could not print debug history: {de}")
            
            return "[Agent Error: Could not generate a response]"
        """
        if not GOOGLE_API_KEY:
            return "[Agent Response Disabled: GOOGLE_API_KEY/GEMINI_API_KEY not set]"
        try:
            result = self.client.chat.completions.create(
                model=os.getenv("HF_MODEL_ID", None),
                messages=[{"role": "user", "content": prompt}],
                max_tokens=256,
                temperature=0.8,
                top_p=0.95,
            )
            msg = result.choices[0].message
            content = msg.get("content", "")
            if isinstance(content, list):
                text = "".join(part.get("text", "") for part in content)
            else:
                text = content or ""
            
            response_text = text.strip()
            self.memory[question] = response_text
            return response_text
        
        except Exception as e:
            print(f"[ERROR] LLM generation failed: {e}")
            return "[Agent Error: Could not generate a response]"
        """

# =========================
# Module/session callbacks
# =========================
def _choices_for_item(st: SessionState, item: Optional[EPSICItem]):
    single_opts, multi_opts, other_opts = [], [], []
    selected_single, selected_multi, selected_other = None, [], []

    if not item:
        return single_opts, multi_opts, other_opts, selected_single, selected_multi, selected_other

    def process_choices(choices_obj):
        s_opts, m_opts = [], []
        if isinstance(choices_obj, list):
            for choice_block in choices_obj:
                names = choice_block.get("names", [])
                if choice_block.get("type", "").startswith("single"):
                    s_opts.extend(names)
                else:
                    m_opts.extend(names)
        elif isinstance(choices_obj, dict):
            names = choices_obj.get("names", [])
            if choices_obj.get("type", "").startswith("single"):
                s_opts = names
            else:
                m_opts = names
        return s_opts, m_opts

    if item.choices:
        single_opts, multi_opts = process_choices(item.choices)
    if item.other_choices and "choices" in item.other_choices:
        _, other_opts = process_choices(item.other_choices["choices"])

    key = str(item.item_number)
    if key in st.selections:
        rec = st.selections[key]
        selected_single = rec.choice
        selected_multi = rec.choices or []
        selected_other = rec.other_choices or []

    return single_opts, multi_opts, other_opts, selected_single, selected_multi, selected_other

def _current_item(st: SessionState) -> Optional[EPSICItem]:
    if not st or not st.items:
        return None
    return st.items[st.idx]

def _format_item_header(item: EPSICItem, module_name: str) -> str:
    if not item:
        return ""
    
    item_num_str = str(item.item_number)
    desc_parts = []
    if item.is_screening:
        desc_parts.append("Screening Item")

    dsm_code = DSM5_MAPPING.get(module_name, {}).get(item_num_str)
    if not dsm_code and '.' in item_num_str:
        dsm_code = DSM5_MAPPING.get(module_name, {}).get(item_num_str.replace('.', ''))

    if dsm_code:
        desc_parts.append(f"DSM-5 Criterion {dsm_code}")

    description = " / ".join(desc_parts)
    
    header = f"### Current EPSI-C Item: {item_num_str}"
    if description:
        header += f" ({description})"
    
    return header + f"\n\n**Assessment:** {item.assessment}"

def _create_module_session(module_name: str, age: int, sex: str, severity: str, 
                         knowledge_base_mode: str, patient_profile: Dict[str, str], 
                         profile_adhd_subtype: Optional[str] = None,    # <-- ADD THIS
                         current_module_adhd_subtype: Optional[str] = None) -> SessionState: # <-- RENAME THIS
    st = SessionState(module_name=module_name, severity=severity, age=age, sex=sex)
    st.items = get_items_for_module(module_name)
    st.idx = 0
    module_data = get_module_data(module_name)
    
    # Fetch the diagnostic rule for the current module
    diagnostic_rule = DIAGNOSTIC_RULES.get(module_name)
    
    retriever = load_retriever(knowledge_base_mode)
    
    # Pass the rule, profile, and subtypes to the agent's constructor
    st.agent = HFInferencePatient(
        module_data, 
        severity, 
        st.age, 
        st.sex, 
        retriever, 
        diagnostic_rule,
        patient_profile=patient_profile,
        profile_adhd_subtype=profile_adhd_subtype, # <-- PASS IT
        current_module_adhd_subtype=current_module_adhd_subtype # <-- PASS IT
    )
    st.transcript = []
    
    return st


def rehydrate_session_for_interview(ds: DatasetState, stored_session: SessionState) -> SessionState:
    current_module_adhd_subtype = ds.patient_adhd_subtype if stored_session.module_name == "ADHD" else None
    profile_adhd_subtype = ds.patient_adhd_subtype
    st = _create_module_session(
        module_name=stored_session.module_name,
        age=ds.patient_age,
        sex=ds.patient_sex,
        severity=stored_session.severity,
        knowledge_base_mode="DSM-5 + ICD-11",
        patient_profile=ds.patient_profile,
        profile_adhd_subtype=profile_adhd_subtype,
        current_module_adhd_subtype=current_module_adhd_subtype,
    )
    st.idx = min(max(int(stored_session.idx), 0), max(len(st.items) - 1, 0)) if st.items else 0
    st.selections = {
        k: (SelectionRecord(**asdict(v)) if isinstance(v, SelectionRecord) else SelectionRecord(**v))
        for k, v in stored_session.selections.items()
    }
    st.transcript = list(stored_session.transcript)
    st.asked_probes = set(stored_session.asked_probes)
    st.module_comments = stored_session.module_comments
    st.module_rating = stored_session.module_rating

    saved_history = getattr(stored_session, "saved_chat_history", None)
    if saved_history and st.agent and st.agent.model:
        try:
            st.agent.chat_session = st.agent.model.start_chat(history=saved_history)
        except Exception as e:
            print(f"[WARN] Could not restore saved chat history for {stored_session.module_name}: {e}")
    return st


def dataset_state_from_checkpoint(payload: Dict[str, Any]) -> Tuple[DatasetState, Optional[SessionState]]:
    ds = init_dataset_state()
    ds.clinician_id = payload.get("clinician_id", ds.clinician_id)
    ds.profession = payload.get("profession", ds.profession)
    ds.yrs_diag_exp = str(payload.get("yrs_diag_exp", ds.yrs_diag_exp))
    ds.yrs_prof_exp = str(payload.get("yrs_prof_exp", ds.yrs_prof_exp))
    ds.consent_given = bool(payload.get("consent_given", ds.consent_given))
    ds.consent_timestamp_utc = payload.get("consent_timestamp_utc", ds.consent_timestamp_utc)
    ds.consent_version = payload.get("consent_version", ds.consent_version)
    ds.feedback_saving_enabled = bool(payload.get("feedback_saving_enabled", ds.feedback_saving_enabled))
    ds.demo_session_id = payload.get("demo_session_id", ds.demo_session_id) or ds.clinician_id
    ds.patient_id = payload.get("patient_id")
    ds.patient_age = payload.get("patient_age")
    ds.patient_sex = payload.get("patient_sex")
    ds.patient_profile = payload.get("patient_profile", {}) or {}
    ds.patient_adhd_subtype = payload.get("patient_adhd_subtype", "none") or "none"
    ds.current_patient_data = payload.get("current_patient_data", {}) or {}
    ds.current_module_index = int(payload.get("current_module_index", -1))
    ds.diag_yes_modules = set(payload.get("diag_yes_modules", []) or [])

    ds.completed_sessions = {}
    for module_name, sess_payload in (payload.get("completed_sessions") or {}).items():
        st = simple_session_from_dict(sess_payload)
        st.saved_chat_history = sess_payload.get("chat_history", [])
        ds.completed_sessions[module_name] = st

    current_session_payload = payload.get("current_session")
    current_session = None
    if current_session_payload:
        current_stored = simple_session_from_dict(current_session_payload)
        current_stored.saved_chat_history = current_session_payload.get("chat_history", [])
        current_session = rehydrate_session_for_interview(ds, current_stored)

    return ds, current_session


def maybe_restore_checkpoint(profile: Optional[gr.OAuthProfile] = None, request: Optional[gr.Request] = None) -> Tuple[DatasetState, Optional[SessionState], Optional[str]]:
    if DEMO_MODE:
        ds = init_dataset_state()
        return ds, None, ds.clinician_id

    anon_id = make_anonymous_clinician_id(profile, request)
    if not anon_id:
        ds = init_dataset_state()
        return ds, None, None

    pointer = load_json_from_hub(active_pointer_repo_path(anon_id))
    if not pointer or pointer.get("status") != "active" or not pointer.get("patient_id"):
        ds = init_dataset_state()
        ds.clinician_id = anon_id
        return ds, None, anon_id

    payload = load_json_from_hub(checkpoint_repo_path(anon_id, pointer["patient_id"]))
    if not payload:
        ds = init_dataset_state()
        ds.clinician_id = anon_id
        return ds, None, anon_id

    ds, current_session = dataset_state_from_checkpoint(payload)
    ds.clinician_id = anon_id
    return ds, current_session, anon_id


def persist_module_session_into_dataset(ds: DatasetState, st: Optional[SessionState], comments: str = "", mod_rating_val: int = 0):
    if not ds or not st:
        return
    st.module_comments = (comments or "").strip()
    st.module_rating = int(mod_rating_val or 0)
    ds.completed_sessions[st.module_name] = st
    mp = MODULE_PREFIX[st.module_name]

    for it in st.items:
        key = str(it.item_number)
        code = item_code(mp, it)
        rating_col = f"RAT_{code}"
        cell = ""
        rec = st.selections.get(key)
        if rec:
            if rec.choice:
                cell = rec.choice
            elif rec.choices:
                cell = "; ".join(rec.choices)
            if rec.other_choices:
                other = "; ".join(rec.other_choices)
                cell = f"{cell} | Other: {other}" if cell else f"Other: {other}"
            ds.current_patient_data[rating_col] = rec.realism_rating if rec.realism_rating is not None else ""
        else:
            ds.current_patient_data.setdefault(rating_col, "")
        ds.current_patient_data[code] = cell

    ds.current_patient_data[f"COM_{mp}"] = st.module_comments
    ds.current_patient_data[f"MOD_RAT_{mp}"] = st.module_rating

def change_module(dataset_state: DatasetState, current_session: Optional[SessionState], comments: str, mod_rating_val: int, nav_type: str, module_name: Optional[str] = None):
    if not dataset_state.patient_id:
        return [dataset_state, None] + [gr.update()] * 23

    if nav_type == "direct" and not module_name:
        return [dataset_state, current_session] + [gr.update()] * 23

    if current_session:
        persist_module_session_into_dataset(dataset_state, current_session, comments, mod_rating_val)
        save_checkpoint(dataset_state, current_session)

    if nav_type == "next":
        dataset_state.current_module_index += 1
    elif nav_type == "prev":
        dataset_state.current_module_index -= 1
    elif nav_type == "direct" and module_name in ALLOWED_MODULES:
        dataset_state.current_module_index = ALLOWED_MODULES.index(module_name)

    dataset_state.current_module_index = int(np.clip(dataset_state.current_module_index, 0, len(ALLOWED_MODULES) - 1))
    target_module_name = ALLOWED_MODULES[dataset_state.current_module_index]
    severity = dataset_state.patient_profile[target_module_name]

    if target_module_name in dataset_state.completed_sessions:
        new_session = rehydrate_session_for_interview(dataset_state, dataset_state.completed_sessions[target_module_name])
    else:
        current_module_adhd_subtype = dataset_state.patient_adhd_subtype if target_module_name == "ADHD" else None
        profile_adhd_subtype = dataset_state.patient_adhd_subtype
        new_session = _create_module_session(
            module_name=target_module_name,
            age=dataset_state.patient_age,
            sex=dataset_state.patient_sex,
            severity=severity,
            knowledge_base_mode="DSM-5 + ICD-11",
            patient_profile=dataset_state.patient_profile,
            profile_adhd_subtype=profile_adhd_subtype,
            current_module_adhd_subtype=current_module_adhd_subtype,
        )

    assessed_count = len(dataset_state.completed_sessions)
    display_severity = SEVERITY_DISPLAY_MAP.get(severity, severity)
    summary_md = (
        f"**Now Assessing:** {target_module_name} ({dataset_state.current_module_index + 1}/{len(ALLOWED_MODULES)})\n"
        f"**Simulated Severity:** {display_severity}\n"
        f"**Completed Modules:** {assessed_count}"
    )

    header, probes_update, single_upd, multi_upd, other_upd, item_rating_upd, saved_md, progress_md = _prepare_item_view(new_session)
    module_data = get_module_data(target_module_name)
    chat_display = _build_chat_display(new_session)

    sev_text_lower = SEVERITY_DISPLAY_MAP.get(severity, severity).lower()
    new_slider_label = f"This patient emulates {sev_text_lower} for the current module. How realistic is the portrayal? (0-5)"
    mod_rating_upd = gr.update(value=new_session.module_rating, label=new_slider_label)

    is_last_module = dataset_state.current_module_index >= len(ALLOWED_MODULES) - 1
    next_btn_text = "Finish Assessment" if is_last_module else "Save and Continue ➡️"

    save_checkpoint(dataset_state, new_session)

    return (
        dataset_state, new_session, summary_md,
        gr.update(value=new_session.module_comments), mod_rating_upd,
        gr.update(value=chat_display),
        header, probes_update, single_upd, multi_upd, other_upd, item_rating_upd,
        saved_md, progress_md,
        f"### Conducting Interview: {target_module_name}", module_data.get("information", "No info."),
        gr.update(interactive=dataset_state.current_module_index > 0),
        gr.update(value=next_btn_text, interactive=True),
        gr.update(value=target_module_name),
        gr.update(variant="secondary", interactive=True),
        gr.update(variant="secondary", interactive=True),
        gr.update(interactive=True), gr.update(interactive=True), gr.update(selected=2),
        gr.update(visible=False)
    )

def nav_item(ds: DatasetState, st: SessionState, direction: str):
    if not st or not st.items:
        return ds, st, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()

    st.asked_probes.clear()
    if direction == "next":
        st.idx = min(st.idx + 1, len(st.items) - 1)
    elif direction == "prev":
        st.idx = max(st.idx - 1, 0)

    header, probes_update, single_upd, multi_upd, other_upd, item_rating_upd, saved_md, progress_md = _prepare_item_view(st)
    save_checkpoint(ds, st)
    return ds, st, header, probes_update, single_upd, multi_upd, other_upd, item_rating_upd, saved_md, progress_md

def ask_selected(ds: DatasetState, st: SessionState, selected_probes: List[str]):
    item = _current_item(st)
    if not st or not st.agent or not item or not selected_probes:
        return ds, st, _build_chat_display(st), gr.update()

    questions_to_ask = [q.lstrip("✔ ").strip() for q in selected_probes]
    for q in questions_to_ask:
        reply = st.agent.respond(q, item)
        st.transcript.append((q, reply))
        st.asked_probes.add(q)

    probes_list = item.probing_questions or []
    formatted_probes = [f"✔ {q}" if q in st.asked_probes else q for q in probes_list]
    save_checkpoint(ds, st)
    return ds, st, _build_chat_display(st), gr.update(choices=formatted_probes, value=[])

def ask_custom(ds: DatasetState, st: SessionState, custom_q: str):
    item = _current_item(st)
    if not st or not st.agent or not item or not custom_q.strip():
        return ds, st, _build_chat_display(st), ""

    reply = st.agent.respond(custom_q, item)
    st.transcript.append((custom_q, reply))
    save_checkpoint(ds, st)
    return ds, st, _build_chat_display(st), ""

def auto_save_selection(ds: DatasetState, st: SessionState, single_sel: Optional[str], multi_sel: List[str], other_sel: List[str], rating: int):
    item = _current_item(st)

    if not st or not item:
        return ds, st, _saved_status_text(st), gr.update(visible=False)

    key = str(item.item_number)
    rec = st.selections.get(key, SelectionRecord(item_number=item.item_number))
    
    def get_base_choice(selection, options):
        if not selection or not options:
            return selection
        for opt in options:
            if opt.endswith(selection):
                return opt
        return selection

    rec.choice = single_sel if single_sel else None
    rec.choices = []
    rec.other_choices = []

    if multi_sel:
        all_multi_opts = _choices_for_item(st, item)[1]
        rec.choices = [get_base_choice(s, all_multi_opts) for s in multi_sel]
    if other_sel:
        all_other_opts = _choices_for_item(st, item)[2]
        rec.other_choices = [get_base_choice(s, all_other_opts) for s in other_sel]
    if rating is not None:
        rec.realism_rating = int(rating)

    st.selections[key] = rec

    modal_update = gr.update(visible=False)
    if item.is_screening and single_sel and str(single_sel).startswith("NO"):
        print("[Debug] Screening 'NO' detected. Triggering modal.")
        modal_update = gr.update(visible=True)

    save_checkpoint(ds, st)
    return ds, st, _saved_status_text(st), modal_update

def _saved_status_text(st: SessionState) -> str:
    if not st or not st.selections:
        return "Saved selections for **0/--** items."
    n = len(st.selections)
    total = len(st.items) if st and st.items else 0
    return f"Saved selections for **{n}/{total}** items."


def _build_chat_display(st: Optional[SessionState]) -> List[Dict[str, str]]:
    chat_display: List[Dict[str, str]] = []
    if not st:
        return chat_display
    for user_msg, assistant_msg in st.transcript:
        chat_display.append({"role": "user", "content": user_msg})
        chat_display.append({"role": "assistant", "content": assistant_msg})
    return chat_display


def _current_item_rating_value(st: Optional[SessionState], item: Optional[EPSICItem]) -> int:
    if not st or not item:
        return 0
    rec = st.selections.get(str(item.item_number))
    if rec and rec.realism_rating is not None:
        return int(rec.realism_rating)
    return 0


def _prepare_item_view(st: Optional[SessionState]):
    if not st or not st.items:
        return (
            gr.update(value="No active item."),
            gr.update(choices=[], value=[]),
            gr.update(choices=[], value=None, visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(value=0),
            _saved_status_text(st),
            "",
        )

    cur = _current_item(st)
    header = _format_item_header(cur, st.module_name)
    single_opts, multi_opts, other_opts, selected_single, selected_multi, selected_other = _choices_for_item(st, cur)
    asked = st.asked_probes or set()
    probes_list = cur.probing_questions or []
    formatted_probes = [f"✔ {q}" if q in asked else q for q in probes_list]
    progress_md = f"**Item {st.idx + 1} of {len(st.items)}**"

    return (
        header,
        gr.update(choices=formatted_probes, value=[]),
        gr.update(choices=single_opts, value=selected_single, visible=bool(single_opts), interactive=True),
        gr.update(choices=multi_opts, value=selected_multi, visible=bool(multi_opts), interactive=True),
        gr.update(choices=other_opts, value=selected_other, visible=bool(other_opts), interactive=True),
        gr.update(value=_current_item_rating_value(st, cur)),
        _saved_status_text(st),
        progress_md,
    )


def serialize_chat_history(agent: Optional["HFInferencePatient"]) -> List[Dict[str, Any]]:
    history: List[Dict[str, Any]] = []
    if not agent or not getattr(agent, "chat_session", None):
        return history
    for msg in getattr(agent.chat_session, "history", []) or []:
        parts_out: List[str] = []
        for part in getattr(msg, "parts", []) or []:
            txt = getattr(part, "text", None)
            if txt is not None:
                parts_out.append(txt)
        history.append({"role": getattr(msg, "role", "user"), "parts": parts_out})
    return history


def session_state_to_dict(st: SessionState) -> Dict[str, Any]:
    return {
        "module_name": st.module_name,
        "severity": st.severity,
        "age": st.age,
        "sex": st.sex,
        "idx": st.idx,
        "selections": {k: asdict(v) for k, v in st.selections.items()},
        "transcript": [[a, b] for a, b in st.transcript],
        "asked_probes": sorted(list(st.asked_probes)),
        "module_comments": st.module_comments,
        "module_rating": st.module_rating,
        "chat_history": serialize_chat_history(st.agent),
    }


def dataset_state_to_checkpoint_dict(ds: DatasetState, current_session: Optional[SessionState]) -> Dict[str, Any]:
    return {
        "version": 1,
        "status": "active",
        "clinician_id": ds.clinician_id,
        "demo_session_id": ds.demo_session_id,
        "consent_given": ds.consent_given,
        "consent_timestamp_utc": ds.consent_timestamp_utc,
        "consent_version": ds.consent_version,
        "feedback_saving_enabled": ds.feedback_saving_enabled,
        "profession": ds.profession,
        "yrs_diag_exp": ds.yrs_diag_exp,
        "yrs_prof_exp": ds.yrs_prof_exp,
        "patient_id": ds.patient_id,
        "patient_age": ds.patient_age,
        "patient_sex": ds.patient_sex,
        "patient_profile": ds.patient_profile,
        "patient_adhd_subtype": ds.patient_adhd_subtype,
        "current_patient_data": ds.current_patient_data,
        "current_module_index": ds.current_module_index,
        "diag_yes_modules": sorted(list(ds.diag_yes_modules)),
        "completed_sessions": {name: session_state_to_dict(st) for name, st in ds.completed_sessions.items()},
        "current_session": session_state_to_dict(current_session) if current_session else None,
    }


def simple_session_from_dict(data: Dict[str, Any]) -> SessionState:
    st = SessionState(
        module_name=data["module_name"],
        severity=data["severity"],
        age=int(data.get("age", 0)),
        sex=str(data.get("sex", "")),
    )
    st.items = get_items_for_module(st.module_name)
    st.idx = int(data.get("idx", 0))
    st.selections = {k: SelectionRecord(**v) for k, v in (data.get("selections") or {}).items()}
    st.transcript = [tuple(x) for x in data.get("transcript", [])]
    st.asked_probes = set(data.get("asked_probes", []))
    st.module_comments = data.get("module_comments", "")
    st.module_rating = int(data.get("module_rating", 0) or 0)
    return st


def save_checkpoint(ds: DatasetState, current_session: Optional[SessionState]):
    if DEMO_MODE and not getattr(ds, "feedback_saving_enabled", False):
        return
    if not ds or not ds.patient_id or not ds.clinician_id:
        return
    payload = dataset_state_to_checkpoint_dict(ds, current_session)
    upload_json_to_hub(payload, checkpoint_repo_path(ds.clinician_id, ds.patient_id))
    upload_json_to_hub({"patient_id": ds.patient_id, "status": "active"}, active_pointer_repo_path(ds.clinician_id))


def close_checkpoint(ds: DatasetState, status: str = "closed"):
    if DEMO_MODE and not getattr(ds, "feedback_saving_enabled", False):
        return
    if not ds or not ds.clinician_id:
        return
    payload = {"patient_id": ds.patient_id, "status": status, "closed_at": datetime.utcnow().isoformat()}
    upload_json_to_hub(payload, active_pointer_repo_path(ds.clinician_id))


def clear_patient_progress(ds: DatasetState):
    ds.current_patient_data = {}
    ds.patient_id = None
    ds.patient_age = None
    ds.patient_sex = None
    ds.patient_profile = {}
    ds.patient_adhd_subtype = "none"
    ds.completed_sessions = {}
    ds.current_module_index = -1
    ds.diag_yes_modules = set()
    return ds

# =========================
# Dataset and Summary Logic
# =========================
def init_dataset_state() -> DatasetState:
    df = load_or_init_main_dataset()
    ds = DatasetState(columns=MASTER_COLUMNS, df=df)
    if DEMO_MODE:
        ds.demo_session_id = make_demo_session_id()
        ds.clinician_id = ds.demo_session_id
        ds.profession = "Not specified"
        ds.yrs_diag_exp = "Not specified"
        ds.yrs_prof_exp = "Not specified"
    return ds

def check_active_session(dataset_state):
    """Checks if a patient is currently active."""
    if dataset_state and dataset_state.patient_id:
        # If there is an ID, show the warning modal
        return gr.update(visible=True)
    else:
        # If no ID, don't show modal, just signal to proceed (handled via event logic below)
        return gr.update(visible=False)

def new_patient(
    dataset_state: DatasetState,
    prof: str,
    y_diag: str,
    y_prof: str,
    feedback_choice: Optional[str] = None,
    profile: Optional[gr.OAuthProfile] = None,
    request: Optional[gr.Request] = None,
):
    if PREGEN_PROFILES is None:
        raise ValueError("Pre-generated profiles not loaded.")

    if DEMO_MODE:
        apply_demo_feedback_choice(dataset_state, feedback_choice)
    else:
        sync_main_artifacts_from_hub()

    if dataset_state.patient_id and dataset_state.clinician_id:
        close_checkpoint(dataset_state, status="discarded")
        clear_patient_progress(dataset_state)

    if DEMO_MODE:
        anon_id = dataset_state.demo_session_id or make_demo_session_id()
        dataset_state.demo_session_id = anon_id
        dataset_state.clinician_id = anon_id
    else:
        anon_id = make_anonymous_clinician_id(profile, request)
        if not anon_id:
            raise gr.Error("Please sign in with Hugging Face OAuth to start an assessment.")

    prof = normalize_optional_background(prof)
    y_diag = normalize_optional_background(y_diag)
    y_prof = normalize_optional_background(y_prof)

    total_profiles = len(PREGEN_PROFILES)
    idx = random.randint(0, total_profiles - 1)

    if DEMO_MODE:
        pid = f"{anon_id}-P{idx}"
    else:
        used_indices = set()
        if os.path.exists(MAIN_DATASET_PATH):
            try:
                df = pd.read_csv(MAIN_DATASET_PATH, usecols=["PatientID"])
                for x in df["PatientID"].dropna().astype(str):
                    if x.startswith("P") and x[1:].isdigit():
                        used_indices.add(int(x[1:]))
            except Exception:
                pass

        if os.path.exists(MAIN_JSON_PATH):
            try:
                with open(MAIN_JSON_PATH, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    if isinstance(data, list):
                        for p in data:
                            pid_old = p.get("patient_id", "")
                            if isinstance(pid_old, str) and pid_old.startswith("P") and pid_old[1:].isdigit():
                                used_indices.add(int(pid_old[1:]))
            except Exception:
                pass

        if len(used_indices) < (total_profiles * 0.8):
            for _ in range(1000):
                candidate = random.randint(0, total_profiles - 1)
                if candidate not in used_indices:
                    idx = candidate
                    break
        else:
            available = list(set(range(total_profiles)) - used_indices)
            if available:
                idx = random.choice(available)
        pid = f"P{idx}"
    bin_map = dict(zip(PPROF_DISORDERS, PREGEN_PROFILES[idx]))
    age = int(np.clip(np.round(np.random.normal(AGE_DISTRIBUTION["mean"], AGE_DISTRIBUTION["sd"])), 13, 18))
    sex = random.choices(["female", "male"], weights=[0.686, 0.314])[0]

    profile_map = {}
    adhd_sub = "none"
    for mod in ALLOWED_MODULES:
        is_sev = False
        sub_sev = "none"
        if mod == "ADHD":
            if bin_map.get("ADHD") == 1 or bin_map.get("ADD") == 1:
                is_sev = True
                sub_sev = "hyperactive_combined" if bin_map.get("ADHD") == 1 else "inattentive"
        elif bin_map.get(mod) == 1:
            is_sev = True

        if is_sev:
            profile_map[mod] = "severe"
            if mod == "ADHD":
                adhd_sub = sub_sev
        else:
            dist = DISTRIBUTIONS.get(mod, {})
            p_mod = dist.get("2", 0)
            p_none = dist.get("0", 0) + dist.get("1", 0) + dist.get("system", 0)
            tot = p_mod + p_none
            profile_map[mod] = random.choices(["moderate", "none"], weights=[p_mod/tot, p_none/tot], k=1)[0] if tot > 0 else "none"
            if mod == "ADHD" and profile_map[mod] == "moderate":
                adhd_sub = random.choice(["inattentive", "hyperactive_combined"])

    dataset_state.clinician_id = anon_id
    dataset_state.profession = prof
    dataset_state.yrs_diag_exp = y_diag
    dataset_state.yrs_prof_exp = y_prof
    dataset_state.patient_id = pid
    dataset_state.patient_age = age
    dataset_state.patient_sex = sex
    dataset_state.patient_profile = profile_map
    dataset_state.patient_adhd_subtype = adhd_sub
    dataset_state.current_patient_data = {
        "DemoSessionID": dataset_state.demo_session_id or anon_id,
        "ClinicianID": anon_id,
        "ConsentGiven": "Yes" if dataset_state.consent_given else "No",
        "ConsentTimestampUTC": dataset_state.consent_timestamp_utc,
        "ConsentVersion": dataset_state.consent_version,
        "Profession": prof,
        "YrsDiagExp": y_diag,
        "YrsProfExp": y_prof,
        "PatientID": pid,
        "Sex": sex,
        "Age": age,
    }
    dataset_state.completed_sessions = {}
    dataset_state.current_module_index = -1
    dataset_state.diag_yes_modules = set()

    saving_status = "Feedback saving is enabled for this demo session." if dataset_state.feedback_saving_enabled else "Feedback saving is disabled for this demo session."
    status = (
        f"🆕 New demo patient initialized: **{pid}** ({age}/{sex}).\n"
        f"Demo session ID: **{dataset_state.clinician_id}**\n"
        f"{saving_status}\n"
        f"**The 'Step 2: Conduct Interview' tab is now unlocked.**"
    )

    save_checkpoint(dataset_state, None)

    return (
        dataset_state,
        status,
        gr.update(value=[]),
        gr.update(value=pid),
        gr.update(value="Start First Module ➡️", variant="primary", interactive=True),
        gr.update(interactive=True),
        gr.update(value="_Patient ready. Click **Start First Module ➡️** to begin._"),
        gr.update(variant="secondary", interactive=False),
        gr.update(variant="secondary", interactive=False),
        gr.update(interactive=True),
        gr.update(interactive=True),
        gr.update(selected=2),
        None,
        gr.update(value=[], choices=[]),
        gr.update(value=None, choices=[], visible=False),
        gr.update(value=[], choices=[], visible=False),
        gr.update(value=[], choices=[], visible=False),
        gr.update(value=0),
        gr.update(value=0),
        gr.update(value=anon_id),
    )

def click_next_or_finish(dataset_state, current_session, comments, mod_rating):
    is_finish_click = dataset_state.current_module_index >= len(ALLOWED_MODULES) - 1

    if is_finish_click:
        print("Finish Assessment clicked. Saving final module.")
        if current_session:
            persist_module_session_into_dataset(dataset_state, current_session, comments, mod_rating)
            save_checkpoint(dataset_state, current_session)

        return (
            dataset_state, current_session, gr.update(value="**Assessment Complete!**"),
            gr.update(value=comments), gr.update(value=mod_rating), gr.update(value=_build_chat_display(current_session)),
            gr.update(value="### ✅ Assessment Complete"),
            gr.update(choices=[], value=[], visible=True),
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value=0),
            gr.update(value="Go to 'Step 3 & 4: Review & Save' tab."),
            gr.update(), gr.update(value="### ✅ Assessment Complete"), gr.update(),
            gr.update(interactive=False), gr.update(value="Finish Assessment", variant="secondary", interactive=False),
            gr.update(interactive=False), gr.update(variant="primary", interactive=True),
            gr.update(variant="secondary", interactive=False),
            gr.update(interactive=False), gr.update(interactive=True), gr.update(selected=3),
            gr.update(visible=False)
        )
    return change_module(dataset_state, current_session, comments, mod_rating, "next", None)

def _compose_row_for_patient(dataset_state: DatasetState, pid: str) -> Dict[str, Any]:
    row: Dict[str, Any] = {col: "" for col in dataset_state.columns}
    row.update(dataset_state.current_patient_data)
    
    # --- MODIFIED: Save Demographics ---
    row["DemoSessionID"] = dataset_state.demo_session_id
    row["ClinicianID"] = dataset_state.clinician_id
    row["ConsentGiven"] = "Yes" if dataset_state.consent_given else "No"
    row["ConsentTimestampUTC"] = dataset_state.consent_timestamp_utc
    row["ConsentVersion"] = dataset_state.consent_version
    row["Profession"] = dataset_state.profession
    row["YrsDiagExp"] = dataset_state.yrs_diag_exp
    row["YrsProfExp"] = dataset_state.yrs_prof_exp
    # -----------------------------------
    
    assessed_modules = dataset_state.completed_sessions.keys()
    
    # 1. Fill in missing data for un-assessed modules
    for mod in ALLOWED_MODULES:
        if mod not in assessed_modules:
            mp = MODULE_PREFIX[mod]
            codes = MODULE_ITEM_CODES[mod]
            for code in codes:
                row[code] = DEFAULT_MISSING
                row[f"RAT_{code}"] = "" 

    # 2. Fill in RATING & Selection data
    for module_name, st in dataset_state.completed_sessions.items():
        mp = MODULE_PREFIX[module_name]
        for item in st.items:
            code = item_code(mp, item)
            rating_col = f"RAT_{code}"
            
            key = str(item.item_number)
            if key in st.selections:
                rec = st.selections[key]
                # Save choice/text (logic handled in current_patient_data updates usually, but ensuring here)
                
                if rec.realism_rating is not None:
                    row[rating_col] = rec.realism_rating

    # 3. Diagnostic Columns
    for mod in ALLOWED_MODULES:
        diag_col = f"DIA_{MODULE_PREFIX[mod]}"
        row[diag_col] = "Yes" if mod in dataset_state.diag_yes_modules else "No"
        
    return row

def finalize_patient(dataset_state):
    pid = dataset_state.patient_id
    if not pid:
        return dataset_state, None, "Error", *[gr.update()]*8
    if DEMO_MODE and not getattr(dataset_state, "feedback_saving_enabled", False):
        return (
            dataset_state, None, f"✅ Demo completed. {feedback_not_saved_message()}",
            gr.update(variant="primary", interactive=True), gr.update(interactive=True),
            gr.update(interactive=True), gr.update(variant="secondary", interactive=True),
            gr.update(variant="primary", interactive=True), gr.update(interactive=True),
            gr.update(interactive=True), gr.update(selected=4)
        )

    final_diagnoses_met = set()
    for module_name, st in dataset_state.completed_sessions.items():
        rule_data = DIAGNOSTIC_RULES.get(st.module_name)
        if not rule_data or "logic" not in rule_data:
            continue
        logic = rule_data["logic"]
        met = False
        if st.module_name == "ADHD":
            for subtype in ["combined", "inattentive", "hyperactive_impulsive"]:
                sub_met, _ = evaluate_diagnostic_rule(logic.get(subtype, {}), st.selections)
                if sub_met:
                    met = True
                    break
        else:
            met, _ = evaluate_diagnostic_rule(logic, st.selections)
        if met:
            final_diagnoses_met.add(module_name)
    dataset_state.diag_yes_modules = final_diagnoses_met

    row = _compose_row_for_patient(dataset_state, pid)
    row["DemoSessionID"] = dataset_state.demo_session_id
    row["ClinicianID"] = dataset_state.clinician_id
    row["ConsentGiven"] = "Yes" if dataset_state.consent_given else "No"
    row["ConsentTimestampUTC"] = dataset_state.consent_timestamp_utc
    row["ConsentVersion"] = dataset_state.consent_version
    row["Profession"] = dataset_state.profession
    row["YrsDiagExp"] = dataset_state.yrs_diag_exp
    row["YrsProfExp"] = dataset_state.yrs_prof_exp
    row["PatientID"] = pid
    row["Sex"] = dataset_state.patient_sex
    row["Age"] = dataset_state.patient_age

    for mod, st in dataset_state.completed_sessions.items():
        mp = MODULE_PREFIX[mod]
        row[f"MOD_RAT_{mp}"] = st.module_rating
        for it in st.items:
            code = item_code(mp, it)
            key = str(it.item_number)
            if key in st.selections and st.selections[key].realism_rating is not None:
                row[f"RAT_{code}"] = st.selections[key].realism_rating

    for mod in ALLOWED_MODULES:
        row[f"DIA_{MODULE_PREFIX[mod]}"] = "Yes" if mod in final_diagnoses_met else "No"

    row_df = pd.DataFrame([row])
    for c in MASTER_COLUMNS:
        if c not in row_df.columns:
            row_df[c] = ""
    row_df = row_df[MASTER_COLUMNS]

    patient_csv_local = os.path.join(DATASET_DIR, f"{pid}.csv")
    row_df.to_csv(patient_csv_local, index=False)
    upload_to_hub(patient_csv_local, path_in_repo=patient_csv_repo_path(pid))

    sync_main_artifacts_from_hub()
    main_df = load_or_init_main_dataset()
    if "PatientID" in main_df.columns:
        main_df = main_df[main_df["PatientID"].astype(str) != str(pid)]
    main_df = pd.concat([main_df, row_df], ignore_index=True)
    main_df.to_csv(MAIN_DATASET_PATH, index=False)
    dataset_state.df = main_df
    upload_to_hub(MAIN_DATASET_PATH, path_in_repo=os.path.basename(MAIN_DATASET_PATH))

    json_file, patient_json_local = export_conversation_json(dataset_state, pid)
    if json_file:
        upload_to_hub(json_file, path_in_repo=os.path.basename(json_file))
    if patient_json_local:
        upload_to_hub(patient_json_local, path_in_repo=patient_json_repo_path(pid))

    dataset_state.last_saved_path = MAIN_DATASET_PATH

    return (
        dataset_state, MAIN_DATASET_PATH, "✅ Record saved and uploaded. You can continue the assessment or start a new patient from Tab 1.",
        gr.update(variant="primary", interactive=True), gr.update(interactive=True),
        gr.update(interactive=True), gr.update(variant="secondary", interactive=True),
        gr.update(variant="primary", interactive=True), gr.update(interactive=True),
        gr.update(interactive=True), gr.update(selected=2)
    )

def download_dataset(dataset_state: DatasetState):
    sync_main_artifacts_from_hub()
    if not os.path.exists(MAIN_DATASET_PATH):
        dataset_state.df.to_csv(MAIN_DATASET_PATH, index=False)
    return MAIN_DATASET_PATH

def export_conversation_json(dataset_state: DatasetState, pid: str):
    """Appends conversation to main JSON list and saves a per-patient JSON artifact."""
    if not pid or not dataset_state.completed_sessions:
        print("[WARN] No data to export.")
        return None, None

    data = {
        "patient_id": pid,
        "clinician_id": dataset_state.clinician_id,
        "demo_session_id": dataset_state.demo_session_id,
        "consent_given": dataset_state.consent_given,
        "consent_timestamp_utc": dataset_state.consent_timestamp_utc,
        "consent_version": dataset_state.consent_version,
        "timestamp": datetime.utcnow().isoformat(),
        "patient_age": dataset_state.patient_age,
        "patient_sex": dataset_state.patient_sex,
        "patient_profile_simulated_severity": dataset_state.patient_profile,
        "assessed_modules": []
    }

    for module_name, st in dataset_state.completed_sessions.items():
        module_data = {
            "module": st.module_name,
            "simulated_severity": st.severity,
            "clinician_comments": st.module_comments,
            "module_realism_rating": st.module_rating,
            "turns": [{"clinician": c, "patient": p} for c, p in st.transcript],
            "selections": [
                {
                    "item_number": rec.item_number,
                    "choice": rec.choice,
                    "choices": rec.choices,
                    "other_choices": rec.other_choices,
                    "realism_rating": rec.realism_rating,
                }
                for rec in st.selections.values()
            ],
        }
        data["assessed_modules"].append(module_data)

    patient_json_local = os.path.join(DATASET_DIR, f"{pid}.json")
    with open(patient_json_local, "w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2, ensure_ascii=False)

    sync_main_artifacts_from_hub()
    all_conversations = []
    if os.path.exists(MAIN_JSON_PATH):
        try:
            with open(MAIN_JSON_PATH, "r", encoding="utf-8") as fh:
                all_conversations = json.load(fh)
                if not isinstance(all_conversations, list):
                    all_conversations = []
        except Exception:
            all_conversations = []

    all_conversations = [entry for entry in all_conversations if str(entry.get("patient_id")) != str(pid)]
    all_conversations.append(data)

    with open(MAIN_JSON_PATH, "w", encoding="utf-8") as fh:
        json.dump(all_conversations, fh, indent=2, ensure_ascii=False)

    return MAIN_JSON_PATH, patient_json_local

# =========================
# Diagnostic Summary Logic
# =========================
def evaluate_diagnostic_rule(logic: Dict[str, Any], selections: Dict[str, SelectionRecord]) -> Tuple[bool, str]:
    def check(item_num_str: str, condition: str, value=None) -> bool:
        rec = selections.get(item_num_str)
        if not rec: return False
        
        if condition == "equals_yes":
            return rec.choice and rec.choice.strip().startswith("YES")
        elif condition == "min_checked":
            count = len(rec.choices or []) + len(rec.other_choices or [])
            return count >= value
        elif condition == "min_yes":
            count = 0
            for item in value:
                rec_inner = selections.get(str(item))
                if rec_inner and rec_inner.choice == "YES":
                    count += 1
            return count >= item_num_str
        return False

    if "all" in logic:
        for i, sub_logic in enumerate(logic["all"]):
            res, reason = evaluate_diagnostic_rule(sub_logic, selections)
            if not res:
                return False, f"A required condition was not met: {reason}"
        return True, "All conditions met."
    
    if "any" in logic:
        any_met = False
        reasons = []
        for sub_logic in logic["any"]:
            res, reason = evaluate_diagnostic_rule(sub_logic, selections)
            reasons.append(reason)
            if res:
                any_met = True
                break
        if any_met:
            return True, f"One of the optional conditions was met: {reason}"
        else:
            return False, "None of the optional conditions were met. " + " ".join(reasons)

    if "item" in logic and "equals" in logic:
        item, val = str(logic["item"]), logic["equals"]
        if val == "YES":
            is_met = check(item, "equals_yes")
            reason = f"Item {item} was {'YES' if is_met else 'not YES'}."
            return is_met, reason

    if "min_checked" in logic and "item" in logic:
        count, item = logic["min_checked"], str(logic["item"])
        is_met = check(item, "min_checked", count)
        num_checked = len((selections.get(item) or SelectionRecord(item_number=None)).choices or [])
        reason = f"Item {item} required at least {count} selections (had {num_checked})."
        return is_met, reason

    if "min_yes" in logic and "items" in logic:
        count, items = logic["min_yes"], logic["items"]
        is_met = check(count, "min_yes", items)
        num_yes = sum(1 for i in items if (selections.get(str(i)) or SelectionRecord(item_number=None)).choice == "YES")
        reason = f"At least {count} of items {', '.join(map(str, items))} needed to be YES (had {num_yes})."
        return is_met, reason

    if "symptom_block" in logic:
        block = logic["symptom_block"]
        items, min_count, must_haves = block["items"], block["count_yes_at_least"], block["must_include_any_of"]
        
        yes_items = {str(i) for i in items if (selections.get(str(i)) or SelectionRecord(item_number=None)).choice == "YES"}
        total_yes = len(yes_items)
        has_must_have = any(str(i) in yes_items for i in must_haves)

        if total_yes >= min_count and has_must_have:
            return True, f"Met criteria with {total_yes}/{min_count} symptoms, including a core symptom."
        else:
            reason = f"Required ≥{min_count} symptoms including one of {', '.join(map(str, must_haves))}. "
            reason += f"Found {total_yes} YES answers. Core symptom present: {has_must_have}."
            return False, reason

    return False, "Unknown rule format."

def generate_full_summary(dataset_state: DatasetState):
    if not dataset_state or not dataset_state.completed_sessions:
        return (
            dataset_state, 
            "No modules have been completed yet. Finish at least one module to generate a summary.",
            gr.update(), # generate_summary_btn
            gr.update()  # finalize_btn
        )

    summary = f"# 📋 Full Diagnostic Summary for Patient {dataset_state.patient_id}\n"
    summary += f"**Patient Profile:** {dataset_state.patient_age}-year-old {dataset_state.patient_sex}\n\n"

    final_diagnoses_met = set()

    for module_name, st in sorted(dataset_state.completed_sessions.items()):
        summary += "__________________________________________________________________\n"
        summary += f"### Module: {module_name}\n"
        summary += f"**Simulated Severity:** {st.severity.capitalize()}\n\n"

        # --- 1. Overall Diagnostic Status ---
        rule_data = DIAGNOSTIC_RULES.get(st.module_name)
        if not rule_data or "logic" not in rule_data:
            summary += "- *No diagnostic algorithm found for this module.*\n"
            continue

        logic = rule_data["logic"]
        rule_text = rule_data.get("rule", "No rule description available.")
        
        met = False
        reason = ""
        
        # Calculate Logic
        if st.module_name == "ADHD":
            adhd_results = []
            for subtype in ["combined", "inattentive", "hyperactive_impulsive"]:
                sub_met, sub_reason = evaluate_diagnostic_rule(logic.get(subtype, {}), st.selections)
                if sub_met:
                    adhd_results.append(f"✅ **Criteria MET for ADHD, {subtype.replace('_', ' ').title()} Type.**\n> *Evaluation:* {sub_reason}")
                    met = True
            if not adhd_results:
                _, reason = evaluate_diagnostic_rule(logic.get('combined',{}), st.selections)
                summary += f"❌ **Diagnostic Criteria NOT MET.**\n> *Rule Checked:* `{rule_text}`\n> *Evaluation:* {reason}\n\n"
            else:
                summary += "\n".join(adhd_results) + "\n\n"
        else:
            met, reason = evaluate_diagnostic_rule(logic, st.selections)
            summary += (
                f"✅ **Diagnostic Criteria MET.**\n> *Rule Checked:* `{rule_text}`\n> *Evaluation:* {reason}\n\n" 
                if met 
                else f"❌ **Diagnostic Criteria NOT MET.**\n> *Rule Checked:* `{rule_text}`\n> *Evaluation:* {reason}\n\n"
            )
        
        if met:
            final_diagnoses_met.add(module_name)

            
        # --- 2. Item-by-Item Breakdown Table ---
        summary += "| Status | Item | DSM-5 | Clinician Selection |\n"
        summary += "| :---: | :--- | :--- | :--- |\n"
        
        # Get mapping for this module
        dsm_map = DSM5_MAPPING.get(module_name, {})
        
        for item in st.items:
            item_num = str(item.item_number)
            dsm_code = dsm_map.get(item_num, "-")
            
            # Get Selection
            rec = st.selections.get(item_num)
            
            display_val = "_(Not answered)_"
            icon = "⚪" # Neutral
            
            if rec:
                # Handle Radio Buttons (YES/NO)
                if rec.choice:
                    display_val = f"**{rec.choice}**"
                    if rec.choice.strip().upper().startswith("YES"):
                        icon = "✅"
                    elif rec.choice.strip().upper().startswith("NO"):
                        icon = "❌"
                    else:
                        icon = "➖"
                
                # Handle Checkboxes
                elif rec.choices:
                    count = len(rec.choices)
                    display_val = f"Checked {count} symptoms"
                    if count > 0:
                        icon = "✅" # Assuming checking anything is a 'positive' finding for that item context
                    else:
                        icon = "❌"
            
            # Special handling for Screening Item
            if item.is_screening:
                dsm_code = "Screen"
            
            summary += f"| {icon} | Item {item_num} | {dsm_code} | {display_val} |\n"

        if st.module_comments:
            summary += f"\n**Clinician Comments:**\n> {st.module_comments}\n"
        
        summary += "\n"

    dataset_state.diag_yes_modules = final_diagnoses_met
            
    return (
        dataset_state, 
        summary,
        gr.update(variant="secondary", interactive=True), # generate_summary_btn
        gr.update(variant="primary", interactive=True)   # finalize_btn
    )



def build_load_view(profile: Optional[gr.OAuthProfile] = None, request: Optional[gr.Request] = None):
    ds, current_session, anon_id = maybe_restore_checkpoint(profile, request)
    anon_upd = gr.update(value=anon_id or "")

    if not anon_id:
        return (
            ds, None, anon_upd,
            gr.update(value="_No demo session ID was created. Reload the page to start a new session._"),
            gr.update(value=""),
            gr.update(value="Start First Module ➡️", variant="secondary", interactive=False),
            gr.update(interactive=False),
            gr.update(value="_Start a new patient and module to see assessment status._"),
            gr.update(variant="secondary", interactive=False),
            gr.update(variant="secondary", interactive=False),
            gr.update(interactive=False), gr.update(interactive=False), gr.update(selected=1),
            gr.update(choices=[], value=[]),
            gr.update(choices=[], value=None, visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(value=0), gr.update(value=0),
            gr.update(value="..."), gr.update(value="_No selections saved yet._"), gr.update(value=""),
            gr.update(value="### Conduct Diagnostic Interview"), gr.update(value="Start a session to see module information."),
            gr.update(interactive=False), gr.update(value=None), gr.update(value=[]), gr.update(interactive=False), gr.update(value=""),
            gr.update(), gr.update(), gr.update()
        )

    if ds.patient_id and current_session is None:
        _prof = ds.profession if ds.profession not in ("Unknown", "", None) else None
        return (
            ds, current_session, anon_upd,
            gr.update(value=f"🔄 Restored active patient **{ds.patient_id}**.\nAnonymous evaluator ID: **{ds.clinician_id}**\nClick **Start First Module ➡️** to continue."),
            gr.update(value=ds.patient_id),
            gr.update(value="Start First Module ➡️", variant="primary", interactive=True),
            gr.update(interactive=True),
            gr.update(value="_Patient ready. Click **Start First Module ➡️** to begin._"),
            gr.update(variant="secondary", interactive=False),
            gr.update(variant="secondary", interactive=False),
            gr.update(interactive=True), gr.update(interactive=True), gr.update(selected=2),
            gr.update(choices=[], value=[]),
            gr.update(choices=[], value=None, visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(value=0), gr.update(value=0),
            gr.update(value="..."), gr.update(value="_No selections saved yet._"), gr.update(value=""),
            gr.update(value="### Conduct Diagnostic Interview"), gr.update(value="Start a session to see module information."),
            gr.update(interactive=False), gr.update(value=None), gr.update(value=[]), gr.update(interactive=False), gr.update(value=""),
            gr.update(value=_prof), gr.update(value=optional_number_value(ds.yrs_diag_exp)), gr.update(value=optional_number_value(ds.yrs_prof_exp))
        )

    if not ds.patient_id or current_session is None:
        _prof = ds.profession if ds.profession not in ("Unknown", "", None) else None
        return (
            ds, current_session, anon_upd,
            gr.update(value="_No active demo patient yet. Click 'Start Demo Patient' to begin._"),
            gr.update(value=""),
            gr.update(value="Start First Module ➡️", variant="secondary", interactive=False),
            gr.update(interactive=False),
            gr.update(value="_Start a new patient and module to see assessment status._"),
            gr.update(variant="secondary", interactive=False),
            gr.update(variant="secondary", interactive=False),
            gr.update(interactive=False), gr.update(interactive=False), gr.update(selected=1),
            gr.update(choices=[], value=[]),
            gr.update(choices=[], value=None, visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(choices=[], value=[], visible=False),
            gr.update(value=0), gr.update(value=0),
            gr.update(value="..."), gr.update(value="_No selections saved yet._"), gr.update(value=""),
            gr.update(value="### Conduct Diagnostic Interview"), gr.update(value="Start a session to see module information."),
            gr.update(interactive=False), gr.update(value=None), gr.update(value=[]), gr.update(interactive=False), gr.update(value=""),
            gr.update(value=_prof), gr.update(value=optional_number_value(ds.yrs_diag_exp)), gr.update(value=optional_number_value(ds.yrs_prof_exp))
        )

    severity = current_session.severity
    display_severity = SEVERITY_DISPLAY_MAP.get(severity, severity)
    module_position = ds.current_module_index + 1 if ds.current_module_index >= 0 else 1
    summary_md = (
        f"**Now Assessing:** {current_session.module_name} ({module_position}/{len(ALLOWED_MODULES)})\n"
        f"**Simulated Severity:** {display_severity}\n"
        f"**Completed Modules:** {len(ds.completed_sessions)}"
    )
    header, probes_update, single_upd, multi_upd, other_upd, item_rating_upd, saved_md, progress_md = _prepare_item_view(current_session)
    module_data = get_module_data(current_session.module_name)
    sev_text_lower = display_severity.lower()
    new_slider_label = f"This patient emulates {sev_text_lower} for the current module. How realistic is the portrayal? (0-5)"
    is_last_module = ds.current_module_index >= len(ALLOWED_MODULES) - 1
    next_btn_text = "Finish Assessment" if is_last_module else "Save and Continue ➡️"

    _prof = ds.profession if ds.profession not in ("Unknown", "", None) else None
    return (
        ds, current_session, anon_upd,
        gr.update(value=f"🔄 Resumed active patient **{ds.patient_id}**.\nAnonymous evaluator ID: **{ds.clinician_id}**"),
        gr.update(value=ds.patient_id),
        gr.update(value=next_btn_text, variant="primary", interactive=True),
        gr.update(interactive=True),
        gr.update(value=summary_md),
        gr.update(variant="secondary", interactive=True),
        gr.update(variant="secondary", interactive=True),
        gr.update(interactive=True), gr.update(interactive=True), gr.update(selected=2),
        probes_update, single_upd, multi_upd, other_upd, item_rating_upd,
        gr.update(value=current_session.module_rating, label=new_slider_label),
        gr.update(value=header), gr.update(value=saved_md), gr.update(value=progress_md),
        gr.update(value=f"### Conducting Interview: {current_session.module_name}"), gr.update(value=module_data.get("information", "No info.")),
        gr.update(interactive=ds.current_module_index > 0), gr.update(value=current_session.module_name), gr.update(value=_build_chat_display(current_session)), gr.update(interactive=False), gr.update(value=current_session.module_comments),
        gr.update(value=_prof), gr.update(value=optional_number_value(ds.yrs_diag_exp)), gr.update(value=optional_number_value(ds.yrs_prof_exp))
    )

# =========================
# UI Definition
# =========================
# General Home Instructions
INSTRUCTIONS_MD = """
### EPSI-C Patient Simulator — Symposium Demo

This demo lets you try the **Electronic Psychiatric Semi-Structured Interview for Children and Adolescents (EPSI-C)** with an AI-simulated adolescent patient. The simulated patient has a psychiatric profile, including possible comorbidities, and responds to your questions based on that profile.

You may try the app without contributing feedback. If you consent to provide feedback, your ratings, comments, SUS responses, and simulated-patient interactions may be saved.

Please do **not** enter names, contact details, or any other personally identifiable information.

**Suggested path:**

1. **Start Demo Patient** — Choose whether you consent to contribute feedback. Optional background information can be left blank.
2. **Conduct Interview** — Work through one or more modules, ask questions, make clinical selections, and rate realism.
3. **Review Demo Record** — Review the diagnostic summary and save only if you consented to feedback collection.
4. **System Evaluation** — Optionally complete SUS usability feedback.
"""

TAB_1_INSTRUCTIONS = """
**Instructions:**
1. Choose whether you consent to contribute demo feedback. You can still use the demo without saving feedback.
2. Optional background information can help interpret aggregate feedback, but it is not required.
3. Click **Start Demo Patient** to generate an AI patient and unlock the interview tab.

> Please do not enter identifiable information in free-text fields.
"""

TAB_2_INSTRUCTIONS = """
**Instructions:**
1. Click **Start First Module ➡️** to begin. The system can take you through all 11 disorder modules in sequence.
2. Select suggested **Probing Questions** or expand **Custom Question** to type your own. 
3. Use the **Clinician's Decision** section to record whether each diagnostic criterion is met.
4. Use the realism sliders to rate the AI patient responses and module-level realism.
5. Click **Save and Continue ➡️** when finished with a module.
"""

TAB_3_INSTRUCTIONS = """
**Instructions:**
1. Review the diagnostic summary generated from your clinical selections.
2. If you consented to contribute demo feedback, click **Save Demo Feedback Record** to upload the record to the separate demo feedback dataset.
3. If you chose not to contribute feedback, you may still review the summary. 
"""

# ========================================================================
# --- START OF MODIFIED UI SECTION (gr.Blocks) ---
# ========================================================================
# ========================================================================
# --- START OF MODIFIED UI SECTION (gr.Blocks) ---
# ========================================================================
css = """
.custom-modal {
    position: fixed !important;
    top: 50% !important;
    left: 50% !important;
    transform: translate(-50%, -50%) !important;
    z-index: 10000 !important;
    max-width: 600px !important;
    width: auto !important;
    margin: 0 !important;
    max-height: 90vh !important;
    overflow-y: auto !important;
    box-shadow: 0px 4px 20px rgba(0,0,0,0.5) !important;
    border-radius: 8px !important;
    background-color: #1f2937 !important;
    padding: 20px !important;
}
"""

with gr.Blocks(title="AI-Patient Demo", fill_height=True, css=css) as demo:
    gr.Markdown("## AI-Patient Demo")
    #gr.Markdown("Feedback is saved only when you explicitly consent, and it is routed to the separate demo feedback dataset configured for this Space.")
    anon_user_display = gr.Textbox(label="Demo Session ID", interactive=False, info="Generated for this browser session.")

    dataset_state = gr.State(init_dataset_state)
    session = gr.State()

    with gr.Tabs() as main_tabs:
        with gr.Tab("Instructions", id=0):
            gr.Markdown(INSTRUCTIONS_MD)

        with gr.Tab("1. Start Patient Evaluation", id=1):
            gr.Markdown(TAB_1_INSTRUCTIONS)
            gr.Markdown("---")
            with gr.Group():
                gr.Markdown("### Consent for Demo Feedback Collection")
                gr.Markdown(
                    "This is a public demonstration version of the EPSI-C simulated patient app. "
                    "If you consent, your ratings, comments, SUS responses, and simulated-patient interactions may be saved "
                    "for research and development purposes in a separate demo feedback dataset. Participation is voluntary. "
                    "Do not enter identifiable information or other sensitive personal data."
                )
                feedback_choice = gr.Radio(
                    choices=[
                        "I consent to save my demo feedback",
                        "Continue without saving feedback",
                    ],
                    value="Continue without saving feedback",
                    label="Feedback participation choice",
                    interactive=True,
                )

            with gr.Accordion("Optional background information", open=False):
                profession_input = gr.Dropdown(
                    choices=["Psychiatrist", "Psychologist", "Resident/Trainee", "Medical Student", "Nurse", "Researcher"],
                    label="Current profession (optional)",
                    allow_custom_value=True,
                    interactive=True,
                )
                with gr.Row():
                    yrs_diag_input = gr.Number(label="Years of experience with clinical diagnoses (optional)", value=None, precision=0)
                    yrs_prof_input = gr.Number(label="Years of professional experience (optional)", value=None, precision=0)

            gr.Markdown("---")
            with gr.Group():
                gr.Markdown("### Start Demo Patient Case")
                with gr.Row():
                    with gr.Column(scale=1):
                        new_patient_btn = gr.Button("Start Demo Patient", variant="primary", size="lg")
                    with gr.Column(scale=2):
                        patient_id = gr.Textbox(label="Current Patient ID", interactive=False)
                        patient_status = gr.Markdown("_No active demo patient yet. Click 'Start Demo Patient' to begin._")

            confirm_new_patient_modal = Modal("⚠️ Assessment in Progress", visible=False, elem_classes=["custom-modal"])
            with confirm_new_patient_modal:
                gr.Markdown(
                    "You have an active demo patient session. Starting a new demo patient now will **discard** unsaved progress for the current patient.\n\n"
                    "What would you like to do?"
                )
                with gr.Row():
                    cancel_new_patient_btn = gr.Button("Continue with Current")
                    confirm_discard_btn = gr.Button("Start New Demo Patient", variant="stop")

        with gr.Tab("2. Conduct Interview", id=2, interactive=False) as interview_tab:
            gr.Markdown(TAB_2_INSTRUCTIONS)
            gr.Markdown("---")
            interview_header = gr.Markdown("### Conduct Diagnostic Interview")
            summary = gr.Markdown("_Start a new patient and module to see assessment status._")

            with gr.Row():
                prev_module_btn = gr.Button("⬅️ Previous Module", interactive=False)
                next_module_btn = gr.Button("Start First Module ➡️", variant="secondary", interactive=False)

            with gr.Accordion("Advanced: Jump to a Specific Module", open=False):
                with gr.Row():
                    module_selector = gr.Dropdown(choices=ALLOWED_MODULES, label="Module", interactive=True)
                    load_module_btn = gr.Button("Load Module", interactive=False)

            with gr.Accordion("Module Information (DSM-5 Criteria & Rules)", open=False):
                interviewer_info = gr.Markdown("Start a session to see module information.")

            progress_status = gr.Markdown("")

            with gr.Row():
                with gr.Column(scale=1):
                    assessment = gr.Markdown("...")
                    with gr.Group():
                        probes = gr.CheckboxGroup(choices=[], label="Probing Questions (check to ask)")
                        ask_probe = gr.Button("Ask Selected Question(s)", interactive=False)
                    with gr.Accordion("Custom Question", open=False):
                        gr.Markdown("Please do not enter names, contact details, or information about real identifiable patients.")
                        custom_q = gr.Textbox(label="Enter a Custom Question", lines=2)
                        ask_custom_btn = gr.Button("Ask Custom Question")

                    gr.Markdown("#### Answer Realism Rating")
                    item_realism_rating = gr.Slider(
                        minimum=0, maximum=5, step=1, value=0,
                        label="Rate the realism of the patient's response to this specific item (0=Not Rated, 1=Poor, 5=Excellent)",
                        interactive=True,
                    )

                    gr.Markdown("#### Clinician's Decision for this Item")
                    single_choice = gr.Radio(choices=[], label="Is the diagnostic criterion met?", visible=False, interactive=True)
                    multi_choice = gr.CheckboxGroup(choices=[], label="Multiple-selection", visible=False, interactive=True)
                    other_choice = gr.CheckboxGroup(choices=[], label="Additional categories", visible=False, interactive=True)
                    saved_status = gr.Markdown("_No selections saved yet._")

                with gr.Column(scale=1):
                    chat = gr.Chatbot(label="Clinician ↔ Simulated Patient", height=700, type="messages")

            with gr.Row():
                module_comments = gr.Textbox(label="Clinician Comments (for current module)", lines=2, scale=3)
                module_rating = gr.Slider(0, 5, 1, value=0, label="Overall Module Realism (0-5)", scale=1)

            with gr.Row():
                prev_btn = gr.Button("← Previous Item", size="sm")
                next_btn = gr.Button("Next Item →", size="sm")

            screen_skip_modal = Modal("Screening Information", visible=False, elem_classes=["custom-modal"])
            with screen_skip_modal:
                gr.Markdown(
                    "**Note:** Clinically, a 'No' on the screening question indicates core symptoms are absent.\n\n"
                    "• To **stop** now: click OK, finish your rating/comments, then click **'Save and Continue'**.\n"
                    "• To **continue**: click OK and proceed to the next item."
                )
                with gr.Row():
                    ok_btn = gr.Button("OK", variant="primary")

        with gr.Tab("3. Review & Save Record", id=3, interactive=False) as review_tab:
            gr.Markdown(TAB_3_INSTRUCTIONS)
            gr.Markdown("---")
            gr.Markdown("### Part A: Review Assessment")
            generate_summary_btn = gr.Button("Generate Full Diagnostic Summary", variant="secondary", interactive=False)
            session_summary_display = gr.Markdown("_Click the button above to generate the summary._")

            gr.Markdown("---")
            gr.Markdown("### Part B: Save Demo Feedback Record")
            finalize_btn = gr.Button("Save Demo Feedback Record", variant="primary", interactive=False)
            finalize_status = gr.Markdown("If you consented to feedback collection, this uploads to the separate demo feedback dataset. Otherwise, nothing is saved.")

            dataset_file = gr.File(label="Saved Demo File", interactive=False, visible=False)
            download_btn = gr.Button("Download CSV Dataset", visible=False)

        with gr.Tab("4. System Evaluation", id=4) as validation_tab:
            gr.Markdown("### Standardized Usability Scale (SUS)")
            gr.Markdown("Please rate the overall system after trying the demo (1 = Strongly Disagree, 5 = Strongly Agree). SUS responses are saved only if you consented to demo feedback collection on the first tab.")
            with gr.Row():
                sus_1 = gr.Slider(1, 5, step=1, value=3, label="1. I think that I would like to use this system frequently.")
                sus_2 = gr.Slider(1, 5, step=1, value=3, label="2. I found the system unnecessarily complex.")
            with gr.Row():
                sus_3 = gr.Slider(1, 5, step=1, value=3, label="3. I thought the system was easy to use.")
                sus_4 = gr.Slider(1, 5, step=1, value=3, label="4. I think that I would need the support of a technical person to use this system.")
            with gr.Row():
                sus_5 = gr.Slider(1, 5, step=1, value=3, label="5. I found the various functions in this system were well integrated.")
                sus_6 = gr.Slider(1, 5, step=1, value=3, label="6. I thought there was too much inconsistency in this system.")
            with gr.Row():
                sus_7 = gr.Slider(1, 5, step=1, value=3, label="7. I would imagine that most people would learn to use this system very quickly.")
                sus_8 = gr.Slider(1, 5, step=1, value=3, label="8. I found the system very cumbersome to use.")
            with gr.Row():
                sus_9 = gr.Slider(1, 5, step=1, value=3, label="9. I felt very confident using the system.")
                sus_10 = gr.Slider(1, 5, step=1, value=3, label="10. I needed to learn a lot of things before I could get going with this system.")

            overall_feedback = gr.Textbox(label="Additional comments about the demo (optional)", lines=3)
            submit_validation_btn = gr.Button("Submit Evaluation", variant="primary")
            validation_status = gr.Markdown("")

            def save_validation(ds, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, comments):
                if not ds:
                    return "Error: No session state."
                if DEMO_MODE and not getattr(ds, "feedback_saving_enabled", False):
                    return feedback_not_saved_message()
                data = {
                    "Timestamp": datetime.utcnow().isoformat(),
                    "DemoSessionID": ds.demo_session_id,
                    "ClinicianID": ds.clinician_id,
                    "ConsentGiven": "Yes" if ds.consent_given else "No",
                    "ConsentTimestampUTC": ds.consent_timestamp_utc,
                    "ConsentVersion": ds.consent_version,
                    "Profession": ds.profession,
                    "YrsDiagExp": ds.yrs_diag_exp,
                    "YrsProfExp": ds.yrs_prof_exp,
                    "SUS_1": s1, "SUS_2": s2, "SUS_3": s3, "SUS_4": s4, "SUS_5": s5,
                    "SUS_6": s6, "SUS_7": s7, "SUS_8": s8, "SUS_9": s9, "SUS_10": s10,
                    "AdditionalComments": comments or "",
                }
                df = pd.DataFrame([data])
                try:
                    sync_main_artifacts_from_hub()
                    if not os.path.exists(SUS_DATASET_PATH):
                        df.to_csv(SUS_DATASET_PATH, index=False)
                    else:
                        df.to_csv(SUS_DATASET_PATH, mode='a', header=False, index=False)
                    upload_to_hub(SUS_DATASET_PATH, path_in_repo=os.path.basename(SUS_DATASET_PATH))
                    return "✅ Demo usability feedback saved to the separate demo feedback dataset."
                except Exception as e:
                    return f"❌ Error saving data: {str(e)}"

            submit_validation_btn.click(
                save_validation,
                inputs=[dataset_state, sus_1, sus_2, sus_3, sus_4, sus_5, sus_6, sus_7, sus_8, sus_9, sus_10, overall_feedback],
                outputs=[validation_status],
            )

    load_outputs = [
        dataset_state, session, anon_user_display, patient_status, patient_id,
        next_module_btn, load_module_btn, summary, generate_summary_btn, finalize_btn,
        interview_tab, review_tab, main_tabs, probes, single_choice, multi_choice, other_choice,
        item_realism_rating, module_rating, assessment, saved_status, progress_status,
        interview_header, interviewer_info, prev_module_btn, module_selector, chat, ask_probe, module_comments,
        profession_input, yrs_diag_input, yrs_prof_input,
    ]
    demo.load(build_load_view, outputs=load_outputs)

    def start_trigger(ds, prof, y_d, y_p, feedback, profile: Optional[gr.OAuthProfile] = None, request: Optional[gr.Request] = None):
        if ds.patient_id:
            return (gr.update(visible=True),) + tuple(gr.skip() for _ in range(20))
        gen_results = new_patient(ds, prof, y_d, y_p, feedback, profile=profile, request=request)
        return (gr.update(visible=False),) + gen_results

    btn_outputs = [
        confirm_new_patient_modal, dataset_state, patient_status, chat, patient_id,
        next_module_btn, load_module_btn, summary, generate_summary_btn, finalize_btn,
        interview_tab, review_tab, main_tabs, session, probes, single_choice, multi_choice,
        other_choice, item_realism_rating, module_rating, anon_user_display,
    ]

    new_patient_btn.click(
        start_trigger,
        inputs=[dataset_state, profession_input, yrs_diag_input, yrs_prof_input, feedback_choice],
        outputs=btn_outputs,
    )

    def confirm_handler(ds, prof, y_d, y_p, feedback, profile: Optional[gr.OAuthProfile] = None, request: Optional[gr.Request] = None):
        return (gr.update(visible=False),) + new_patient(ds, prof, y_d, y_p, feedback, profile=profile, request=request)

    confirm_discard_btn.click(
        confirm_handler,
        inputs=[dataset_state, profession_input, yrs_diag_input, yrs_prof_input, feedback_choice],
        outputs=btn_outputs,
    )

    cancel_new_patient_btn.click(lambda: gr.update(visible=False), outputs=[confirm_new_patient_modal])

    change_module_outputs = [
        dataset_state, session, summary,
        module_comments, module_rating,
        chat, assessment, probes, single_choice, multi_choice, other_choice,
        item_realism_rating,
        saved_status, progress_status,
        interview_header, interviewer_info,
        prev_module_btn, next_module_btn, module_selector,
        generate_summary_btn, finalize_btn,
        interview_tab, review_tab, main_tabs,
        screen_skip_modal,
    ]

    next_module_btn.click(
        click_next_or_finish,
        inputs=[dataset_state, session, module_comments, module_rating],
        outputs=change_module_outputs,
    )
    prev_module_btn.click(
        lambda ds, cs, c, mr: change_module(ds, cs, c, mr, "prev", None),
        inputs=[dataset_state, session, module_comments, module_rating],
        outputs=change_module_outputs,
    )
    load_module_btn.click(
        lambda ds, cs, c, mr, mn: change_module(ds, cs, c, mr, "direct", mn),
        inputs=[dataset_state, session, module_comments, module_rating, module_selector],
        outputs=change_module_outputs,
    )

    nav_outputs = [dataset_state, session, assessment, probes, single_choice, multi_choice, other_choice, item_realism_rating, saved_status, progress_status]
    next_btn.click(lambda ds, st: nav_item(ds, st, "next"), inputs=[dataset_state, session], outputs=nav_outputs)
    prev_btn.click(lambda ds, st: nav_item(ds, st, "prev"), inputs=[dataset_state, session], outputs=nav_outputs)

    ask_probe.click(ask_selected, inputs=[dataset_state, session, probes], outputs=[dataset_state, session, chat, probes])
    ask_custom_btn.click(ask_custom, inputs=[dataset_state, session, custom_q], outputs=[dataset_state, session, chat, custom_q])
    probes.change(lambda x: gr.update(interactive=bool(x)), inputs=[probes], outputs=[ask_probe])

    for inp in [single_choice, multi_choice, other_choice, item_realism_rating]:
        inp.change(
            auto_save_selection,
            inputs=[dataset_state, session, single_choice, multi_choice, other_choice, item_realism_rating],
            outputs=[dataset_state, session, saved_status, screen_skip_modal],
        )

    generate_summary_btn.click(
        generate_full_summary,
        inputs=[dataset_state],
        outputs=[dataset_state, session_summary_display, generate_summary_btn, finalize_btn],
    )

    review_tab.select(
        generate_full_summary,
        inputs=[dataset_state],
        outputs=[dataset_state, session_summary_display, generate_summary_btn, finalize_btn],
    )

    finalize_btn.click(
        finalize_patient,
        inputs=[dataset_state],
        outputs=[
            dataset_state, dataset_file, finalize_status,
            new_patient_btn, next_module_btn, load_module_btn,
            generate_summary_btn, finalize_btn,
            interview_tab, review_tab, main_tabs,
        ],
    )

    download_btn.click(download_dataset, inputs=[dataset_state], outputs=[dataset_file])
    ok_btn.click(lambda: gr.update(visible=False), outputs=[screen_skip_modal])

if __name__ == "__main__":
    demo.queue().launch()