#!/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
from pprint import pprint
import pandas as pd
import pickle


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

intro_text = """Vielen Dank dass Sie bei unserem Experiment zum menschlichen Lernen teilnehmen!

Im folgenden müssen sie verschiedene Rechenaufgaben lösen.
Verwenden Sie die gewohnten Rechenregeln und geben sie Ihre Lösung bitte immer als zweistellige Zahl ein, und bestätigen mit Enter.
Die Werte der Variablen werden oben am Bildschirm angezeigt. Manche Variablen haben mehrere mögliche Werte; "_2" besagt z.B. das der zweite Wert zu verwenden ist.

Nach jeder Aufgabe können Sie kurz pausieren.

Drücken Sie die Leertaste um zu beginnen"""


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


WIN = visual.Window((2560, 1440), fullscr=True, units="pix")
MONITOR_FPS = 60
# TRAIN_TRIALS = 1
# TEST_TRIALS = 1
TRAIN_TRIALS = 75
TEST_TRIALS = 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 generate_all_watersamples(n):
    samples = []
    for _ in range(n):
        samples.append(frensch_procedures.constrained_WaterSample())
    return samples


def run_trial(water_sample, procedure_keys: list, condition):
    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
            if "\n" in answer:
                if answer[0].isdigit() and answer[1].isdigit():
                    not_finished = False
                else:
                    stim_answer_box.text = answer[:-1]
            if len(answer) > 2:
                stim_answer_box.text = stim_answer_box.text[:2]
        answer_time = core.monotonicClock.getTime() - start_time
        answers.append((answer.replace("\n", ""), answer_time))

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

    return answers, procedure_keys


pause = visual.TextBox2(
    WIN,
    """Drücken Sie die Leertaste um mit der nächsten Aufgabe fortzufahren""",
    letterHeight=50,
    alignment="center",
)

intro = visual.TextBox2(WIN, intro_text, letterHeight=30, alignment="center")
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("test/train")
print(train_procedures)
test_procedures = train_procedures
test_procedures[2] = transfer_procedure
print(test_procedures)

all_samples = generate_all_watersamples(TRAIN_TRIALS + TEST_TRIALS)
pprint(all_samples)


results = {}
for i in range(TRAIN_TRIALS):
    answer, procedure_keys = run_trial(
        all_samples[i], train_procedures, ORDER_CONDITIONS[0]
    )
    answer_dict = {}
    answer_dict["procedure_order"] = procedure_keys
    for proc, key in zip(answer, procedure_keys):
        answer_dict[key] = {"answer": proc[0], "time": proc[1]}

    results[f"train_{i}"] = answer_dict

    pause.draw()
    WIN.flip()
    event.waitKeys(keyList=["space"])


train_procedures[2] = transfer_procedure
for i in range(TEST_TRIALS):
    answer, procedure_keys = run_trial(
        all_samples[TRAIN_TRIALS + i], train_procedures, ORDER_CONDITIONS[0]
    )
    answer_dict = {}
    answer_dict["procedure_order"] = procedure_keys
    for proc, key in zip(answer, procedure_keys):
        answer_dict[key] = {"answer": proc[0], "time": proc[1]}

    results[f"test_{i}"] = answer_dict

    pause.draw()
    WIN.flip()
    event.waitKeys(keyList=["space"])

pprint(results)

df = pd.DataFrame.from_dict(results, orient="index")
df.to_csv("vp_results.csv")

with open("vp.pkl", "wb") as file:
    pickle.dump(results, file)