app.py

import json
import os
from dataclasses import dataclass, field
from typing import List, Optional, Dict
from PIL import Image

import streamlit as st
from markdown import markdown
st.set_page_config(layout="wide", page_title="Therapy Session Annotation", initial_sidebar_state="collapsed")

import re, textwrap, html as py_html
from pathlib import Path
from streamlit.components.v1 import html as st_html

import numpy as np
import pandas as pd
from fsspec.implementations.local import LocalFileSystem
from huggingface_hub import HfFileSystem

import streamlit.components.v1 as components

@dataclass
class Field:
    type: str
    title: str
    name: str = None
    mandatory: bool = True
    # if value of field is in the list of those values, makes following siblings mandatory
    following_mandatory_values: list = False
    skip_mandatory: bool = False
    help: Optional[str] = None
    children: Optional[List['Field']] = None
    other_params: Optional[Dict[str, object]] = field(default_factory=lambda: {})

# Function to get user ID from URL
def get_param_from_url(param):
    user_id = st.query_params.get(param, "")
    return user_id


########################################################################################

# 'local' or 'hf'. hf is for Hugging Face file system but has limits on the number of access per hour
filesystem = 'hf'
input_repo_path = 'datasets/emvecchi/therapy_annotation'
output_repo_path = 'datasets/emvecchi/therapy_annotation/pilot'

to_annotate_file_name = 'to_annotate.csv'  # CSV file to annotate
COLS_TO_SAVE = ['dialogue_id','dialogue_name','generatedPatient']

agreement_labels = ['strongly disagree', 'disagree', 'neither agree no disagree', 'agree', 'strongly agree']
mentioned_y_n = ['yes, explicitly', 'yes, implicitly', 'no', 'not sure', 'other']
degree_labels = ['Not at all', 'Slightly', 'Moderately', 'Very much', 'Extremely']
quality_labels = ['very poor', 'poor', 'acceptable', 'good', 'very good']
priority_labels = ['not a priority', 'low priority', 'neutral', 'moderate priority', 'high priority']
yes_no_labels = ['no','yes']
yes_no_other_labels = ['no','yes','other']
default_labels = agreement_labels

MAX_RUPTURE_MARKERS = 8  # 1 initial + up to 7 extra
RUPTURE_LINE_INPUT_TYPE = "number"  # or "text" if your line numbers aren't pure ints

rupture_choices = ["",
                   "Denial",
                   "Minimal response",
                   "Abstract communication",
                   "Avoidant storytelling and/or topic-shifting",
                   "Deferential / appeasing behavior",
                   "Content/affect split (i.e. talking content vs emotional tone mismatch)",
                   "Self-criticism and/or hopelessness",
                   "Complaints or concerns about the therapist",
                   "Patient rejects therapist intervention",
                   "Complaints/concerns about the activities of therapy",
                   "Complaints/concerns about the parameters of therapy (e.g. scheduling, rules)",
                   "Complaints/concerns about progress in therapy",
                   "Patient defends self against therapist",
                   "Efforts to control or pressure the therapist",
                    "Other"]
          
default_choices = rupture_choices

consent_text = '''
## Consent Form
You will be asked to take part in a research study. Before you decide to take part in this study, it is important that you understand why the study is being done and what it involves. Please read the following information carefully.
________________________________________________________________________________________
Project title: [...]\\
Researchers: [...]\\
Institute: Institute for Machine Speech Processing (IMS)\\
University: University of Stuttgart\\
Contact: eva-maria.vecchi@ims.uni-stuttgart.de
_________________________________________________________________________________________
### Description of the research study

In this study, we investigate [...]

The intended use of the results of this study includes [...]

**Time required:** Your participation will take up to an estimated [...] hours. The time required may vary on an individual basis.

**Risks and benefits:** The risks to your participation in this online survey are those associated with basic computer tasks, including boredom, fatigue, mild stress, or breach of confidentiality. Some of the topics discussed in the online posts to be annotated may include violence, suicide or rape. The only benefit to you is the learning experience from participating in a research study. The benefit to society is the contribution to scientific knowledge

**Compensation:** You will be compensated for participating in this study. If you are interested, we will also be more than happy to share more information about our research with you.

**Voluntary participation:** Your participation in this study is voluntary. It is your decision whether or not to participate in this study. If you decide to participate in this study, you will be asked to confirm this consent form ("I agree."). Even after signing the consent form, you can withdraw from participation at any time and without giving any reason. Partial data will not be analysed.

**Confidentiality:** Your responses to this experiment will be anonymous. Please do not share any

Information that can be used to identify you. The researcher(s) will make every effort to maintain your confidentiality.

**Contact:** If at any time you have questions about this study or would like to report any adverse effects due to this study, please contact the researcher(s).

**Trigger Warning:** [...]

### Consent:
Please indicate, in the box below, that you are at least 18 years old, have read and understood this consent form, are comfortable using the English and German languages to complete the survey, and you agree to participate in this online research survey.
- *I am age 18 or older.*
- *I have read this consent form or had it read to me.*
- *I am comfortable using the English language to participate in this survey.*
- *I agree to participate in this research and I want to continue with the survey.*
'''

