# setting A tensors from the environment parameters
A = env.A
A_dependencies = env.A_dependencies # dependencies allow you to specify which state factors each observation modality depends on, so you dont have to store all the conditional dependencies between all state factors and each modality

# setting B tensors from the environment parameters
B = env.B
B_dependencies = env.B_dependencies

# creating C tensors filled with zeros for [location], [reward], [cue] based on A shapes
C = [jnp.zeros(a.shape[0], dtype=jnp.float32) for a in A] 
# setting preferences for outcomes only
C[1] = C[1].at[1].set(2.0)    # prefer reward
C[1] = C[1].at[2].set(-3.0)   # avoid punishment

# creating D tensors [location], [reward] based on B shapes
D = []
# D[0]: location - all zeros except location 0 (centre) because the agent always starts in the centre
D_loc = jnp.zeros(B[0].shape[0], dtype=jnp.float32) 
D_loc = D_loc.at[0].set(1.0)  # set centre location to 1.0
D.append(D_loc)

# D[1]: reward location - uniform distribution
D_reward = jnp.ones(B[1].shape[0], dtype=jnp.float32) 
D_reward = D_reward / jnp.sum(D_reward, axis=0, keepdims=True)  # normalise to get uniform distribution
D.append(D_reward)


# initialising the agent
agent = Agent(
    A, B, C, D, 
    policy_len=2, # how long the action sequence is that the agent is evaluating
    A_dependencies=A_dependencies, 
    B_dependencies=B_dependencies,
    learn_A=False,
    learn_B=False
)

# you may print the agent's generative model parameters to see the shapes of the tensors and the values by editing and uncommenting the following lines and running the code: 

print(f'Number of generated agents: {agent.batch_size}\n')
print([a.shape for a in agent.A]) # shape of all A tensors
print(agent.A[1][0][:,:,1]) # likelihood of observing no outcome, reward, or punishment (rows), in each location (columns), when the reward condition is 1 (right arm)
print(agent.C[1]) # preferences for outcomes

frames = []
for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
    # get observations for this timestep
    observations_t = [
        info["observation"][0][:, t, :],
        info["observation"][1][:, t, :],  
        info["observation"][2][:, t, :]   
    ]
       
    frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
    frame = np.asarray(frame, dtype=np.uint8)
    plt.close()  # close the figure to prevent memory leak
    frames.append(frame)

frames = np.array(frames, dtype=np.uint8)
mediapy.show_video(frames, fps=1)

# # uncomment the following lines to save the video as a gif
# os.makedirs("figures", exist_ok=True)
# pil_frames = [Image.fromarray(frame) for frame in frames]
# reward_location = "random" if reward_condition is None else ("left" if reward_condition == 0 else "right")
# filename = os.path.join("figures", f"tmaze_{batch_size}_{reward_location}.gif")
# pil_frames[0].save(
#     filename,
#     save_all=True,
#     append_images=pil_frames[1:],
#     duration=1000,  # 1000ms per frame
#     loop=0
# )

batch_size = 9 # number of agents to run in parallel

# setting A tensors from the environment parameters
A, B = env.A, env.B
A_dependencies, B_dependencies = env.A_dependencies, env.B_dependencies # dependencies allow you to specify the state factors the observation modality depends on so you dont have to compute the full tensor using all state factors

# now expand A and B to have a new batch_size-length leading dimension
A = [jnp.broadcast_to(a, (batch_size,) + a.shape) for a in A]
B = [jnp.broadcast_to(b, (batch_size,) + b.shape) for b in B]

# creating C tensors filled with zeros for [location], [reward], [cue] based on A shapes
C = [jnp.zeros((batch_size, a.shape[1]), dtype=jnp.float32) for a in A] 
# setting preferences for outcomes only
C[1] = C[1].at[:,1].set(2.0)    # prefer reward
C[1] = C[1].at[:,2].set(-3.0)   # avoid punishment


# creating D tensors [location], [reward] based on B shapes
D = []
# D[0]: location - all zeros except location 0 (centre) because the agent always starts in the centre
D_loc = jnp.zeros((batch_size, B[0].shape[1]), dtype=jnp.float32) 
D_loc = D_loc.at[0,0].set(1.0)  # set centre location to 1.0
D.append(D_loc)

# D[1]: reward location - uniform distribution
D_reward = jnp.ones((batch_size, B[1].shape[1]), dtype=jnp.float32) 
D_reward = D_reward / jnp.sum(D_reward, axis=1, keepdims=True)  # normalise to get uniform distribution
D.append(D_reward)


