# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Regarding this Notebook
# This is a replication of the original analysis performed in the paper by [<NAME> 2020](missing). This replication script will not be updated as it is intended for reproducibility. Any deviations from the paper is marked with bold for transparency.
# Footnotes and internal documentation references are removed from this example to avoid confusion.
#
# ---
#

# # 2.2 Using tomsup
# One of the advantages of computational models of cognitive processes is that the implications of the model can be worked out by simulating the model’s behavior in a variety of situations. tomsup in particular, allows to test the k-ToM model as it plays a wide set of game-theoretical situations (e.g.  Matching Pennies or Prisoner’s Dilemma), in interaction with a variety of different agents (e.g.  other k-ToM or less sophisticated agents), within different possible settings (e.g. repeated interactions with the same opponent, or round robin tournaments). In order to better understand the setup of the tomsup package, we start with the case of two simple agents interacting, followed by a simple exampleusing k-ToM agents, which will also illustrate how one might implement tomsup in an experiment. Lastly, we will show how to run a simulation using multiple agents as well as how to plot the evolving internal states of a k-ToM agent. In this simple scenario two agents are playing the Matching Pennies game. One agent hides a penny in one hand: let’s say chooses 0 for hiding in the left hand, and 1 in the right. The other agent has to guess where the penny is. If the second agent guesses (chooses the same hand as the first), it wins and the first loses. In other words, the first agent wants to choose the hand that the second will not choose and the second wants to choose the hand that the first chooses. In this example, one of the agents implements the Random Bias strategy (e.g. has a 60 percent probability of choosing right over left), while the other implements a classic Q-learning strategy (a model free reinforcement learning mechanism updating the expected reward of choosing a specific option on a trial by trial basis). The full list of strategies already implemented in tomsup is accessible using the function `valid_agents()`. The user first has to install the tomsup package developed using python 3.6 (Van Rossum & Drake, 2009). The package can be downloaded and installed using pip:
#
# ```pip3  install  tomsup```

# **However, in this notebook we will assume the user simply downloaded the git. Feel free to skip the next code chunk if that is not the case.**

# +
# assuming you are in the github folder change the path - not relevant if tomsup is installed via. pip
import os

os.chdir("..")  # go out of the tutorials folder
# -

# Both approaches will also install the required dependencies. Now tomsup can be imported into Python following the lines;

import tomsup as ts

# We will also set a arbitrary seed for to ensure reproducibility;

# +
import random
import numpy as np

np.random.seed(1995)
random.seed(1995)  # The year of birth of the first author
# -

# First we need to set up the Matching Pennies game. As different games are defined by different payoff matrices, we set up the game by creating the appropriate payoff matrix using the ```PayoffMatrix``` class. 

# +
# initiate the competitive matching pennies game
penny = ts.PayoffMatrix(name="penny_competitive")

# print the payoff matrix
print(penny)
# -

#
# The Matching Pennies game is a zero sum game, meaning that for one agent to get a reward, the opponent has to lose. Agents have thus to predict their opponents' behavior, which is ideal for investigating \gls{tom}. Note that to explore other payoff matrices included in the package, or to learn how to specify a custom payoff matrix, the user can type the `help(ts.PayoffMatrix)` command.
#
# Then we create the first of the two competing agents:
#

# +
# define the random bias agent, which chooses 1 70 percent of the time, and call the agent "jung"
jung = ts.RB(bias=0.7)

# Examine Agent
print(f"jung is a class of type: {type(jung)}")
if isinstance(jung, ts.Agent):
    print(f"but jung is also an instance of the parent class ts.Agent")

# let us have Jung make a choice
choice = jung.compete()

print(f"jung chose {choice} and its probability for choosing 1 was {jung.get_bias()}.")
# -

#
# Note that it is possible to create one or more agents simultaneously using the convenient `create\_agents()` and passing any starting parameters to it in the form of a dictionary. 

# create a reinforcement learning agent
skinner = ts.create_agents(agents="QL", start_params={"save_history": True})

#
# Now that both agents are created, we have them play against each other. 
#

# +
# have the agents compete for 30 rounds
results = ts.compete(jung, skinner, p_matrix=penny, n_rounds=30)

# examine results
print(results.head())  # inspect the first 5 rows of the dataframe
# -

# ** Note: you can remove the print() to get a nicer printout of the dataframe **

results.head()  # inspect the first 5 rows of the dataframe