fields: List[Field] = [
    Field(name="patient", type="input_col", title=" "),
    Field(type="expander", title="**Session Transcription:** *(expand)*", children=[
        Field(name="dialogue_name", type="input_col", title=""),
    ]),

    Field(type="container", title="##### Session-/Patient-Specific Properties", children=[
        Field(name="diagnosis", type="likert_radio",
                   title="Is there mention of a specific [possible] clinical diagnosis?", other_params={'labels': mentioned_y_n}, mandatory=True),
        Field(name="stressor", type="likert_radio",
                    title="Does the patient mention sources of psychological pressure that contribute to the development or worsening of the patient's state?", 
                    other_params={'labels': mentioned_y_n}, mandatory=True),
    ]),

    Field(type="container", title="#### Rupture Markers", children=[
        Field(type="expander", title="**Overview of rupture markers** *(expand)*", children=[
            Field(type="markdown", title="", other_params={"path": "rupture_marker_overview.md"}),
        ]),
        Field(name="rupture_marker", type="rupture_markers",
          title="Select rupture markers noted in the session, include line numbers where rupture is found.", mandatory=False),
    ]),

    Field(type="container", title="#### True-To-Patient-Prompt Features", children=[
        Field(type="expander", title="**Patient Role Description:** *(expand)*", children=[
        Field(name="role_name", type="input_col", title=""),
    ]),
        Field(name="emotion", type="likert_radio",
                title="To what extent did the simulated patient’s emotional expressions (e.g., frustration, sadness, guardedness) align with those expected for this role?", 
                other_params={"labels": degree_labels}, mandatory=True),
        Field(name="conversationalStyle", type="likert_radio",
                title="How well did the patient consistently exhibit the conversational style they were assigned, and did this feel realistic compared to actual patients with this style?", 
                other_params={"labels": degree_labels}, mandatory=True),
        Field(name="maladaptive", type="likert_radio",
                title="How accurately did the patient’s statements reflect distorted thoughts consistent with the underlying cognitive model?", 
                other_params={"labels": degree_labels}, mandatory=True),
        Field(name="fidelity", type="likert_radio",
                title="Overall, how closely did the patient reflect real patients in terms of realism, variability, and challenge during the session?", 
                other_params={"labels": degree_labels}, mandatory=True),
    ]),

    Field(type="container", title="#### Interpersonal therapeutic competence", children=[
        Field(name="v_expressiveness", type="likert_radio",
                  title="**Verbal expressiveness.** How fluent, clear, and confident is the person’s verbal communication (regardless of content)?", 
                  other_params={'labels': ['Very unclear','Unclear','Average','Clear','Very Clear']}, 
                  mandatory=True, following_mandatory_values=[1]),
        Field(name="hope", type="likert_radio",
                title="**Hope and positive expectation.** To what extent does the person convey optimism, confidence, and positive expectations for the patient’s future or therapy?", 
                other_params={'labels': ['Hopeless','Discouraging','Neutral','Hopeful','Very hopeful']}, 
                mandatory=True, following_mandatory_values=[1]),
        Field(name="credible", type="likert_radio",
                title="**Ability to be credible and convincing.** How credible and convincing is the person when presenting a new perspective or interpretation to the patient?", 
                other_params={'labels': ['Unconvincing','Weak','Neutral','Convincing','Very convincing']}, 
                mandatory=True, following_mandatory_values=[1]),
        Field(name="e_expressiveness", type="likert_radio",
                title="**Emotional expressiveness.** How well does the person’s voice and manner convey emotional engagement and intensity (without being hostile)?", 
                other_params={'labels': ['Flat','Low','Moderate','Expressive','Very expressive']}, 
                mandatory=True, following_mandatory_values=[1]),
        Field(name="acceptance", type="likert_radio",
                title="**Acceptance, understanding, and warmth.** To what extent does the person show care, acceptance, and warmth toward the patient (without judgment)?", 
                other_params={'labels': ['Cold','Limited','Neutral','Warm','Very warm']}, 
                mandatory=True, following_mandatory_values=[1]),
        Field(name="empathy", type="likert_radio",
                title="**Empathy.** How well does the person demonstrate an accurate, deeper understanding of the patient’s subjective experience?", 
                other_params={'labels': ['Distorted','Superficial','Partial','Accurate','Deep']}, 
                mandatory=True, following_mandatory_values=[1]),
        Field(name="alliance", type="likert_radio",
                title="**Ability to form alliances.** How effectively does the person create a collaborative “working together” atmosphere with the patient?", 
                other_params={'labels': ['Blocking','Weak','Moderate','Cooperative','Strong']}, 
                mandatory=True, following_mandatory_values=[1]),
        Field(name="tensions", type="likert_radio",
                title="**Addressing tensions and crises.** How effectively does the person acknowledge and work through interpersonal tensions or crises in the interaction?", 
                other_params={'labels': ['Avoiding','Poor','Partial','Productive','Resolving']}, 
                mandatory=True, following_mandatory_values=[1]),
    ]),
    
    Field(type="container", title="##### Additional Comments *(optional)*", children=[
        Field(name="other_comments", type="text", title="Please provide any additional details or information:", mandatory=False),
    ]),
]
url_conditional_fields = [
    Field(name="skip", type="skip_checkbox",
          title="*I am uncomfortable annotating this text and  voluntarily skip this instance*", mandatory=False)
]
INPUT_FIELD_DEFAULT_VALUES = {'slider': 0,
                              'text': '',
                              'textarea': '',
                              'checkbox': False,
                              'radio': None,
                              'select_slider': 0,
                              'multiselect': [],
                              'likert_radio': None,
                              'y_n_radio': None}