# initialising the agent
agent = Agent(
    A, B, C, D, 
    policy_len=2,
    A_dependencies=A_dependencies, 
    B_dependencies=B_dependencies,
    batch_size=batch_size, # note we have to pass in this batch_size parameter so the class knows to treat the first dimension as the batch size
    learn_A=False,
    learn_B=False
)

_, info = rollout(agent, env, num_timesteps=T, rng_key=key) # running the aif loop

# rendering the task to visualise the agent's behaviour 
frames = []
for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
    # get observations for this timestep
    observations_t = [
        info["observation"][0][:, t, :],
        info["observation"][1][:, t, :],  
        info["observation"][2][:, t, :]   
    ]
       
    frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
    frame = np.asarray(frame, dtype=np.uint8)
    plt.close()  # close the figure to prevent memory leak
    frames.append(frame)

frames = np.array(frames, dtype=np.uint8)
mediapy.show_video(frames, fps=1)

# THE GENERATIVE MODEL

# Get A, B, and dependencies from the environment
A, B = env.A, env.B
A_dependencies, B_dependencies = env.A_dependencies, env.B_dependencies # dependencies allow you to specify

# now expand A and B to have a new batch_size-length leading dimension
A = [jnp.broadcast_to(a, (batch_size,) + a.shape) for a in A]
B = [jnp.broadcast_to(b, (batch_size,) + b.shape) for b in B]

# creating C tensors filled with zeros for [location], [reward], [cue] based on A shapes
C = [jnp.zeros((batch_size, a.shape[1]), dtype=jnp.float32) for a in A] 
# setting preferences for outcomes only
C[1] = C[1].at[:,1].set(3.0)    # prefer reward
C[1] = C[1].at[:,2].set(-4.0)   # avoid punishment


# creating D tensors [location], [reward] based on B shapes
D = []
# D[0]: location - all zeros except location 0 (centre) because the agent always starts in the centre
D_loc = jnp.zeros((batch_size, B[0].shape[1]), dtype=jnp.float32) 
D_loc = D_loc.at[0,0].set(1.0)  # set centre location to 1.0
D.append(D_loc)

# D[1]: reward location - uniform distribution
D_reward = jnp.ones((batch_size, B[1].shape[1]), dtype=jnp.float32) 
D_reward = D_reward / jnp.sum(D_reward, axis=1, keepdims=True)  # normalise to get uniform distribution
D.append(D_reward)


# initialising the agent
agent = Agent(
    A, B, C, D, 
    policy_len=3, # how long the action sequence is that the agent is evaluating
    A_dependencies=A_dependencies, 
    B_dependencies=B_dependencies,
    batch_size=batch_size,
    learn_A=False,
    learn_B=False
)

key = jr.PRNGKey(0) # random key for the aif loop
T = 20 # number of timesteps to rollout the aif loop for
_, info = rollout(agent, env, num_timesteps=T, rng_key=key) # running the aif loop

frames = []
for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
    # get observations for this timestep
    observations_t = [
        info["observation"][0][:, t, :],
        info["observation"][1][:, t, :],  
        info["observation"][2][:, t, :]   
    ]
       
    frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
    frame = np.asarray(frame, dtype=np.uint8)
    plt.close()  # close the figure to prevent memory leak
    frames.append(frame)

frames = np.array(frames, dtype=np.uint8)
mediapy.show_video(frames, fps=1)

#  initializing the prior over the A tensors (pA) based on the environment parameters
num_obs = [a.shape[0] for a in env.A]  # number of observations for each modality
num_states = [b.shape[0] for b in env.B]  # number of states for each factor

# initializing pA, pB to be identical to the environmental parameters
pA, pB = [jnp.copy(a) for a in env.A], [jnp.copy(b) for b in env.B] # have to copy them here because we will be modifying pA below

# adding non-zero Dirichlet parameter values to pA flatten the expected value of the A tensor (make more uncertain) 
alpha_scale_pA = 0.3 
for m in [1, 2]:  # only modifying the outcome (m=1) and cue (m=2) observation likelihood mappings
    pA[m] = pA[m] + (alpha_scale_pA * jnp.ones_like(pA[m]))

expected_A = [dirichlet_expected_value(pa_m, event_dim=0) for pa_m in pA]# compute the expected value of the Dirichlet parameters in pA 
expected_B = [dirichlet_expected_value(pb_f, event_dim=0) for pb_f in pB]# compute the expected value of the Dirichlet parameters in pB 

# creating C tensors filled with zeros for [location], [reward], [cue] based on modality shapes
C = [jnp.zeros(no, dtype=jnp.float32) for no in num_obs] 
# setting preferences for outcomes only
C[1] = C[1].at[1].set(2.0)    # prefer reward
C[1] = C[1].at[2].set(-3.0)   # avoid punishment

