#!/usr/bin/env python3

import pyactr as actr

# https://stackoverflow.com/a/39644726
def get_digit(number, n):
    return number // 10**n % 10


def init():
    env = actr.Environment()
    Model = actr.ACTRModel(environment=env)

    DM = Model.decmem
    Model.visualBuffer("visual", "visual_location", DM, finst=4)
    goal = Model.goal
    imaginal = Model.set_goal(name="imaginal", delay=0.2)

    actr.chunktype("number", ("number", "next", "ones", "tens", "hundreds"))
    actr.chunktype(
        "procedure",
        (
            "proc",
            "op1",
            "arg1_1",
            "arg1_1_idx"
            "arg1_2",
            "arg1_2_idx"
            "op2",
            "arg2_1",
            "arg2_1_idx"
            "arg2_2",
            "arg2_2_idx"
        ),
    )
    actr.chunktype(
        "math_goal",
        (
            "task",
            "result1",
            "result2",
            "op",
            "nextop",
            "arg1",
            "arg1_idx"
            "arg2",
            "arg2_idx"
            "result",
            "expand_slot",
            "hundreds1",
            "tens1",
            "ones1",
            "hundreds2",
            "tens2",
            "ones2",
            "hundreds_ans",
            "tens_ans",
            "ones_ans",
            "mul_counter",
            "mem_arg1",
            "mem_arg2",
            "ones_carry",
            "tens_carry",
            "vis_counter",
        ),
    )
    actr.chunktype(
        "math_op",
        (
            "op",
            "arg1",
            "arg2",
            "result",
        ),
    )

    # Add procedures
    DM.add(
        actr.makechunk(
            "procedure1",
            "procedure",
            proc="proc1",
            op1="sub",
            arg1_1="Sandstein",
            arg1_1_idx=4,
            arg1_2="Sandstein",
            arg1_2_idx=2,
            op2="mul",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2="Mineralien",
            arg2_2_idx=1,
        )
    )

    DM.add(
        actr.makechunk(
            "procedure2",
            "procedure",
            proc="proc2",
            op1="mul",
            arg1_1=2,
            arg1_1_idx=0,
            arg1_2="Algen",
            arg1_2_idx=1,
            op2="add",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2="Sandstein",
            arg2_2_idx="min",
        )
    )
    DM.add(
        actr.makechunk(
            "procedure3",
            "procedure",
            proc="proc3",
            op1="add",
            arg1_1="Gifte",
            arg1_1_idx="max",
            arg1_2="Gifte",
            arg1_2_idx="min",
            op2="add",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2=0,
            arg2_2_idx=0,
        )
    )
    DM.add(
        actr.makechunk(
            "procedure4",
            "procedure",
            proc="proc4",
            op1="mul",
            arg1_1="Mineralien",
            arg1_1_idx=1,
            arg1_2=2,
            arg1_2_idx=0,
            op2="sub",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2="Gifte",
            arg2_2_idx="4",
        )
    )
    DM.add(
        actr.makechunk(
            "procedure5",
            "procedure",
            proc="proc5",
            op1="sub",
            arg1_1="Gifte",
            arg1_1_idx=3,
            arg1_2="Gifte",
            arg1_2_idx=2,
            op2="greater",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2="Sandstein",
            arg2_2_idx=3,
        )
    )
    DM.add(
        actr.makechunk(
            "procedure6",
            "procedure",
            proc="proc6",
            op1="add",
            arg1_1="Sandstein",
            arg1_1_idx=1,
            arg1_2="Gifte",
            arg1_2_idx=1,
            op2="lesser",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2="Algen",
            arg2_2_idx=1,
        )
    )
    DM.add(
        actr.makechunk(
            "procedure_overall",
            "procedure",
            proc="proc_overall",
            op1="sub",
            arg1_1=100,
            arg1_1_idx=0,
            arg1_2="results",
            arg1_2_idx="max",
            op2="add",
            arg2_1="result1",
            arg2_1_idx=0,
            arg2_2=0,
            arg2_2_idx=0,
        )
    )

    # Add numbers 0-999 to decmem
    for i in range(0, 1000):
        DM.add(
            actr.makechunk(
                f"number{str(i)}",
                "number",
                number=i,
                next=i + 1,
                ones=get_digit(i, 0),
                tens=get_digit(i, 1),
                hundreds=get_digit(i, 2),
            )
        )

    # Add comparison relations to single digit numbers
    for i in range(0, 11):
        for j in range(0, 11):
            DM.add(
                actr.makechunk(
                    f"greater{i}{j}",
                    "math_op",
                    op="greater",
                    arg1=i,
                    arg2=j,
                    result=max(i, j),
                )
            )
            DM.add(
                actr.makechunk(
                    f"lesser{i}{j}",
                    "math_op",
                    op="lesser",
                    arg1=i,
                    arg2=j,
                    result=min(i, j),
                )
            )
    for i in range(0, 21):
        for j in range(0, 21):
            DM.add(
                actr.makechunk(
                    f"plus{i}{j}", "math_op", op="add", arg1=i, arg2=j, result=i + j
                )
            )
    return Model, DM, goal, imaginal, env