SHOW_HELP_ICON = False
SHOW_VALIDATION_ERROR_MESSAGE = True 

########################################################################################
if filesystem == 'hf':
    HF_TOKEN = os.environ.get("HF_TOKEN_WRITE")
    print("is none?", HF_TOKEN is None)
    hf_fs = HfFileSystem(token=HF_TOKEN)
else:
    hf_fs = LocalFileSystem()

def get_start_index():
    if hf_fs.exists(output_repo_path + '/' + get_base_path()):
        files = hf_fs.ls(output_repo_path + '/' + get_base_path())
        return len(files) #-1
    else:
        return -2

def read_data():
    with hf_fs.open(input_repo_path + '/' + to_annotate_file_name) as f:
        return pd.read_csv(f)

def read_saved_data():
    _path = get_path()
    if hf_fs.exists(output_repo_path + '/' + _path):
        with hf_fs.open(output_repo_path + '/' + _path) as f:
            try:
                return json.load(f)
            except json.JSONDecodeError as e:
                print(e)
    return None
    
# Write a remote file
def save_data(data):
    if not hf_fs.exists(f"{output_repo_path}/{get_base_path()}"):
        hf_fs.mkdir(f"{output_repo_path}/{get_base_path()}")
    with hf_fs.open(f"{output_repo_path}/{get_path()}", "w") as f:
        f.write(json.dumps(data))

def get_base_path():
    return f"{st.session_state.user_id}"

def get_path():
    return f"{get_base_path()}/{st.session_state.current_index}.json"

def display_dialogue(hf_path: str):
    with hf_fs.open(hf_path, "rb") as f:
        txt = f.read().decode("utf-8")
    txt = load_text(hf_path)
    st.markdown(txt,
    unsafe_allow_html=True)

def display_image(image_path):
    with hf_fs.open(image_path) as f:
        img = Image.open(f)
        st.image(img, caption='', use_column_width=True)

def _ensure_key(key: str, default):
    if key not in st.session_state:
        st.session_state[key] = default
    return st.session_state[key]
    
SPEAKER_RE = re.compile(r'^\s*\*\*(T|P):\*\*\s*(.*)$')

def read_md(path: str) -> str:
    """Read .md file from HF dataset or local disk."""
    if filesystem == 'hf':
        with hf_fs.open(path, "rb") as f:
            return f.read().decode("utf-8")
    else:
        # local path
        return Path(path).read_text(encoding="utf-8")

def wrap_para(text: str, width: int) -> list[str]:
    if not text.strip():
        return [""]
    return textwrap.wrap(
        text.strip().replace("\u00A0", " "),
        width=width,
        break_long_words=False,
        break_on_hyphens=False,
        drop_whitespace=True,
        replace_whitespace=True,
    ) or [""]

