#!/usr/bin/env python3
from __future__ import annotations

from psychopy import constants, core, event, gui, visual
from collections import namedtuple
import frensch_procedures
import random


DisplayVariable = namedtuple("DisplayVariable", ["name", "values"])
DisplayProcedure = namedtuple("DisplayProcedure", ["procedure", "solution"])


def experiment_shutdown():
    WIN.close()
    core.quit()


WIN = visual.Window((2560, 1440), fullscr=True, units="pix")
MONITOR_FPS = 60
TRAIN_TRIALS = 1  # 75
TEST_TRIALS = 1  # 50
ORDER_CONDITIONS = ["fixed", "random", "blocked"]
PROCEDURE_KEYS = ["1", "2", "3", "4", "5", "6", "overall"]

# Cancel experiment anytime with Esc
event.globalKeys.add(key="escape", func=experiment_shutdown, name="shutdown")


def generate_variable_display(varx: list[DisplayVariable], x_positions: list[int]):
    assert len(varx) == len(x_positions)

    stims = []

    def gen_value_stims(values, x, y, offset):
        for value in values:
            y -= offset
            value_stim = visual.TextBox2(
                WIN,
                pos=(x, y),
                text=value,
                # size=200,
                letterHeight=100,
                alignment="center",
            )
            stims.append(value_stim)

    y = 650
    offset = 120

    for var, x_pos in zip(varx, x_positions):
        stim_var = visual.TextBox2(
            WIN,
            pos=[x_pos, y],
            text=var.name,
            # size=[1000, 1000],
            letterHeight=50,
            alignment="center",
        )
        stims.append(stim_var)

        gen_value_stims(var.values, x_pos, y, offset)

    return stims


def generate_procedure_display(procedure: DisplayProcedure, position):
    stim_procedure = visual.TextBox2(
        WIN,
        pos=position,
        text=procedure.procedure,
        # size=[1000, 1000],
        letterHeight=50,
        alignment="center",
    )
    return stim_procedure


def run_trial(procedure_keys: list, condition):
    water_sample = frensch_procedures.constrained_WaterSample()
    water_sample.print_all()

    if condition == "random":
        random.shuffle(procedure_keys)

    solid = DisplayVariable("SOLID", [water_sample.solid])
    algae = DisplayVariable("ALGAE", [water_sample.algae])
    lime = DisplayVariable("LIME", water_sample.lime)
    toxin = DisplayVariable("TOXIN", water_sample.toxin)
    x_positions = [-800, -400, 400, 800]

    stims = generate_variable_display([solid, algae, lime, toxin], x_positions)

    procedures = water_sample.procedure_dict()

    answers = []
    proc_x = -800
    proc_y = -100
    answ_x = 0
    answ_y = -100
    y_offset = 80
    for proc in procedure_keys:
        p = DisplayProcedure(procedures[proc][1](), procedures[proc][0]())
        p = generate_procedure_display(p, (proc_x, proc_y))
        stims.append(p)
        proc_y -= y_offset

        stim_answer_box = visual.TextBox2(
            WIN,
            "",
            letterHeight=50,
            pos=(answ_x, answ_y),
            size=[150, 70],
            editable=True,
            fillColor="white",
            color="black",
            alignment="center",
        )
        stims.append(stim_answer_box)
        answ_y -= y_offset

        not_finished = True
        answer = "not answered"
        start_time = core.monotonicClock.getTime()
        while not_finished:
            stim_answer_box.hasFocus = True
            for stim in stims:
                stim.draw()
            WIN.flip()
            answer = stim_answer_box.text
            # TODO: prevent nondigits
            # answer.replace("\n", "")
            # if not answer.isdigit():
            #     stim_answer_box.text = ""
            #     answer = ""
            if "\n" in stim_answer_box.text:
                not_finished = False
        answer_time = core.monotonicClock.getTime() - start_time
        print(f"time to answer: {answer_time}")
        answers.append(answer.replace("\n", ""))

    event.waitKeys(keyList=["space"])

    return answers


intro = visual.TextBox2(WIN, "Introduction")
intro.draw()
WIN.flip()
event.waitKeys(keyList=["space"])

train_procedures = PROCEDURE_KEYS[:-1]
random.shuffle(train_procedures)
transfer_procedure = train_procedures[-1]
train_procedures = train_procedures[:-1]
train_procedures.append(PROCEDURE_KEYS[-1])
print(train_procedures)


for _ in range(TRAIN_TRIALS):
    print(f"Answer: {run_trial(train_procedures, ORDER_CONDITIONS[0])}")

train_procedures[2] = transfer_procedure
print(train_procedures)
for _ in range(TEST_TRIALS):
    print(f"Answer: {run_trial(train_procedures, ORDER_CONDITIONS[0])}")