# Neural Correlates of Memory and Attention Study 2025
# Vividness of Visual Imagery Questionnaire (VVIQ)
# as described by: MARKS, D.F. (1973), VISUAL IMAGERY DIFFERENCES IN THE RECALL OF PICTURES. British Journal of Psychology, 64: 17-24. https://doi.org/10.1111/j.2044-8295.1973.tb01322.x

import datetime
import pandas as pd
import pyglet
from pyglet import shapes
from pyglet.window import key
from pyglet import event
from pyglet import gl
import sys
import os

# ---- Resolving script directory ----
try:
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
    SCRIPT_DIR = os.path.dirname(os.path.abspath(sys.argv[0])) if sys.argv and sys.argv[0] else os.getcwd()

# ---- Participant ID ----
PID = input("Participant ID: ").strip()
if not PID:
    sys.exit(0)
TIMESTAMP = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# ---- Layout constants ----
FONT_STACK = ["Arial", "Helvetica", "DejaVu Sans", "Liberation Sans", "Sans-Serif"]
FG_WHITE = (255, 255, 255, 255)
BG_BLACK = (0, 0, 0, 255)
ACCENT_WHITE = (255, 255, 255)
ACCENT_BLACK = (0, 0, 0)

# ---- VVIQ Questions ----
# as described by: MARKS, D.F. (1973), VISUAL IMAGERY DIFFERENCES IN THE RECALL OF PICTURES. British Journal of Psychology, 64: 17-24. https://doi.org/10.1111/j.2044-8295.1973.tb01322.x

VVIQ_QUESTIONS = [
    "The exact contour of their face, head, shoulders and body.",
    "Their characteristic poses of head, attitudes of body, etc.",
    "Their precise carriage, length of step, etc., in walking.",
    "The different colours worn in some familiar clothes.",
    "The sun is rising above the horizon into a hazy sky.",
    "The sky clears and surrounds the sun with blueness.",
    "Clouds. A storm blows up, with flashes of lightning.",
    "A rainbow appears.",
    "The overall appearance of the shop from the opposite side of the road.",
    "A window display including colors, shape, and details of individual items for sale.",
    "You are near the entrance. The color, shape, and details of the door.",
    "You enter the shop and go to the counter. The counter assistant serves you. Money changes hands.",
    "The contours of the landscape.",
    "The color and shape of the trees.",
    "The color and shape of the lake.",
    "A strong wind blows on the tree and on the lake causing waves."
]

PROMPTS = {
    range(0, 4):  "Think of some relative or friend whom you frequently see (but who is not with you at present).",
    range(4, 8):  "Visualize a rising sun.",
    range(8, 12): "Think of the front of a shop which you often go to.",
    range(12, 16):"Think of a country scene which involves trees, mountains, and a lake."
}

def header_for_idx(i):
    for r, txt in PROMPTS.items():
        if i in r: return txt
    return ""

labels = [
    "No image at all",
    "Dim and vague",
    "Moderately clear",
    "Clear and reasonably vivid",
    "Perfectly clear and as vivid as normal vision"
]

# ---- Saving data functions ----
def save_excel_or_csv(df, xlsx_path, csv_path):
    try:
        try:
            with pd.ExcelWriter(xlsx_path, engine="openpyxl") as w:
                df.to_excel(w, index=False, sheet_name="VVIQ")
        except Exception:
            with pd.ExcelWriter(xlsx_path, engine="xlsxwriter") as w:
                df.to_excel(w, index=False, sheet_name="VVIQ")
        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)}")

def save_results(interrupted=False):
    data_dir = os.path.join(SCRIPT_DIR, "data")
    participant_dir = os.path.join(data_dir, PID)
    os.makedirs(participant_dir, exist_ok=True)

    # this is only for the VVIQ summary = Total score reported at end of questionaire
    rows = list(data_rows)
    label = "TOTAL_SCORE (partial)" if interrupted else "TOTAL_SCORE"
    rows.append({
        "test": "VVIQ_Summary",
        "participant_id": PID, 
        "timestamp": TIMESTAMP,
        "question_index": None, 
        "question": label,
        "rating": total_score
    })
    df = pd.DataFrame(rows)

    suffix = "_INTERRUPTED" if interrupted else ""
    base = os.path.join(participant_dir, f"{PID}_VVIQ{suffix}")
    xlsx = base + ".xlsx"
    csv  = base + ".csv"
    save_excel_or_csv(df, xlsx, csv)

# ---- Drawing the experiment window ----
win = pyglet.window.Window(fullscreen=True, caption="VVIQ", vsync=True)
win.set_mouse_visible(True)
try:
    gl.glClearColor(0.0, 0.0, 0.0, 1.0)
except Exception:
    pass

key_handler = key.KeyStateHandler()
win.push_handlers(key_handler)
_esc_latched = False

batch_ui = pyglet.graphics.Batch()
batch_rate = pyglet.graphics.Batch()
batch_modal = pyglet.graphics.Batch()