def md_dialogue_to_visual_lines(md_text: str, width: int) -> list[str]:
    """Return a list of wrapped visual lines with <strong>T:/P:</strong> on first line."""
    md_text = md_text.replace("\r\n", "\n").replace("\r", "\n").strip("\n")
    paragraphs = re.split(r"\n\s*\n", md_text)
    out = []
    for p in paragraphs:
        p = p.strip()
        if not p:
            out.append("")
            continue
        m = SPEAKER_RE.match(p)
        if m:
            speaker, content = m.group(1), m.group(2)
            wrapped = wrap_para(content, width)
            out.append(f"<strong>{speaker}:</strong> {py_html.escape(wrapped[0])}".rstrip())
            for w in wrapped[1:]:
                out.append(py_html.escape(w))
        else:
            for w in wrap_para(p, width):
                out.append(py_html.escape(w))
        out.append("")  # blank between paragraphs
    if out and out[-1] == "": out.pop()
    return out

def render_markdown_simple(md_path, height_px=520):
    md = read_md(md_path)
    st.markdown(
        f"<div style='max-height:{height_px}px; overflow:auto'>{markdown(md, extensions=['extra','codehilite'])}</div>",
        unsafe_allow_html=True
    )

def render_dialogue(md_path: str,
                    width_chars: int = 100,
                    height_px: int = 520,
                    font_family: str = "system-ui, -apple-system, 'Segoe UI', Roboto, Helvetica, Arial, 'Noto Sans', sans-serif",
                    font_size: str = "1.05rem",
                    show_border: bool = False):
    """Render the .md dialogue with line numbers inside an iframe to avoid Streamlit CSS quirks."""
    text = read_md(md_path)
    lines = md_dialogue_to_visual_lines(text, width_chars)

    border_css = "1px solid #e6e6e6" if show_border else "none"
    radius_css = ".6rem" if show_border else "0"
    padding_css = ".8rem 1rem" if show_border else "0"

    # Build stable HTML (table keeps gutter aligned)
    rows = "\n".join(
        f"<tr><td class='num'>{i}</td><td class='txt'>{ln or '&nbsp;'}</td></tr>"
        for i, ln in enumerate(lines, 1)
    )

    html_doc = f"""
<!doctype html>
<meta charset="utf-8">
<style>
  :root {{
    --font-text: {font_family};
    --font-size: {font_size};
  }}
  html, body {{
    margin:0; padding:0; background:transparent;
    font-family: var(--font-text); font-size: var(--font-size);
    -webkit-font-smoothing: antialiased; -moz-osx-font-smoothing: grayscale;
    font-synthesis: none;  /* avoid auto fake-bold */
  }}
  .root {{
    border:{border_css}; border-radius:{radius_css}; padding:{padding_css};
    height:{height_px}px; overflow:auto; background:transparent;
    line-height:1.5;
  }}
  table {{ border-collapse:collapse; table-layout:fixed; width:max-content; max-width:100%; }}
  td {{ padding:0; vertical-align:top; }}
  .num {{
    width:4ch; padding-right:1ch; text-align:right; color:rgba(0,0,0,.55);
    user-select:none; font-variant-numeric: tabular-nums;
    font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, 'Liberation Mono', monospace;
  }}
  .txt {{
    white-space: pre-wrap; word-break: break-word; max-width:{width_chars}ch;
  }}
  strong {{ font-weight:700; }}
</style>
<div class="root">
  <table>
    {rows}
  </table>
</div>
"""
    # iframe height = inner height; no external scrollbars flashing
    st_html(html_doc, height=height_px + (16 if show_border else 0), scrolling=False)


def render_rupture_markers_widget(index: int, choices: list[str]):
    """Dynamic list of (rupture_marker_i, line_a_i, line_b_i) with Add/Remove.
       Works inside a st.form by using form_submit_button instead of callbacks."""
    count_key = f"rupture_count_{index}"
    if count_key not in st.session_state:
        st.session_state[count_key] = 1

    count = st.session_state[count_key]

    st.markdown("**What rupture markers are found?** *(one per row; max 8)*")

    for i in range(1, count + 1):
        col1, col2, col3 = st.columns([3, 1, 1])

        marker_key = f"rupture_marker_{i}_{index}"
        a_key = f"rupture_marker_{i}_line_a_{index}"
        b_key = f"rupture_marker_{i}_line_b_{index}"

        with col1:
            current = st.session_state.get(marker_key, "")
            if marker_key in st.session_state:
                st.selectbox(f"Rupture marker {i}", options=choices, key=marker_key)
            else:
                default_idx = choices.index(current) if current in choices else 0
                st.selectbox(f"Rupture marker {i}", options=choices, key=marker_key, index=default_idx)
