using MuZero

function MuZero.getconfig(::Type{Tac})
    config = make_board_game_config(
        Tac, 
        network_type=Incremental,
        max_moves=9, 
        root_dirichlet_alpha=0.3, 
        action_space_size=9, # Number of all possible actions in the game 
        support_size = 10, # Value and reward are scaled (with almost sqrt) and encoded on a vector with a range of -support_size to support_size. Choose it so that support_size <= sqrt(max(abs(discounted reward)))
        observation_shape = (3, 3, 3),
        stacked_observations = 0, # Number of previous observations and previous actions to add to the current observation
        hidden_state_size = 32, # The 1-dimensional size of the encoded game state
        checkpoint_interval = 1000, #  Number of training steps before using test model is updated.
        buffer_capacity = 100,
        num_actors = 2,
        weight_decay = 1.0f-1,
        lr_init=8e-4,
        save_interval = 1000, # Number of training steps before using play model is updated.
        num_unroll_steps = 2,
    )
    config.max_games = 55_000
    config.batch_size = 16
    config.num_simulations = 100
    config.training_steps = 200001
    config
end

function tacschedule!(l::AbstractLearner)
    step = get_steps(l)
    opt = l.optimizer
    setonce = (at, value) -> begin
        if step > at && opt[3].eta > value
            opt[3].eta = value
            println("Change learning rate to: $value")
        end
    end
    setonce(0, 8e-4)
    setonce(10000, 1e-4)
    setonce(20000, 9e-5)
    setonce(23000, 4e-5)
    setonce(25000, 3e-5)
    setonce(33000, 1e-5)
    setonce(140000, 1e-5)
end