W = H = 0
CENTER_X = CENTER_Y = 0
start_title = None
start_intro = None
start_btn_rect = start_btn_border = None
start_btn_label = None
rating_prompt = None
POST_CLICK_DELAY = 0.35
_pending_next_index = None

def make_center_label(text, y, size=24, bold=False, width_frac=0.8, batch=batch_ui):
    return pyglet.text.Label(
        text, font_size=size, bold=bold,
        font_name=FONT_STACK,
        x=CENTER_X, y=y,
        anchor_x='center', anchor_y='center',
        color=FG_WHITE,
        width=int(W*width_frac), multiline=True, align='center',
        batch=batch
    )

class RatingButton:
    def __init__(self, idx, label_text):
        self.idx = idx
        self.value = idx + 1
        self.cx = self.cy = 0
        self.size = 0
        self.rect = None
        self.number = None
        self.label = None
        self.label_text = label_text
        self.col_width = 0

    def ensure_drawables(self):
        if self.rect is None:
            self.rect = shapes.BorderedRectangle(0, 0, 10, 10, border=3,
                                                 color=ACCENT_BLACK, border_color=ACCENT_WHITE,
                                                 batch=batch_rate)
        if self.number is None:
            self.number = pyglet.text.Label(str(self.value), font_size=36, bold=True,
                                            font_name=FONT_STACK,
                                            x=0, y=0, anchor_x='center', anchor_y='center',
                                            color=FG_WHITE, batch=batch_rate)
        if self.label is None:
            self.label = pyglet.text.Label(
                self.label_text, font_size=18,
                font_name=FONT_STACK,
                x=0, y=0,
                anchor_x='center', anchor_y='center',
                color=FG_WHITE,
                width=200, multiline=True, align='center',
                batch=batch_rate
            )

    def place(self, cx, cy, size, label_y, col_width):
        self.ensure_drawables()
        self.cx, self.cy, self.size, self.col_width = cx, cy, size, col_width
        self.rect.x = int(cx - size/2)
        self.rect.y = int(cy - size/2)
        self.rect.width = self.rect.height = int(size)
        self.number.x, self.number.y = int(cx), int(cy)
        self.label.x, self.label.y = int(cx), int(label_y)
        self.label.width = int(col_width * 0.9)

    def hit(self, mx, my):
        return (self.cx - self.size/2 <= mx <= self.cx + self.size/2 and
                self.cy - self.size/2 <= my <= self.size/2 + self.cy)

buttons = [RatingButton(i, labels[i]) for i in range(5)]

def layout():
    global W, H, CENTER_X, CENTER_Y
    global start_title, start_intro, start_btn_rect, start_btn_border, start_btn_label, rating_prompt

    W, H = win.get_size()
    CENTER_X, CENTER_Y = W//2, H//2

    for b in (batch_ui, batch_rate, batch_modal):
        b._draw_list = []

    start_title = make_center_label("Vividness of Visual Imagery Questionaire", int(CENTER_Y + 0.24*H), size=max(16, int(0.025*H)), bold=True)

    intro_text = ("The following test asks you to imagine different scenarios "
                  "and create a respective image in your head.\n\n"
                  "Press the button to begin.")
    start_intro = make_center_label(intro_text, int(CENTER_Y + 0.08*H),
                                    size=max(16, int(0.025*H)), width_frac=0.9, bold=False)

    btn_font = max(16, int(0.025*H))
    start_btn_label = pyglet.text.Label(
        "Begin test",
        font_size=btn_font,
        font_name=FONT_STACK,
        x=0, y=0, anchor_x='center', anchor_y='center',
        color=FG_WHITE, batch=batch_ui
    )
    padding_x = int(btn_font * 1.2)
    padding_y = int(btn_font * 0.8)
    text_w = start_btn_label.content_width
    text_h = start_btn_label.content_height
    btn_w = text_w + 2 * padding_x
    btn_h = text_h + 2 * padding_y
    sx = int(CENTER_X - btn_w/2)
    sy = int(CENTER_Y - 0.12*H - btn_h/2)
    start_btn_rect   = shapes.Rectangle(sx, sy, btn_w, btn_h, color=ACCENT_BLACK, batch=batch_ui)
    start_btn_border = shapes.BorderedRectangle(sx, sy, btn_w, btn_h, border=3,
                                                color=ACCENT_BLACK, border_color=ACCENT_WHITE, batch=batch_ui)
    start_btn_label.x = sx + btn_w/2
    start_btn_label.y = sy + btn_h/2

    rating_prompt = make_center_label(
        "How vivid was your image?\nClick a number to rate 1–5",
        int(CENTER_Y + 0.22*H),
        size=max(16, int(0.025*H)),
        bold=False
    )

    total_width = min(W * 0.88, 5 * min(W, H) * 0.22)
    left = CENTER_X - total_width/2
    step = total_width / 5
    col_width = step
    size = min(W, H) * 0.12
    cy = int(CENTER_Y - 0.02*H)
    label_y = int(cy - size*1.05)

    for i, b in enumerate(buttons):
        cx = int(left + (i + 0.5) * step)
        b.place(cx, cy, size, label_y, col_width)

