#!/usr/bin/env python3

import pyactr as actr
import prod_numbers

Model = actr.ACTRModel()
DM = Model.decmem
goal = Model.goal
imaginal = Model.set_goal(name="imaginal", delay=0.2)

_ = prod_numbers.number(Model)

def get_digit(number, n):
    return number // 10**n % 10

actr.chunktype("number", ("number", "next", "ones", "tens", "hundreds"))
actr.chunktype(
    "math_goal",
    (
        "task",
        "result1",
        "result2",
        "op",
        "nextop",
        "nextnextop",
        "arg1",
        "arg1_idx",
        "arg2",
        "arg2_idx",
        "result",
        "expand_slot",
        "hundreds1",
        "tens1",
        "ones1",
        "hundreds2",
        "tens2",
        "ones2",
    ),
)
actr.chunktype(
    "math_op",
    (
        "op",
        "arg1",
        "arg2",
        "result",
    ),
)

# 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),
        )
    )


Model.productionstring(
    name="continue_with_next_op",
    string="""
        =g>
        isa     math_goal
        op      done
        nextop  ~None
        nextop  =nextop
        nextnextop  =nextnextop
        ==>
        =g>
        isa     math_goal
        op      =nextop
        nextop  =nextnextop
        nextnextop  None
    """,
)

Model.productionstring(
    name="reset",
    string="""
    =g>
    isa         math_goal
    op          reset
    ==>
    =g>
    isa         math_goal
    op          expand
    nextop      reset
    """
)

goal.add(
    actr.makechunk(
        "",
        "math_goal",
        op="expand",
        nextop="reset",
        expand_slot=None,
        ones1=1,
        tens1=2,
        hundreds1=3,
        ones2=4,
        tens2=5,
        hundreds2=6,
    )
)
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

sim = Model.simulation(
    gui=False,
)
sim.run(max_time=2)