#            st.selectbox(
#                f"Rupture marker {i}", options=choices, key=marker_key,
#                index=(choices.index(st.session_state.get(marker_key))
#                       if st.session_state.get(marker_key) in choices else 0)
#                if st.session_state.get(marker_key) is not None else 0
#            )
        with col2:
            if RUPTURE_LINE_INPUT_TYPE == "number":
                st.number_input("Lines: from", min_value=0, step=1, key=a_key)
            else:
                st.text_input("Lines: from", key=a_key, value=st.session_state.get(a_key, ""))
        with col3:
            if RUPTURE_LINE_INPUT_TYPE == "number":
                st.number_input("Lines: to", min_value=0, step=1, key=b_key)
            else:
                st.text_input("Lines: to", key=b_key, value=st.session_state.get(b_key, ""))

        st.markdown("---")

    a_col, b_col = st.columns([1, 1])
    with a_col:
        if st.button("➕ Add another rupture marker", disabled=(count >= MAX_RUPTURE_MARKERS)):
            st.session_state[count_key] = min(MAX_RUPTURE_MARKERS, count + 1)
            st.rerun()
    with b_col:
        if st.button("➖ Remove last", disabled=(count <= 1)):
            st.session_state[count_key] = max(1, count - 1)
            st.rerun()

    st.caption(f"{st.session_state[count_key]}/{MAX_RUPTURE_MARKERS} markers")

def _get_value_for_field(f: Field, index: int):
    """Return the session value for input fields or None."""
    if f.type in INPUT_FIELD_DEFAULT_VALUES:
        key = f.name + str(index)
        return st.session_state.get(key, INPUT_FIELD_DEFAULT_VALUES[f.type])
    return None

def _is_default_value(f: Field, val):
    """Check if current value equals the default for this widget type."""
    return val == INPUT_FIELD_DEFAULT_VALUES.get(f.type)

def prime_rupture_markers_from_saved(index: int, data_collected: dict | None):
    """
    Copy saved JSON values (rupture_marker_1, *_line_a, *_line_b) into
    st.session_state keys used by the widget (…_{index}), and set the count.
    Runs once per index to avoid clobbering user edits on reruns.
    """
    if not data_collected:
        return

    primed_flag = f"rupture_primed_{index}"
    if st.session_state.get(primed_flag):
        return  # already restored for this page

    slots = []
    for i in range(1, MAX_RUPTURE_MARKERS + 1):
        m = data_collected.get(f"rupture_marker_{i}")
        a = data_collected.get(f"rupture_marker_{i}_line_a")
        b = data_collected.get(f"rupture_marker_{i}_line_b")
        if any(v is not None for v in (m, a, b)):
            # widget's selectbox uses "" for "no selection"
            slots.append(((m or ""), a, b))

    if not slots:
        return

    # set widget count (at least 1, capped to MAX)
    count_key = f"rupture_count_{index}"
    st.session_state[count_key] = min(MAX_RUPTURE_MARKERS, max(1, len(slots)))

    for i, (m, a, b) in enumerate(slots, start=1):
        st.session_state[f"rupture_marker_{i}_{index}"] = m if m in rupture_choices else (m or "")
        st.session_state[f"rupture_marker_{i}_line_a_{index}"] = a
        st.session_state[f"rupture_marker_{i}_line_b_{index}"] = b

    st.session_state[primed_flag] = True

def validate_current_page(fields: List[Field], index: int) -> bool:
    """
    Walk the field tree and verify mandatory inputs are filled.
    Honors `following_mandatory_values`: if a field's value is in that list,
    subsequent siblings become mandatory.
    """
    ok = True

    def walk(nodes: List[Field], following_required: bool = False):
        nonlocal ok
        # local flag that can be turned on by a node for its following siblings
        req_for_following = following_required

        for f in nodes:
            if f.children:
                # containers/expanders: recurse; pass the current rule downward
                walk(f.children, req_for_following)
                continue

            if f.type in INPUT_FIELD_DEFAULT_VALUES:
                val = _get_value_for_field(f, index)

                # decide if THIS field is required
                required = bool(f.mandatory or req_for_following)

                # if required and default → error
                if required and _is_default_value(f, val):
                    ok = False

                # update the "following required" rule for siblings
                if f.following_mandatory_values:
                    try:
                        if val in f.following_mandatory_values:
                            req_for_following = True
                    except Exception:
                        pass

            # structural types (markdown/input_col/etc.) don’t affect validation

    walk(fields, False)

    if not ok:
        st.error("Please fill in all mandatory fields.")
    return ok


