#!/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!

In diesem Experiment arbeiten Sie in einem Labor, welches die Wasserqualität analysiert.
Dafür bekommen Sie einige Wasserproben und müssen für jede Probe verschiedene Kennwerte ermitteln.

Jede Wasserprobe besitzt bereits verschiedene gemessene Werte, wie der Algengehalt, welche für die Berechnungen benutzt werden.


(Leertaste zum Fortfahren)
"""

intro2_text = """Im folgenden müssen sie verschiedene Rechenaufgaben lösen um die Kennwerte zu berechnen.
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 (z.B. Algen) werden oben am Bildschirm angezeigt.

Manche Variablen haben mehrere mögliche Werte; "Gifte_2" besagt z.B., dass der zweite Wert der Giftwerte zu verwenden ist.
"_max/_min" besagt, dass der maximale/minimale Wert dieser Variable zu verwenden ist.

Als letzte Berechnung müssen Sie den Gesamtwert der Wasserqualität aus ihren Ergebnissen berechnen.

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 = 75
TEST_TRIALS = 50
# TRAIN_TRIALS = 1
# TEST_TRIALS = 1
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=[1000000, 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("Mineralien", [water_sample.solid])
    algae = DisplayVariable("Algen", [water_sample.algae])
    lime = DisplayVariable("Sandstein", water_sample.lime)
    toxin = DisplayVariable("Gifte", 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 = -600
    proc_y = -100
    answ_x = 200
    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_equals = visual.TextBox2(
            WIN,
            "=",
            letterHeight=50,
            pos=(answ_x - 100, answ_y),
            size=[150, 70],
            alignment="center",
        )
        stims.append(stim_answer_equals)

        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 Wasserprobe fortzufahren""",
    letterHeight=50,
    alignment="center",
)

intro = visual.TextBox2(
    WIN, intro_text, letterHeight=40, alignment="center", size=(100000, 100000)
)
intro.draw()
WIN.flip()
event.waitKeys(keyList=["space"])

intro2 = visual.TextBox2(
    WIN, intro2_text, letterHeight=40, alignment="center", size=(100000, 100000)
)
intro2.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])

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


results = {}
for i in range(TRAIN_TRIALS):
    print(train_procedures)
    answer, procedure_keys = run_trial(
        all_samples[i], train_procedures, ORDER_CONDITIONS[0]
    )
    answer_dict = {}
    answer_dict["procedure_order"] = procedure_keys
    answer_dict["water_sample"] = all_samples[i].water_sample_dict()
    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"])


phase = visual.TextBox2(
    WIN,
    """Sie haben den ersten Teil geschafft! Der zweite Teil ist etwas kürzer als der erste

Drücken Sie die Leertaste um anzufangen.""",
    letterHeight=50,
    alignment="center",
)
phase.draw()
WIN.flip()
event.waitKeys(keyList=["space"])

train_procedures[2] = transfer_procedure
for i in range(TEST_TRIALS):
    print(train_procedures)
    answer, procedure_keys = run_trial(
        all_samples[TRAIN_TRIALS + i], train_procedures, ORDER_CONDITIONS[0]
    )
    answer_dict = {}
    answer_dict["procedure_order"] = procedure_keys
    answer_dict["water_sample"] = all_samples[TRAIN_TRIALS + i].water_sample_dict()
    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)

outro_text = "Das Experiment ist nun vorüber.\n\nVielen Dank für Ihre Teilnahme!"
outro = visual.TextBox2(
    WIN, outro_text, letterHeight=40, alignment="center", size=(100000, 100000)
)
outro.draw()
WIN.flip()
event.waitKeys(keyList=["space"])