#
# The data frame stores the choice of each agent as well as their resulting payoff. Simply summing the payoff columns would determine the winner.
#
#
# ## k-ToM
# Here we will present some simple examples of the k-ToM agent. For a more in-depth description we recommend checking the expanded introduction on the [Github repository](https://github.com/KennethEnevoldsen/tomsup/blob/master/tutorials/introduction_to_tom.ipynb).
#
# We will start of by creating a 0-ToM with default priors and `save_history=True` to examine the workings of it. Notice that setting `save_history` is turned off by default to save on memory which is especially problematic for ToM agents with high sophistication level.

# +
# Creating a simple 1-ToM with default parameters
tom_1 = ts.TOM(level=1, dilution=None, save_history=True)

# Extract the parameters
tom_1.print_parameters()
# -

#
# Note that k-ToM agents as default uses agnostic starting beliefs. These can be shown in detail and specified as desired, as shown in **appendix in the paper**.
#
# To increase the agent's tendency to choose one we could simply increase its bias. Similarly, if we want the agent to behave in a more more deterministic fashion we can decrease the behavioural temperature. When the parameter values are set, we can play the agent against an opponent using the `.compete()` method. Where `agent` denote the agent in the payoff matrix (0 or 1) and the `op_choice` denote the choice of the opponent during the previous round. 

tom_2 = ts.TOM(
    level=2,
    volatility=-2,
    b_temp=-2,  # more deterministic
    bias=0,
    dilution=None,
    save_history=True,
)
choice = tom_2.compete(p_matrix=penny, agent=0, op_choice=None)
print("tom_2 chose:", choice)

# The user is recommended to have the 1-ToM and the 2-ToM agents compete using the previously presented `ts.compete()` function for simplicity. However, to make the process more transparent for the user in the following we create a simple for-loop:
#

# +
tom_2.reset()  # reset before start


prev_choice_1tom = None
prev_choice_2tom = None
for trial in range(1, 4):
    # note that op_choice is choice on previous turn
    # and that agent is the agent you respond to in the payoff matrix
    choice_1 = tom_1.compete(p_matrix=penny, agent=0, op_choice=prev_choice_1tom)
    choice_2 = tom_2.compete(p_matrix=penny, agent=1, op_choice=prev_choice_2tom)

    # update previous choice
    prev_choice_1tom = choice_1
    prev_choice_2tom = choice_2

    print(
        f"Round {trial}",
        f"  1-ToM choose {choice_1}",
        f"  2-ToM choose {choice_2}",
        sep="\n",
    )
# -

# A for loop like this can be used to implement k-ToM in an experimental setting by replacing the agent with the behavior of a participant. Examples of such implementations (interfacing with PsychoPy are available in the [documentation](https://github.com/KennethEnevoldsen/tomsup/tree/master/tutorials/psychopy_experiment)). 
#

tom_2.print_internal(
    keys=["p_k", "p_op"], level=[0, 1]  # print these two states
)  # for the agent simulated opponents 0-ToM and 1-ToM

#
# For instance, we can note that the estimate of the opponent's sophistication level (\texttt{p\_k}) slightly favors a 1-ToM as opposed to a 0-ToM and that the average probability of the opponent choosing one (`p_op`) slightly favors 1 (which was indeed the option the opponent chose). These estimates are quite uncertain due to the few rounds played. More information on how to interpret the internal states of the ToM agent is available in the documentation of the package, e.g. by using the help function `help(tom_2.print_internal)`
#
# ## Multiple Agents and Visualizing Results
# The above syntax is useful for small setups. However, the user might want to build larger simulations involving several agents to simulate data for experimental setup or test underlying assumptions. The package provides syntax for quickly iterating over multiple agents, rounds and even simulations. We will here show a quick example along with how to visualize the results and internal states of ToM agents.

# +
# Create a list of agents
agents = ["RB", "QL", "WSLS", "1-TOM", "2-TOM"]
# And set their starting parameters. An empty dictionary denotes default values
start_params = [{"bias": 0.7}, {"learning_rate": 0.5}, {}, {}, {}]

group = ts.create_agents(agents, start_params)  # create a group of agents

# Specify the environment
# round_robin e.g. each agent will play against all other agents
group.set_env(env="round_robin")

# Finally, we make the group compete 20 simulations of 30 rounds
results = group.compete(p_matrix=penny, n_rounds=30, n_sim=20, save_history=True)
# -

# Following the simulation, a data frame can be extracted as before, with additional columns reporting simulation number, competing agent pair (`agent0` and `agent1`) and if `save_history=True` it will also add two columns denoting the internal states of each agent, e.g. estimates and expectations at each trial.
#