#################################### Streamlit App ####################################

# Function to navigate rows
def navigate(index_change):
    st.session_state.current_index += index_change
    # only works consistently if done before rerun
    js = '''
    <script>
        setTimeout(function() {
            var titleElement = window.parent.document.querySelector("h1");
            if (titleElement) {
                titleElement.scrollIntoView({behavior: "smooth", block: "start"});
            }
        }, 100);
    </script>
    '''
    st.components.v1.html(js, height=0)
    # https://discuss.streamlit.io/t/click-twice-on-button-for-changing-state/45633/2

    # disable text input enter to submit
    # https://discuss.streamlit.io/t/text-input-how-to-disable-press-enter-to-apply/14457/6
    components.html(
        """
        <script>
        const inputs = window.parent.document.querySelectorAll('input');
        inputs.forEach(input => {
            input.addEventListener('keydown', function(event) {
                if (event.key === 'Enter') {
                    event.preventDefault();
                }
            });
        });
        </script>
        """,
        height=0
    )
    st.rerun()

def show_field(f: Field, index: int, data_collected):
    if f.type not in INPUT_FIELD_DEFAULT_VALUES.keys():
        st.session_state.following_mandatory = False
        match f.type:
            case 'input_col':
                value = (
                    st.session_state.data.iloc[index][f.name]
                    if f.name and f.name in st.session_state.data.columns
                    else None
                )

                if f.name == 'image_name' and value:
                    display_image(os.path.join(input_repo_path, 'images', value))

                elif f.name == 'dialogue_name' and value:
                    render_dialogue(
                        os.path.join(input_repo_path, 'dialogues', value),
                        width_chars=115, height_px=520, show_border=False
                    )

                elif f.name == 'role_name' and value:
                    render_markdown_simple(os.path.join(input_repo_path, 'role_descriptions', value), height_px=520)
                        #,
                        #width_chars=115, height_px=520, show_border=False
                    #)

                elif f.name == 'patient' and value:
                    st.markdown(f"#### Patient:&nbsp;&nbsp;{value}")

                elif value not in (None, np.nan, ""):
                    if f.title:
                        st.write(f.title)
                    st.write(value)

            case 'markdown':
                path = f.other_params.get("path") if f.other_params else None
                if path:
                    content = load_text(os.path.join(input_repo_path, path))
                    st.markdown(content, unsafe_allow_html=True)
                else:
                    st.markdown(f.title)

            case 'expander':
                with (st.expander(f.title) if f.type == 'expander' else st.container(border=True)):
                    for child in f.children:
                        show_field(child, index, data_collected)

            case 'container':
                with st.container(border=True):
                    st.markdown(f.title)
                    for child in (f.children or []):
                        show_field(child, index, data_collected)

            case 'skip_checkbox':
                st.checkbox(f.title, key=f.name, value=False)

            case 'rupture_markers':
                prime_rupture_markers_from_saved(index, data_collected)
                render_rupture_markers_widget(index, rupture_choices)

        return

    else:
        key = f.name + str(index)
        # track the logical field name once per page; prep_and_save_data relies on this
        st.session_state.data_inputs_keys.append(f.name)
        default_val = (
            data_collected.get(f.name, INPUT_FIELD_DEFAULT_VALUES[f.type])
            if data_collected else INPUT_FIELD_DEFAULT_VALUES[f.type]
        )
        value = _ensure_key(key, default_val)
        #value = st.session_state[key] if key in st.session_state else \
        #    (data_collected[f.name] if data_collected else INPUT_FIELD_DEFAULT_VALUES[f.type])
        if not SHOW_HELP_ICON:
            f.title = f'**{f.title}**\n\n{f.help}' if f.help else f.title

        validation_error = False
        if st.session_state.form_displayed == st.session_state.current_index:
            if st.session_state.following_mandatory and f.skip_mandatory:
                st.session_state.following_mandatory = False
            if f.following_mandatory_values and st.session_state[key] in f.following_mandatory_values:
                st.session_state.following_mandatory = True
            if f.mandatory or st.session_state.following_mandatory:
                if value == INPUT_FIELD_DEFAULT_VALUES[f.type]:
                #if st.session_state[key] == INPUT_FIELD_DEFAULT_VALUES[f.type]:
                    st.session_state.valid = False
                    validation_error = True
                elif f.following_mandatory_values and st.session_state[key] in f.following_mandatory_values:
                    st.session_state.following_mandatory = True

        if f.mandatory or st.session_state.following_mandatory:
            f.title += " :red[* required!]" if (validation_error and not SHOW_VALIDATION_ERROR_MESSAGE) else' :red[*]'
        f.help = None
        
    value = st.session_state[key] if key in st.session_state else \
            (data_collected[f.name] if data_collected else INPUT_FIELD_DEFAULT_VALUES[f.type])

    match f.type:
        case 'checkbox':
            st.checkbox(f.title, key=key, value=value, help=f.help)

        case 'radio':
            labels = (f.other_params.get('labels')
                      if f.other_params and f.other_params.get('labels')
                      else default_labels)
            st.radio(
                f.title,
                options=range(len(labels)),
                format_func=lambda x: labels[x],
                key=key,
                index=value,
                help=f.help,
                horizontal=False
            )

        case 'slider':
            st.slider(f.title, min_value=0, max_value=6, step=1, key=key,
                      value=value, help=f.help)

        case 'select_slider':
            labels = (f.other_params.get('labels')
                      if f.other_params and f.other_params.get('labels')
                      else default_labels)
            st.select_slider(
                f.title,
                options=[0, 20, 40, 60, 80, 100],
                format_func=lambda x: labels[x // 20],
                key=key,
                value=value,
                help=f.help
            )

        case 'multiselect':
            choices = (f.other_params.get('choices')
                       if f.other_params and f.other_params.get('choices')
                       else default_choices)
            st.multiselect(
                f.title,
                options=choices,
                format_func=lambda x: x,
                key=key,
                default=value,
                max_selections=3,
                help=f.help
            )

        case 'likert_radio':
            labels = (f.other_params.get('labels')
                      if f.other_params and f.other_params.get('labels')
                      else default_labels)
            st.radio(
                f.title,
                options=[0, 1, 2, 3, 4],
                format_func=lambda x: labels[x],
                key=key,
                index=value,
                help=f.help,
                horizontal=True
            )

        case 'y_n_radio':
            labels = (f.other_params.get('labels')
                      if f.other_params and f.other_params.get('labels')
                      else yes_no_labels)
            st.radio(
                f.title,
                options=[0, 1],
                format_func=lambda x: labels[x],
                key=key,
                index=value,
                help=f.help,
                horizontal=True
            )

        case 'text':
            st.text_input(f.title, key=key, value=(value if value is not None else ""),
                          max_chars=None, help=f.help)

        case 'textarea':
            st.text_area(f.title, key=key, value=(value if value is not None else ""),
                         max_chars=None, help=f.help)



def show_fields(fields: List[Field]):
    index = st.session_state.current_index
    data_collected = read_saved_data()
    st.session_state.data_inputs_keys = []  # rebuilt each render
    st.session_state.following_mandatory = False

    for field in fields:
        show_field(field, index, data_collected)

    # mark that the page has been rendered at least once (if you still use this elsewhere)
    st.session_state.form_displayed = st.session_state.current_index

def prep_and_save_data(index, skip_sample):
    payload = {
        'user_id': st.session_state.user_id,
        'index': st.session_state.current_index,
        **(st.session_state.data.iloc[index][COLS_TO_SAVE].to_dict() if 0 <= index < len(st.session_state.data) else {}),
        **{k: st.session_state[k + str(index)] for k in st.session_state.data_inputs_keys},
        'skip': skip_sample
    }
    # normalize rupture markers -> always write 8 slots
    count_key = f"rupture_count_{index}"
    count = st.session_state.get(count_key, 1)
    for i in range(1, MAX_RUPTURE_MARKERS + 1):
        m_key = f"rupture_marker_{i}_{index}"
        a_key = f"rupture_marker_{i}_line_a_{index}"
        b_key = f"rupture_marker_{i}_line_b_{index}"

        # export names WITHOUT the trailing _{index}, so your JSON uses a stable schema:
        out_m = f"rupture_marker_{i}"
        out_a = f"rupture_marker_{i}_line_a"
        out_b = f"rupture_marker_{i}_line_b"

        if i <= count:
            marker_val = st.session_state.get(m_key)
            line_a_val = st.session_state.get(a_key)
            line_b_val = st.session_state.get(b_key)

            # Basic sanity: if numbers, coerce order (optional)
            if RUPTURE_LINE_INPUT_TYPE == "number":
                try:
                    if line_a_val is not None and line_b_val is not None and line_a_val > line_b_val:
                        line_a_val, line_b_val = line_b_val, line_a_val
                except Exception:
                    pass

            payload[out_m] = marker_val if marker_val not in ("", None) else None
            payload[out_a] = line_a_val if line_a_val not in ("", None) else None
            payload[out_b] = line_b_val if line_b_val not in ("", None) else None
        else:
            payload[out_m] = None
            payload[out_a] = None
            payload[out_b] = None

    save_data(payload)

#st.set_page_config(layout='wide')
# Title of the app
st.title("Therapy Session Annotation Task")

st.markdown(
    """<style>
div[data-testid="stMarkdownContainer"] > p {
    font-size: 1rem;
}
    section.main > div {max-width:80rem}
    </style>
    """, unsafe_allow_html=True)

def add_annotation_guidelines():
    guidelines_text = load_text("texts/guidelines.md")
    #st.write(f"username is {st.session_state.user_id}")
    st.markdown(
        f"<details><summary><b>Annotation Guidelines</b></summary><div>{guidelines_text}</div></details><br>",
    unsafe_allow_html=True)

if 'unacceptable_response' in st.session_state and st.session_state.unacceptable_response:
    #error_message = "You are not eligible for this study.<br><br>"
    error_message = "<br>"
    if st.session_state.current_index >= -5:
        error_message += (
            "Thank you for your time! You will receive a small compensation for your contribution up to now. <br><br>"
            f'Please return to the study and copy/paste this code: <b>{failed_sanity_check_code}</b>, or '
            f'<a href="https://app.prolific.com/submissions/complete?cfc={failed_sanity_check_code}">Click Here</a>'
        )
    # Display the error message using custom HTML
    st.markdown(
        f"""
        <div style="background-color: #f8d7da; color: #721c24; padding: 10px; border-radius: 5px; border: 1px solid #f5c6cb;">
            <h4>Error: You are not eligible for this study.</h4> {error_message}
        </div>
        """,
        unsafe_allow_html=True
    )
    st.stop()


# Load the data to annotate
if 'data' not in st.session_state:
    st.session_state.data = read_data()

# user id
user_id_from_url = get_param_from_url("user_id")
if user_id_from_url and not st.session_state.get("user_id"):
    st.session_state.user_id = user_id_from_url
    
# current index
if 'current_index' not in st.session_state:
    start_index = get_start_index()
    if start_index < len(st.session_state.data)-1:
        st.session_state.current_index = start_index
    else:
        st.session_state.current_index = start_index+1
    st.session_state.form_displayed = -2

if get_param_from_url('show_extra_fields'):
    fields += url_conditional_fields
else:
    fields += url_conditional_fields

def add_validated_submit(fields, message):
    st.session_state.form_displayed = st.session_state.current_index
    if st.form_submit_button("Submit"):
        if all(not x for x in fields):
            st.error(message)
        else:
            navigate(1)

def add_checked_submit():
    check = st.checkbox('I agree', key='consent')
    add_validated_submit([check], "Please agree to give your consent to proceed")

def load_text(path: str) -> str:
    if filesystem == 'hf':
        with hf_fs.open(path, "rb") as f:
            return f.read().decode("utf-8")
    else:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()

if st.session_state.current_index == -2:
    with st.form("data_form"):
        st.markdown(consent_text)
        add_checked_submit()

elif st.session_state.current_index == -1:
    if st.session_state.get('user_id'):
        navigate(1)
    else:
        with st.form("user_id_form"):
            st.session_state.user_id = st.text_input("User ID", value="")
            add_validated_submit([st.session_state.user_id], "Please enter a valid user ID")

elif st.session_state.current_index < len(st.session_state.data):
    show_fields(fields)

    # Action buttons
    c1, c2, c3 = st.columns([1,1,6])
    with c1:
        if st.button("Submit"):
            if validate_current_page(fields, st.session_state.current_index):
                with st.spinner("saving"):
                    prep_and_save_data(st.session_state.current_index, ('skip' in st.session_state and st.session_state['skip']))
                st.success("Saved!")
                navigate(1)

elif st.session_state.current_index == len(st.session_state.data):
    st.write(f"**Thank you for taking part in this study!** \n ")
    if st.button("Previous"):
        navigate(-2)

# Navigation buttons
if 0 < st.session_state.current_index < len(st.session_state.data):
    if st.button("Previous"):
        navigate(-1)

if 0 <= st.session_state.current_index < len(st.session_state.data):
    st.write(f"Page {st.session_state.current_index + 1} out of {len(st.session_state.data)}")

st.markdown(
    """<style>
    div[data-testid="InputInstructions"] {
        visibility: hidden;
    }
    </style>""", unsafe_allow_html=True
)