def instruction_label_for(idx):
    header = header_for_idx(idx)
    question = VVIQ_QUESTIONS[idx]
    txt = f"{header}\n\nClose your eyes and imagine:\n\n{question}\n\n(Press SPACE when ready to rate)"
    return make_center_label(txt, CENTER_Y, size=max(16, int(0.024*H)), width_frac=0.85, bold=False)

state = "start"
q_index = 0
instr_label = None
data_rows = []
total_score = 0
prev_state = None

def build_confirm_overlay():
    cover = shapes.Rectangle(0, 0, W, H, color=(64, 64, 64), batch=batch_modal)
    cover.opacity = 180
    msg = pyglet.text.Label(
        "Do you really want to quit?\n\nPress Y to confirm, N to continue.",
        font_size=max(16, int(0.025 * H)),
        font_name=FONT_STACK,
        x=CENTER_X, y=CENTER_Y,
        anchor_x="center", anchor_y="center",
        color=FG_WHITE,
        width=int(W * 0.8), multiline=True, align="center",
        batch=batch_modal
    )
    return {"cover": cover, "msg": msg}

def set_question(i):
    global q_index, state, instr_label
    if i >= len(VVIQ_QUESTIONS):
        save_results(interrupted=False)
        state = "thanks"
        return
    q_index = i
    instr_label = instruction_label_for(i)
    state = "instruction"

@win.event
def on_draw():
    win.clear()
    if state == "start":
        start_title.draw(); start_intro.draw()
        start_btn_rect.draw(); start_btn_border.draw(); start_btn_label.draw()
    elif state == "instruction":
        if instr_label: instr_label.draw()
    elif state == "rating":
        rating_prompt.draw()
        for b in buttons:
            b.rect.draw(); b.number.draw(); b.label.draw()
    elif state == "thanks":
        make_center_label("Thank you!\nYou may now close this window.",
                          CENTER_Y, size=max(18, int(0.024*H))).draw()
    elif state == "confirm_exit":
        confirm_ui["cover"].draw()
        confirm_ui["msg"].draw()
    elif state == "waiting":
        rating_prompt.draw()
        for b in buttons:
            b.rect.draw(); b.number.draw(); b.label.draw()

@win.event
def on_mouse_press(x, y, button, modifiers):
    global state, total_score, _pending_next_index
    if state == "start":
        if start_btn_rect.x <= x <= start_btn_rect.x + start_btn_rect.width and \
           start_btn_rect.y <= y <= start_btn_rect.y + start_btn_rect.height:
            set_question(0)
    elif state == "rating":
        for b in buttons:
            if b.hit(x, y):
                data_rows.append({
                    # report results of each question
                    "test": "VVIQ",
                    "participant_id": PID,
                    "timestamp": TIMESTAMP,
                    "question_index": q_index + 1, # question number
                    "question": VVIQ_QUESTIONS[q_index],
                    "rating": b.value
                })
                total_score += b.value
                _pending_next_index = q_index + 1
                state = "waiting"
                pyglet.clock.schedule_once(lambda dt: set_question(_pending_next_index), POST_CLICK_DELAY)
                break

@win.event
def on_key_press(symbol, modifiers):
    global state, prev_state, confirm_ui
    if symbol == key.ESCAPE and state != "confirm_exit":
        prev_state = state
        state = "confirm_exit"
        confirm_ui = build_confirm_overlay()
        return event.EVENT_HANDLED

    if state == "confirm_exit":
        if symbol == key.Y:
            save_results(interrupted=True)
            pyglet.app.exit()
            return event.EVENT_HANDLED
        elif symbol in (key.N, key.ESCAPE):
            state = prev_state
            return event.EVENT_HANDLED
        return event.EVENT_HANDLED

    if state == "instruction" and symbol == key.SPACE:
        state = "rating"
        return event.EVENT_HANDLED

@win.event
def on_close():
    if state != "thanks":
        save_results(interrupted=True)

@win.event
def on_resize(width, height):
    layout()
    if state in ("instruction", "rating", "waiting"):
        globals()['instr_label'] = instruction_label_for(q_index)
    global confirm_ui
    confirm_ui = build_confirm_overlay()

def update(dt):
    global _esc_latched, state, prev_state, confirm_ui
    esc_down = bool(key_handler[key.ESCAPE])
    if state != "confirm_exit":
        if esc_down and not _esc_latched:
            _esc_latched = True
            prev_state = state
            state = "confirm_exit"
            confirm_ui = build_confirm_overlay()
    if not esc_down and _esc_latched and state != "confirm_exit":
        _esc_latched = False

layout()
confirm_ui = build_confirm_overlay()
instr_label = None
pyglet.clock.schedule_interval(update, 1/120)
pyglet.app.run()