# Neural Correlates of Memory and Attention Study 2025
# Demographics + Visual Imagery Language Elements Survey (single dialog)
# Collects: Participant ID, Age, Sex assigned at Birth, Ethnicity (+ optional "Other" write-in), Handedness,
#           Mother tongue, Visual imagery language-elements (Y/N/Unsure) + If yes, which language
# Saves: ./data/<participantID>/<participantID>_demographics.xlsx  (fallback CSV if Excel engine missing)

from psychopy import gui, core
import pandas as pd
import datetime
import os
import sys

def show_alert(title, message):
    """Simple alert dialog replacement for gui.alertDlg"""
    alert = gui.Dlg(title=title)
    alert.addText(message)
    alert.show()

def save_excel_or_csv(df, xlsx_path, csv_path):
    """Save as Excel if possible, otherwise fall back to CSV."""
    try:
        try:
            with pd.ExcelWriter(xlsx_path, engine="openpyxl") as w:
                df.to_excel(w, index=False, sheet_name="Demographics")
        except Exception:
            with pd.ExcelWriter(xlsx_path, engine="xlsxwriter") as w:
                df.to_excel(w, index=False, sheet_name="Demographics")
        print(f"✅ Saved: {os.path.abspath(xlsx_path)}")
    except Exception as e:
        print(f"⚠️ Excel write failed ({e}). Writing CSV fallback.")
        df.to_csv(csv_path, index=False, encoding="utf-8")
        print(f"💾 Saved: {os.path.abspath(csv_path)}")

# --- Determine script directory ---
try:
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
    # __file__ may not exist in some environments (e.g., interactive)
    SCRIPT_DIR = os.path.dirname(os.path.abspath(sys.argv[0])) if sys.argv and sys.argv[0] else os.getcwd()

# --- Dropdown options ---
ethnicity_options = [
    "Please select",
    "American Indian or Alaska Native",
    "Asian",
    "Black or African American",
    "Hispanic or Latino",
    "Middle Eastern or North African (MENA)",
    "Native Hawaiian or Pacific Islander",
    "White",
    "Prefer not to answer",
    "Other"
]
sex_options = ["Please select", "Female", "Male", "Intersex", "Prefer not to answer", "Another / Self-describe"]
hand_options = ["Please select", "Right", "Left", "Ambidextrous", "Prefer not to say"]

# NEW: imagery yes/no options
yesno_options = ["Please select", "Yes", "No", "Not sure / Prefer not to say"]

# --- Single dialog ---
dlg = gui.Dlg(title="Demographics & Visual Imagery Survey")
dlg.addText("Please fill in the following information.")
dlg.addText("")  # extra spacer

dlg.addField("Participant ID:", "")  # 0
dlg.addField("Age (years):", "")     # 1
dlg.addField("Sex assigned at birth:", choices=sex_options, initial=0)  # 2
dlg.addField("Ethnicity:", choices=ethnicity_options, initial=0)        # 3
dlg.addField("If other, please specify:", "")                           # 4
dlg.addField("Handedness:", choices=hand_options, initial=0)            # 5
dlg.addField("Mother tongue / first language:", "")                     # 6
dlg.addText("When you visualize something in your mind (for example, a scene, an object, or a person), \n"
            "do you experience any language-related elements (such as written words, signs, or labels)? \n" \
            "If yes, which language do these elements most often appear in?")
dlg.addField("Language-related elements present in visual imagery?", choices=yesno_options, initial=0)  # 7
dlg.addField("If YES: Which language most often?", "")                 # 8

ok = dlg.show()
if not ok:
    core.quit()

# --- Extract responses ---
pid         = (dlg.data[0] or "").strip()
age_str     = (dlg.data[1] or "").strip()
sex         = dlg.data[2]
eth_cat     = dlg.data[3]
eth_other   = (dlg.data[4] or "").strip()
hand        = dlg.data[5]
mother_tongue = (dlg.data[6] or "").strip()
imagery_has_lang = dlg.data[7]  # "Yes"/"No"/"Not sure..." or "Please select"
imagery_lang_text = (dlg.data[8] or "").strip()

# --- Validation ---
if not pid:
    show_alert("Missing Participant ID", "Please enter a Participant ID.")
    core.quit()

if not (age_str.isdigit() and int(age_str) > 0):
    show_alert("Invalid Age", "Please enter a positive integer for age.")
    core.quit()

if sex == "Please select":
    show_alert("Missing Selection", "Please select a value for 'Sex assigned at birth'.")
    core.quit()

if eth_cat == "Please select":
    show_alert("Missing Selection", "Please select a value for 'Ethnicity'.")
    core.quit()

if hand == "Please select":
    show_alert("Missing Selection", "Please select a value for 'Handedness'.")
    core.quit()

if eth_cat == "Other" and not eth_other:
    show_alert("Missing 'Other' ethnicity", "You selected 'Other'. Please specify the ethnicity (or choose a different category).")
    core.quit()

# Mother tongue can be optional, but if you want to require it, uncomment below:
# if not mother_tongue:
#     show_alert("Missing Mother Tongue", "Please enter your mother tongue / first language.")
#     core.quit()

if imagery_has_lang == "Please select":
    show_alert("Missing Selection", "Please select an option for language-related elements in visual imagery.")
    core.quit()

# If they said YES, require a language
if imagery_has_lang == "Yes" and not imagery_lang_text:
    show_alert("Missing Language", "You indicated that language-related elements are present. Please specify which language they most often appear in.")
    core.quit()

# --- Final ethnicity text ---
final_eth = eth_other if eth_cat == "Other" else eth_cat

# --- Save data ---
timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
row = {
    "test": "Demographics",  # specify which test/questionnaire the data comes from
    "participant_id": pid,
    "timestamp": timestamp,
    "age": int(age_str),
    "sex_assigned_at_birth": sex,
    "ethnicity_category": eth_cat,
    "ethnicity_text": eth_other if eth_other else None,
    "ethnicity": final_eth,
    "handedness": hand,

    # NEW fields
    "mother_tongue": mother_tongue if mother_tongue else None,
    "visual_imagery_language_elements": imagery_has_lang,  # Yes/No/Not sure...
    "visual_imagery_language_most_often": imagery_lang_text if imagery_lang_text else None,
}
df = pd.DataFrame([row])

# --- Build the save folder structure ---
data_dir = os.path.join(SCRIPT_DIR, "data")
participant_dir = os.path.join(data_dir, pid)
os.makedirs(participant_dir, exist_ok=True)  # create if not exist

# --- Build full file paths ---
base = os.path.join(participant_dir, f"{pid}_demographics")
xlsx = base + ".xlsx"
csv  = base + ".csv"

save_excel_or_csv(df, xlsx, csv)

core.quit()