app.py

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

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
import streamlit as st
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'
# path to repo or local file system TODO rename
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']
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="to_moderate", type="y_n_radio",
                   title="Explicit mention of clinical diagnosis?", mandatory=True),
        Field(name="priority_level", type="likert_radio",
                   title="Level of priority", other_params={'labels': priority_labels}, mandatory=True),
    ]),

    Field(type="container", title="**Rupture Markers**", children=[
        Field(name="rupture_markers", type="rupture_markers",
          title="Select rupture markers noted in the session, include line numbers where rupture is found.", mandatory=False),
    #    Field(name="rupture_marker", type="multiselect",
    #                 title="What type of rupture markers are found? *(Select only one)*", 
    #                 other_params={'choices': rupture_choices}, mandatory=True, following_mandatory_values=['Other (please specify)']),
    #    Field(name="rupture_marker", type="text", title="What lines demonstrate the rupture markers you notice?", mandatory=False),
    #    Field(name="rupture_marker_other", type="text", title="*If Other, please specify:*", mandatory=False),
    ]),

    Field(type="container", title="**True-To-Patient-Prompt Properties**", children=[
        Field(name="helpful", type="y_n_radio",
                title="Did the patient remain true to the provided patient prompt?", mandatory=True, 
                following_mandatory_values=[1]),
    ]),

    Field(type="container", title="**Rate the therapist**", children=[
        Field(name="empathy", type="y_n_radio",
                title="Was the therapist empathetic?", mandatory=True, 
                following_mandatory_values=[1]),
    ]),
    
    Field(type="container", title="**Other**", children=[
        Field(name="other_comments", type="text", title="Please provide any additional details or information: *(optional)*", 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) -2
    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='8 most contributing properties', use_column_width=True)




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_dialogue(md_path: str,
                    width_chars: int = 80,
                    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  # start with one block

    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:
            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(f"Lines: {i} – from", min_value=1, step=1, key=a_key)
            else:
                st.text_input(f"Lines: {i} – from", key=a_key, value=st.session_state.get(a_key, ""))

        with col3:
            if RUPTURE_LINE_INPUT_TYPE == "number":
                st.number_input(f"Lines: {i} – to", min_value=1, step=1, key=b_key)
            else:
                st.text_input(f"Lines: {i} – to", key=b_key, value=st.session_state.get(b_key, ""))

        st.markdown("---")

    # Use SUBMIT buttons (legal inside forms) to add/remove rows.
    a_col, b_col = st.columns([1, 1])
    with a_col:
        add_clicked = st.form_submit_button(
            "➕ Add another rupture marker",
            key=f"add_rup_{index}",
            disabled=(count >= MAX_RUPTURE_MARKERS),
        )
    with b_col:
        remove_clicked = st.form_submit_button(
            "➖ Remove last",
            key=f"remove_rup_{index}",
            disabled=(count <= 1),
        )

    if add_clicked and count < MAX_RUPTURE_MARKERS:
        st.session_state[count_key] = count + 1
        st.rerun()
    if remove_clicked and count > 1:
        st.session_state[count_key] = count - 1
        st.rerun()

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

#################################### 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 value and value is not np.nan:
                    st.write(f.title)
                    if f.name == 'image_name':
                        display_image(os.path.join(input_repo_path, 'images', value))
                    elif f.name == 'dialogue_name':
                        #display_dialogue(os.path.join(input_repo_path, 'dialogues', value))
                        render_dialogue(os.path.join(input_repo_path, 'dialogues', value), width_chars=115, height_px=520, show_border=False)
                    elif f.name == 'patient':
                        st.markdown(f"**Patient:**&nbsp;&nbsp;{value}")
                    else:
                        st.write(value)
            case 'markdown':
                st.markdown(f.title)
            case 'expander' | 'container':
                with (st.expander(f.title) if f.type == 'expander' else st.container(border=True)):
                    if f.type == 'container':
                        st.markdown(f.title)
                    for child in f.children:
                        show_field(child, index, data_collected)
            case 'skip_checkbox':
                st.checkbox(f.title, key=f.name, value=False)
            case 'rupture_markers':
                render_rupture_markers_widget(index, rupture_choices)
    else:
        key = f.name + str(index)
        st.session_state.data_inputs_keys.append(f.name)
        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
        
        # form is not displayed for first time
        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 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
        
            # check for any unaccepted values
            if (
                    (f.other_params.get('accepted_values') and
                     value not in f.other_params.get('accepted_values')) or
                    (f.other_params.get('accepted_values_per_sample') and
                     index in f.other_params.get('accepted_values_per_sample') and
                     value not in f.other_params.get('accepted_values_per_sample').get(index))
            ):
                st.session_state.unacceptable_response = 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

        match f.type:
            case 'checkbox':
                st.checkbox(f.title,
                            key=key,
                            value=value, help=f.help)
            case 'radio':
                labels = default_labels if not f.other_params.get('labels') else f.other_params.get('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 = default_labels if not f.other_params.get('labels') else f.other_params.get('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 = default_choices if not f.other_params.get('choices') else f.other_params.get('choices')
                st.multiselect(f.title,
                                options = choices,
                                format_func=lambda x: x,
                                key=key, max_selections=3,
                                default=value, help=f.help)
            case 'likert_radio':
                labels = default_labels if not f.other_params.get('labels') else f.other_params.get('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 = yes_no_labels if not f.other_params.get('labels') else f.other_params.get('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, max_chars=None)
            case 'textarea':
                st.text_area(f.title, key=key, value=value, max_chars=None)

        if validation_error:
            st.session_state.unacceptable_response = False
            st.error(f"Mandatory field")

def show_fields(fields: List[Field]):
    st.session_state.valid = True
    index = st.session_state.current_index
    data_collected = read_saved_data()
    st.session_state.data_inputs_keys = []
    st.session_state.following_mandatory = False

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

    submitted = st.form_submit_button("Submit")
    if submitted:
        if 'unacceptable_response' in st.session_state and st.session_state.unacceptable_response:
            prep_and_save_data(index, ('skip' in st.session_state and st.session_state['skip']))
            st.rerun()
        skip_sample = ('skip' in st.session_state and st.session_state['skip'])
        if not skip_sample and not st.session_state.valid:
            st.error("Please fill in all mandatory fields")
            # st.rerun() # filed-out values are not shown otherwise
        else:
            with st.spinner(text="saving"):
                prep_and_save_data(index, skip_sample)
            st.success("Feedback submitted successfully!")
            navigate(1)

    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()

# batch-specific data subset
#st.session_state.batch = get_param_from_url("batch")

# 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:
#    st.session_state.user_id = user_id_from_url    
# If user_id is in URL → set it
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):
    #add_annotation_guidelines()
    with st.form("data_form"+str(st.session_state.current_index)):
        show_fields(fields)

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
)