# Neural Correlates of Memory and Attention Study 2025
# Demographics Survey (single dialog)
# Collects: Participant ID, Age, Sex assigned at Birth, Ethnicity (+ optional "Other" write-in), Handedness
# 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"]

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

dlg.addField("Participant ID:", "") # empty text field
dlg.addField("Age (years):", "") # empty text field
dlg.addField("Sex assigned at birth:", choices=sex_options, initial=0)
dlg.addField("Ethnicity:", choices=ethnicity_options, initial=0)
dlg.addField("If other, please specify:", "")   # optional unless "Other" selected
dlg.addField("Handedness:", choices=hand_options, initial=0)

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]

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

# --- 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/questionaire the data comes from 
    "participant_id": pid, # participant ID 
    "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,
}
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()