# creating D tensors [location], [reward] based on B shapes
D = [jnp.zeros(ns, dtype=jnp.float32) for ns in num_states]  # creating D tensors filled with zeros for [location], [reward] based on B shapes
# D[0]: location - all zeros except location 0 (centre) because the agent always starts in the centre
D_loc = jnp.zeros(num_states[0], dtype=jnp.float32) 
D_loc = D_loc.at[0].set(1.0)  # set centre location to 1.0
D[0] = D_loc

# D[1]: reward location - uniform distribution
D_reward = jnp.ones(num_states[1], dtype=jnp.float32) 
D_reward = D_reward / jnp.sum(D_reward, axis=0, keepdims=True)  # normalise to get uniform distribution
D[1] = D_reward


# initialising the agent
# NOTE: We initialize the first value of A and B based on the expected values of the Dirichlet parameters in pA and pB, respectively.
agent = Agent(
    expected_A, expected_B, C, D, 
    pA=pA,
    pB=pB, 
    policy_len=5, # how long the action sequence is that the agent is evaluating
    A_dependencies=env.A_dependencies, 
    B_dependencies=env.B_dependencies,
    learn_A=True,
    learn_B=False,
    gamma=1.0,
    action_selection="stochastic"
)

print("the environment's A tensor")
print(env.A[1][:,:,1])
print()
print("the agent's A tensor at t=0")
print(agent.A[1][0][:,:,1])
print()
print(f"the agent's A tensor at t={T}")
print(info["A"][1][0,-1,:,:,1])

frames = []
for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
    # get observations for this timestep
    observations_t = [
        info["observation"][0][:, t, :],
        info["observation"][1][:, t, :],  
        info["observation"][2][:, t, :]   
    ]
       
    frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
    frame = np.asarray(frame, dtype=np.uint8)
    plt.close()  # close the figure to prevent memory leak
    frames.append(frame)

frames = np.array(frames, dtype=np.uint8)
mediapy.show_video(frames, fps=1)

# Random initialization test
key = jr.PRNGKey(24)

# setting A tensors from the environment parameters, but randomize the prior over the transitions
pA, pB = [jnp.copy(a) for a in env.A], [jr.uniform(subkey, shape=b.shape) for subkey, b in zip(jr.split(key, num=2), env.B)]

# adding noise to outcome and cue modalities to make priors more uncertain
key1, key2 = jr.split(key, num=2)
for(key_m, modality) in zip([key1, key2], [1, 2]):  # only modifying the outcome (m=1) and cue (m=2) observation likelihood mappings
    pA[modality] = jr.uniform(key_m, shape=pA[modality].shape) # generating random values between 0 and 1

expected_A = [dirichlet_expected_value(pa_m, event_dim=0) for pa_m in pA] # compute the expected value of the Dirichlet parameters in pA
expected_B = [dirichlet_expected_value(pb_f, event_dim=0) for pb_f in pB]# compute the expected value of the Dirichlet parameters in pB

# initialising the agent
agent = Agent(
    expected_A, expected_B, C, D, 
    pA=pA,
    pB=pB, 
    policy_len=5, # how long the action sequence is that the agent is evaluating
    A_dependencies=env.A_dependencies, 
    B_dependencies=env.B_dependencies,
    batch_size=batch_size, 
    learn_A=True,
    learn_B=True,
    gamma=0.1,
    action_selection="stochastic"
)

# running the active inference simulation
key = jr.PRNGKey(0) # random key for the aif loop
T = 50 # number of timesteps to rollout the aif loop for
_, info = rollout(agent, env, num_timesteps=T, rng_key=key) # running the aif loop

print("the environment's A tensor")
print(env.A[1][:,:,1])
print()
print("the agent's A tensor at t=0")
print(agent.A[1][0][:,:,1])
print()
print(f"the agent's A tensor at t={T}")
print(info["A"][1][0,-1,:,:,1])

print("the environment's B tensor")
print(env.B[1][:,:,1])
print()
print("the agent's B tensor at t=0")
print(agent.B[1][0][:,:,1])
print()
print(f"the agent's B tensor at t={T}")
print(info["B"][1][0,-1,:,:,1])

frames = []
for t in range(info["observation"][0].shape[1]):  # iterate over timesteps
    # get observations for this timestep
    observations_t = [
        info["observation"][0][:, t, :],
        info["observation"][1][:, t, :],  
        info["observation"][2][:, t, :]   
    ]
       
    frame = env.render(mode="rgb_array", observations=observations_t) # render the environment using the observations for this timestep
    frame = np.asarray(frame, dtype=np.uint8)
    plt.close()  # close the figure to prevent memory leak
    frames.append(frame)

frames = np.array(frames, dtype=np.uint8)
mediapy.show_video(frames, fps=1)