tasks/shape_recall.py

import streamlit as st
import sqlite3
import random
import time
import json

def shape_recall(session_id):
    st.subheader("🔷 Shape Recall Task")

    # Load NCPT-style config for shape_recall
    with open("data/ncpt_tasks.json", "r", encoding="utf-8") as f:
        ncpt = json.load(f)
    sr_config = next(task for task in ncpt if task["task"] == "shape_recall")

    # Extract shapes and distractor from JSON
    default_shapes = sr_config["stimuli"]
    default_distractor = sr_config["distractor"]

    # Step 1: Initialize
    if "sr_stage" not in st.session_state:
        st.session_state.sr_stage = "memorize"
        st.session_state.sr_shapes = default_shapes
        st.session_state.sr_start_time = None
        st.session_state.sr_selected = None
        st.session_state.sr_distractor = default_distractor
        st.session_state.sr_options = default_shapes + [default_distractor]
        random.shuffle(st.session_state.sr_options)
        st.session_state.sr_result_displayed = False

    # Step 2: Show shapes to memorize
    if st.session_state.sr_stage == "memorize":
        st.info("Memorize these shapes:")
        st.markdown(" ".join(st.session_state.sr_shapes))
        if st.button("I'm Ready"):
            st.session_state.sr_stage = "recall"
            st.session_state.sr_start_time = time.time()
            st.rerun()

    # Step 3: Recall stage
    elif st.session_state.sr_stage == "recall":
        st.warning("Which one of these shapes was NOT shown?")
        st.session_state.sr_selected = st.radio("Choose:", st.session_state.sr_options, key="sr_radio")

        if st.button("Submit") and not st.session_state.sr_result_displayed:
            rt = round(time.time() - st.session_state.sr_start_time, 2)
            score = 1 if st.session_state.sr_selected == st.session_state.sr_distractor else 0

            # Save to DB
            conn = sqlite3.connect("data/mindmentor.db")
            conn.execute("INSERT INTO cognitive_logs (session_id, task_name, score, reaction_time) VALUES (?, ?, ?, ?)",
                         (session_id, "Shape Recall", score, rt))
            conn.commit()
            conn.close()

            st.success(f"✅ Score: {score} | 🕒 Reaction Time: {rt} seconds")
            st.session_state.sr_result_displayed = True

    # Step 4: Retry option
    if st.session_state.sr_result_displayed:
        if st.button("Try Again"):
            for key in list(st.session_state.keys()):
                if key.startswith("sr_"):
                    del st.session_state[key]
            st.rerun()