res = group.get_results()
print(res.head(1))  # print the first row

# **Again, removing the print statement gives you a more readable output**

res.head(1)

# ** to allow other authors to examine these results we have also saved the results to a new lines delimited .ndjson**

res.to_json("tutorials/paper.ndjson", orient="records", lines=True)

# The package also provides convenient functions for plotting the agent's choices and performance.
#

# > for nicer plots we will increase the figure size using the following code. This is excluded from the paper for simplicity

# +
import matplotlib.pyplot as plt

# Set figure size
plt.rcParams["figure.figsize"] = [10, 10]

# +
# plot a heatmap of the rewards for all agent in the tournament
group.plot_heatmap(cmap="RdBu_r")

plt.rcParams["figure.figsize"] = [5, 5]

# plot the choices of the 1-ToM agent when competing against the WSLS agent
group.plot_choice(agent0="WSLS", agent1="1-TOM", agent=1)
# plot the choices of the 1-ToM agent when competing against the WSLS agent
group.plot_choice(agent0="RB", agent1="1-TOM", agent=1)

# plot the score of the 1-ToM agent when competing against the WSLS agent
group.plot_score(agent0="WSLS", agent1="1-TOM", agent=1)
# plot the score of the 2-ToM agent when competing against the WSLS agent
group.plot_score(agent0="WSLS", agent1="2-TOM", agent=1)
# -

# As seen in the heatmap we see that the k-ToM model compares favorably against simpler agents
# such as the QL. Furthermore notice that the 1-ToM and 2-ToM compares especially favorably
# against the WSLS agent as this agent act as a deterministic 0-ToM. Similarly, we see that the
# 2-ToM agent incurs a cost for being more complex by being less able to take advantage of the
# deterministic nature of WSLS. We can examine this further in the figures, where we see that the
# 1-ToM is almost perfectly able to predict the behaviour of the WSLS agent after a turn 5
# across simulations while the 2-ToM, take longer to estimate the behaviour. The figures also show
# that 1-ToM differs in behavioural patterns figures when playing against a RB agents showing
# a bias estimation behaviour, while when playing against the WSLS it shows a oscillating
# choice pattern. Ultimately these are meant for initial investigation and more elaborate plots
# can be constructed from the results data frame.
#
# > here we just refer to the figures, for more exact references please see the paper

# Besides these general plots the package also contains a series of shortcuts for plotting $k$-ToM's internal states such as its estimate of its opponent's sophistication level, in which it is seen that the 2-ToM correctly estimates the opponents estimates as having a sophistication level of 1 on average.
#

# plot 2-ToM estimate of its opponent sophistication level
group.plot_p_k(agent0="1-TOM", agent1="2-TOM", agent=1, level=0)
group.plot_p_k(agent0="1-TOM", agent1="2-TOM", agent=1, level=1)

#
# It is also easy to plot k-ToM's estimates of its opponent's model parameters. As an example, the following code plots the 2-ToM's estimate of 1-ToM's volatility and bias. We see that the ToM agent approaches a correct estimate of the default volatility of -2 as well as correctly estimated its opponent as having no inherent bias.
#

# +
# plot 2-ToM estimate of its opponent's volatility while believing the opponent to be level 1.
group.plot_tom_op_estimate(
    agent0="1-TOM", agent1="2-TOM", agent=1, estimate="volatility", level=1, plot="mean"
)

# plot 2-ToM estimate of its opponent's bias while believing the opponent to be level 1.
group.plot_tom_op_estimate(
    agent0="1-TOM", agent1="2-TOM", agent=1, estimate="bias", level=1, plot="mean"
)
# -

#
# Use `help(ts.AgentGroup.plot_tom_op_estimate)` for information on how to plot the other estimated parameters or k-ToM's uncertainty in these parameters. 
# Additional information can be found in the history column in the results data frame, if needed. This includes all k-ToM's internal states (the changing variables in the model) which for example include choice probability, gradient, estimate uncertainties as well as k-ToM's estimates of its opponent's internal states. Documentation, examples and further tutorials can be found on the Github repository, this also includes a more in-depth description of the dynamics of **the k-ToM model implementation**.

# ---
# ## Are you left with any questions?
# Feel free to open a github issue with questions and or bug reports.
#
# Best, 
#
# *<NAME>*

StarcoderdataJupyterScriptsDedupFiltered