#!/usr/bin/env python3

import pyactr as actr
from pprint import pprint

from model_init import init
import prod_addition
import prod_subtraction
import prod_comp
import prod_multi
import prod_numbers
import prod_procedure
import prod_motor
import prod_vis
import model_env


def add_goal(goal, op, arg1, arg2):

    goal.add(actr.makechunk("", "math_goal", op=op, task=op, arg1=arg1, arg2=arg2))


def wait_input():
    op = input("op\n")
    arg1 = input("arg1\n")
    arg2 = input("arg2\n")
    return op, int(arg1), int(arg2)


def add_proc(goal, proc):
    # input()
    goal.add(actr.makechunk("", "math_goal", proc=proc, ones_carry=""))


def start():
    condition = "fixed"
    training_N = 1
    test_N = 1

    stimuli = model_env.Stimuli(condition, training_N=training_N, test_N=test_N)

    pprint(stimuli.training_order)
    pprint(stimuli.test_order)

    stimuli.generate_stimuli()

    pprint(stimuli.training_stimuli)
    pprint(stimuli.test_stimuli)

    stimulus = stimuli.next_stimulus()

    # op, arg1, arg2 = wait_input()
    Model, DM, goal, imaginal, env = init()
    # Model.model_parameters["subsymbolic"] = True
    # Model.model_parameters["partial_matching"] = True
    # Model.model_parameters["activation_trace"] = True
    # Model.model_parameters["instantaneous_noise"] = .2
    # Model.model_parameters["utility_alpha"] = 0.5
    # Model.model_parameters["utility_noise"] = 0.5
    # Model.model_parameters["production_compilation"] = True
    # Model.model_parameters["utility_learning"] = True
    # add_goal(goal, op, arg1, arg2)
    add_proc(goal, "next_proc")

    prod_procedure.procedures(Model)
    number_prods = prod_numbers.number(Model)
    add_prods = prod_addition.addition(Model)
    sub_prods = prod_subtraction.subtraction(Model)
    greater_prods = prod_comp.greater_than(Model)
    less_prods = prod_comp.lesser_than(Model)
    multi_prods = prod_multi.multiplication(Model)
    motor_prods = prod_motor.procedures(Model)
    visual_prods = prod_vis.procedures(Model)

    for p in Model.productions:
        Model.productions[p].__setitem__("utility", 1)

    print("goal: ", goal)
    # print("imaginal: ", imaginal)
    sim = Model.simulation(
        gui=False,
        environment_process=env.environment_process,
        stimuli=stimulus,
        realtime=False,
        # triggers="space",
        triggers="b",
    )
    # pprint(sim._Simulation__pr.rules)
    # pprint(sim._Simulation__pr.ordered_rulenames)
    i = 1
    j = 1
    phase = "training"
    userinput = {"training": {}, "test": {}}
    while True:
        if j > training_N:
            phase = "test"
        sim.step()
        # print(goal)
        if sim.current_event.time >= 900:
            pprint(sim.current_event)
            break
        if "KEY PRESSED" in sim.current_event.action:
            if not stimuli.current_stimulus_id in userinput[phase]:
                userinput[phase][stimuli.current_stimulus_id] = {}
            if not str(i) in userinput[phase][stimuli.current_stimulus_id]:
                userinput[phase][stimuli.current_stimulus_id][str(i)] = []
            userinput[phase][stimuli.current_stimulus_id][str(i)].append(
                sim.current_event.action.split(":")[1].strip()
            )
        if (
            "RULE FIRED:" in sim.current_event.action
            and "number_expand_done and continue_with_next_op"
            in sim.current_event.action
        ):
            print("FIRED COMPILED RULE!!!!")
            # print(Model.productions[sim.current_event.action[12:]])
            # break
        if sim.current_event.action == "NO RULE FOUND":
            print(goal)
        if sim.current_event.action == "KEY PRESSED: SPACE":
            sim._Simulation__env.stimulus = stimuli.update_current_stimulus(
                f"Answer{i}",
                int(
                    "".join(userinput[phase][stimuli.current_stimulus_id][str(i)][:-1])
                ),
            )
            i += 1
            pprint(userinput)
            pprint("NEW PROC")
            if condition != "blocked":
                if i <= 6:
                    goal.add(
                        actr.makechunk("", "math_goal", proc="next_proc", ones_carry="")
                    )
                elif i >= 7 and j < training_N + test_N:
                    print("Stimulus Done, next stimulus")
                    sim._Simulation__env.stimulus = stimuli.next_stimulus()
                    goal.add(
                        actr.makechunk("", "math_goal", proc="next_proc", ones_carry="")
                    )
                    i = 1
                    j += 1
                elif i >= 7 and j >= training_N + test_N:
                    print("DONE")
                    break
            elif condition == "blocked":
                if i >= 6 * (training_N + test_N):
                    print("DONE")
                    break
                sim._Simulation__env.stimulus = stimuli.next_stimulus()

    # sim.run(max_time=25)
    print("training order: ", stimuli.training_order_list)
    print("test order: ", stimuli.test_order_list)
    pprint(userinput)
    print("Simulation time: ", sim.show_time())
    print("goal: ", goal)
    print(sim.current_event)
    # pprint(vars(sim))
    pprint(vars(sim._Simulation__env))
    # print(sim.__env)
    # sim.__printenv__()
    # print(envs.stimulus)
    # print("imaginal: ", imaginal)
    math_goals = [sim for sim in list(DM) if sim.typename == "procedure"]
    # print(math_goals)
    # pprint(Model.productions)

    # for p in Model.productions:
    #     print(p)
    #     print(Model.productions[p])#.__setitem__("utility", 1)
    #     print("\n")
    # print("#######used prods########")
    # print(sim.ordered_rulenames)

    # pprint(sim._Simulation__pr.rules)
    # pprint(sim._Simulation__pr.ordered_rulenames)


if __name__ == "__main__":
    start()