# 最初の行は1度だけ実行
!pip install -q --upgrade inferactively-pymdp==1.0.4 jax==0.6.2 jaxlib==0.6.2 equinox==0.13.8 numpyro==0.19.0

from importlib.metadata import version
import equinox as eqx
import jax
import jax.numpy as jnp
from jax import random as jr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import numpyro
import numpyro.distributions as dist
from numpyro.diagnostics import effective_sample_size, gelman_rubin
from numpyro.infer import MCMC, NUTS

from pymdp.agent import Agent

SEED = 7
plt.style.use("seaborn-v0_8-whitegrid")
print("pymdp:", version("inferactively-pymdp"))
print("jax:", jax.__version__, "numpyro:", numpyro.__version__)

MOTION_LEFT, MOTION_RIGHT = 0, 1
RESPONSE_NONE, RESPONSE_LEFT, RESPONSE_RIGHT = 0, 1, 2
EVIDENCE_LEFT, EVIDENCE_RIGHT = 0, 1
NEUTRAL, ERROR, CORRECT = 0, 1, 2

A_evidence = jnp.array([[0.8, 0.2], [0.2, 0.8]])
A_outcome = np.zeros((3, 2, 3))
A_outcome[NEUTRAL, :, RESPONSE_NONE] = 1.0
for motion in (MOTION_LEFT, MOTION_RIGHT):
    for response in (RESPONSE_LEFT, RESPONSE_RIGHT):
        result = CORRECT if response - 1 == motion else ERROR
        A_outcome[result, motion, response] = 1.0

B_motion = jnp.eye(2)[:, :, None]
B_response = np.zeros((3, 3, 2))
B_response[RESPONSE_LEFT, :, 0] = 1.0
B_response[RESPONSE_RIGHT, :, 1] = 1.0

# A：evidence←motion、outcome←motion＋response
A_dependencies = [[0], [0, 1]]
# B：次のmotion←現在のmotion、次のresponse←現在のresponse
B_dependencies = [[0], [1]]

agent = Agent(
    A=[A_evidence, jnp.asarray(A_outcome)],
    B=[B_motion, jnp.asarray(B_response)],
    C=[jnp.zeros(2), jnp.array([0.0, -2.0, 2.0])],
    D=[jnp.array([0.5, 0.5]), jnp.array([1.0, 0.0, 0.0])],
    A_dependencies=A_dependencies, B_dependencies=B_dependencies,
    # 状態因子1（反応状態）だけを行動で制御する
    control_fac_idx=[1],
    policy_len=1, batch_size=1,
    use_states_info_gain=False, action_selection="deterministic",
)

observation = [jnp.array([EVIDENCE_LEFT]), jnp.array([NEUTRAL])]
qs = agent.infer_states(observation, empirical_prior=agent.D)
q_pi, neg_efe = agent.infer_policies(qs)
action = agent.sample_action(q_pi)

print("q(motion) =", np.asarray(qs[0][0, -1]))
print("q(response policy) =", np.asarray(q_pi[0]))
print("selected action =", np.asarray(action))

posterior_motion = np.asarray(qs[0][0, -1])
response_policy = np.asarray(q_pi[0])

fig, axes = plt.subplots(1, 2, figsize=(8, 3.5))
axes[0].bar(["motion left", "motion right"], posterior_motion)
axes[0].set_title("Posterior belief about motion")
axes[1].bar(["respond left", "respond right"], response_policy)
axes[1].set_title("Response policy probability")
for ax in axes: ax.set_ylim(0, 1)
fig.tight_layout()
plt.show()

NULL, LOSS, REWARD = 0, 1, 2
UNDECIDED, ARM_A, ARM_B = 0, 1, 2
hazard = 0.04

A_outcome = np.zeros((3, 2, 3))
A_outcome[NULL, :, UNDECIDED] = 1.0
reward_prob = np.array([[0.8, 0.2], [0.2, 0.8]])
for context in range(2):
    for arm in range(2):
        p = reward_prob[context, arm]
        A_outcome[REWARD, context, arm + 1] = p
        A_outcome[LOSS, context, arm + 1] = 1 - p

B_context = np.array([[1-hazard, hazard],
                      [hazard, 1-hazard]])[:, :, None]
B_choice = np.zeros((3, 3, 2))
B_choice[ARM_A, :, 0], B_choice[ARM_B, :, 1] = 1.0, 1.0

agent = Agent(
    A=[jnp.asarray(A_outcome), jnp.eye(3)],
    B=[jnp.asarray(B_context), jnp.asarray(B_choice)],
    C=[jnp.array([0., -1., 3.]), jnp.zeros(3)],
    D=[jnp.array([.5, .5]), jnp.array([1., 0., 0.])],
    A_dependencies=[[0, 1], [1]], B_dependencies=[[0], [1]],
    control_fac_idx=[1], policy_len=1, batch_size=1,
    use_states_info_gain=True, action_selection="deterministic",
)

def true_reward_probability(trial, arm):
    good_arm = 0 if trial < 40 else 1
    return 0.8 if arm == good_arm else 0.2

rng = np.random.default_rng(SEED)
qs = agent.infer_states([jnp.array([NULL]), jnp.array([UNDECIDED])],
                        empirical_prior=agent.D)
rows = []
for trial in range(80):
    q_pi, neg_efe = agent.infer_policies(qs)
    action = agent.sample_action(q_pi)
    arm = int(action[0, 1])
    prior = agent.update_empirical_prior(action, qs)
    reward = rng.random() < true_reward_probability(trial, arm)
    outcome = REWARD if reward else LOSS
    qs = agent.infer_states(
        [jnp.array([outcome]), jnp.array([arm + 1])], prior
    )
    rows.append((trial + 1, arm, int(reward),
                 true_reward_probability(trial, arm),
                 float(qs[0][0, -1, 0]), float(q_pi[0, 0]),
                 float(neg_efe[0, 0]), float(neg_efe[0, 1])))
reversal_df = pd.DataFrame(rows, columns=[
    "trial", "arm", "reward", "p_reward_true",
    "q_A_good", "q_choose_A", "neg_efe_A", "neg_efe_B"
])

true_prob_A = np.array([true_reward_probability(t, 0) for t in range(80)])
fig, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
axes[0].plot(reversal_df["trial"], reversal_df["q_A_good"])
axes[0].axhline(0.5, color="0.5", lw=1)
axes[0].set_ylabel("q(A good)")
axes[1].step(reversal_df["trial"], reversal_df["arm"], where="mid")
axes[1].set_yticks([0, 1], ["A", "B"])
axes[1].invert_yaxis()  # Aを上、Bを下に表示する
axes[1].set_ylabel("choice")
axes[2].scatter(reversal_df["trial"], reversal_df["reward"],
                c=reversal_df["reward"], cmap="RdYlGn", vmin=0, vmax=1,
                label="observed reward")
axes[2].step(reversal_df["trial"], true_prob_A, where="mid",
             color="black", ls="--", lw=2,
             label="environment P(reward | A)")
axes[2].set(xlabel="trial", ylabel="reward / probability")
axes[2].legend()
fig.tight_layout()
plt.show()

TRUE_GAMMAS = np.array([0.35, 0.60, 1.00, 1.80, 3.00])
delta_neg_efe_A = (reversal_df["neg_efe_A"]
                   - reversal_df["neg_efe_B"]).to_numpy()
rng = np.random.default_rng(20260906)

choices_A = []
for gamma in TRUE_GAMMAS:
    p_choose_A = jax.nn.sigmoid(
        gamma * jnp.asarray(delta_neg_efe_A)
    )
    choices_A.append(rng.binomial(1, np.asarray(p_choose_A)))
choices_A = np.asarray(choices_A)

q_A_before = np.r_[0.5, reversal_df["q_A_good"].to_numpy()[:-1]]
colors = plt.cm.viridis(np.linspace(.12, .88, len(TRUE_GAMMAS)))
fig, axes = plt.subplots(5, 1, figsize=(11, 10), sharex=True, sharey=True)
for person, (gamma, color, ax) in enumerate(
    zip(TRUE_GAMMAS, colors, axes), start=1
):
    p_choose_A = np.asarray(jax.nn.sigmoid(
        gamma * jnp.asarray(delta_neg_efe_A)
    ))
    chosen_A = choices_A[person - 1]
    ax.plot(np.arange(1, 81), q_A_before, "--", color=".35",
            label="shared q(A good)")
    ax.plot(np.arange(1, 81), p_choose_A, color=color,
            label="P(choice=A)")
    ax.scatter(np.arange(1, 81), chosen_A, color=color, s=10, alpha=.45)
    ax.set(ylabel=f"P{person}", ylim=(-.08, 1.08))
    ax.set_title(f"true gamma = {gamma:.2f}", loc="left", fontsize=10)
axes[0].legend(ncol=2, fontsize=8)
axes[-1].set_xlabel("trial")
fig.tight_layout()
plt.show()

def gamma_choice_model(delta_neg_efe_A, choices_A):
    n = choices_A.shape[0]
    gamma = numpyro.sample(
        "gamma",
        dist.LogNormal(jnp.log(1.0), 1.0).expand([n]).to_event(1),
    )
    logits = gamma[:, None] * delta_neg_efe_A[None, :]
    numpyro.sample(
        "choice_A", dist.Bernoulli(logits=logits).to_event(2),
        obs=choices_A,
    )

mcmc = MCMC(
    NUTS(gamma_choice_model, target_accept_prob=0.90),
    num_warmup=400, num_samples=800, num_chains=4,
    chain_method="sequential",
)
mcmc.run(jr.PRNGKey(20260906), jnp.asarray(delta_neg_efe_A),
         jnp.asarray(choices_A))
gamma_chains = mcmc.get_samples(group_by_chain=True)["gamma"]
gamma_samples = np.asarray(gamma_chains).reshape(-1, len(TRUE_GAMMAS))
gamma_rhat = np.asarray(gelman_rubin(gamma_chains))
gamma_ess = np.asarray(effective_sample_size(gamma_chains))
divergences = int(np.asarray(
    mcmc.get_extra_fields(group_by_chain=True)["diverging"]
).sum())

fig, axes = plt.subplots(5, 1, figsize=(11, 9), sharex=True)
for person, ax in enumerate(axes):
    for chain in range(gamma_chains.shape[0]):
        ax.plot(gamma_chains[chain, :, person], lw=.55, alpha=.75,
                label=f"chain {chain+1}" if person == 0 else None)
    ax.axhline(TRUE_GAMMAS[person], color="black", ls="--", lw=1)
    ax.set_ylabel(f"P{person+1}\ngamma")
    ax.set_title(f"R-hat={gamma_rhat[person]:.3f}, "
                 f"ESS={gamma_ess[person]:.0f}", loc="right", fontsize=9)
axes[0].legend(ncol=4, fontsize=8)
axes[-1].set_xlabel("post-warmup draw")
fig.tight_layout()
plt.show()
print("divergences:", divergences)

gamma_mean = gamma_samples.mean(axis=0)
gamma_low = np.quantile(gamma_samples, .05, axis=0)
gamma_high = np.quantile(gamma_samples, .95, axis=0)
correlation = np.corrcoef(TRUE_GAMMAS, gamma_mean)[0, 1]
rmse = np.sqrt(np.mean((TRUE_GAMMAS - gamma_mean)**2))
coverage = np.mean((TRUE_GAMMAS >= gamma_low) &
                   (TRUE_GAMMAS <= gamma_high))

fig, ax = plt.subplots(figsize=(6, 5))
yerr = np.vstack([gamma_mean-gamma_low, gamma_high-gamma_mean])
ax.errorbar(TRUE_GAMMAS, gamma_mean, yerr=yerr, fmt="o", capsize=4)
limit = max(TRUE_GAMMAS.max(), gamma_high.max()) * 1.08
ax.plot([0, limit], [0, limit], "--", color="black")
ax.set(xlim=(0, limit), ylim=(0, limit), xlabel="true gamma",
       ylabel="estimated gamma")
ax.set_aspect("equal", adjustable="box")
plt.show()
print({"correlation": correlation, "RMSE": rmse,
       "coverage_90": coverage, "max_Rhat": gamma_rhat.max(),
       "min_ESS": gamma_ess.min(), "divergences": divergences})

def normalize_columns(x):
    return x / x.sum(axis=0, keepdims=True)

def build_learning_agent():
    # outcome | choice：A/Bについては弱い対称事前
    pA_outcome = np.ones((3, 3))
    pA_outcome[:, UNDECIDED] = [20.0, 0.1, 0.1]
    pA_outcome[:, ARM_A:] = [[0.1, 0.1],
                                [1.0, 1.0],
                                [1.0, 1.0]]
    # choice観測は既知に近い強い事前
    pA_choice = np.full((3, 3), 0.1)
    np.fill_diagonal(pA_choice, 20.0)
    B_choice = np.zeros((3, 3, 2))
    B_choice[ARM_A, :, 0] = 1.0
    B_choice[ARM_B, :, 1] = 1.0

    agent = Agent(
        A=[jnp.asarray(normalize_columns(pA_outcome)),
           jnp.asarray(normalize_columns(pA_choice))],
        B=[jnp.asarray(B_choice)],
        C=[jnp.array([0., -1., 3.]), jnp.zeros(3)],
        D=[jnp.array([1., 0., 0.])],
        pA=[jnp.asarray(pA_outcome), jnp.asarray(pA_choice)],
        A_dependencies=[[0], [0]], B_dependencies=[[0]],
        control_fac_idx=[0], policy_len=1, batch_size=1,
        learn_A=True, use_param_info_gain=True,
        # alpha=1なら、q(pi)を追加で尖らせずに確率的に行動を選ぶ
        action_selection="stochastic", gamma=4.0, alpha=1.0,
    )
    base_pA = [x.copy() for x in agent.pA]
    return agent, base_pA

learning_agent, base_pA = build_learning_agent()
print("initial P(reward | A/B):",
      np.asarray(learning_agent.A[0][0, REWARD, ARM_A:]))

def apply_dirichlet_forgetting(agent, base_pA, retention):
    new_pA = [base + retention * (current - base)
              for base, current in zip(base_pA, agent.pA)]
    new_A = [p / p.sum(axis=1, keepdims=True) for p in new_pA]
    return eqx.tree_at(lambda x: (x.pA, x.A), agent, (new_pA, new_A))

def good_arm(trial):
    if trial < 40: return 0
    if trial < 80: return 1
    if trial < 120: return ((trial - 80) // 10) % 2
    return 0

def reward_probability(trial, arm):
    return 0.8 if arm == good_arm(trial) else 0.2

def run_volatile_learning(retention, seed=22, n_trials=160):
    agent, base_pA = build_learning_agent()
    rng, key = np.random.default_rng(seed), jr.PRNGKey(seed)
    qs = agent.infer_states(
        [jnp.array([NULL]), jnp.array([UNDECIDED])], agent.D
    )
    rows = []
    for trial in range(n_trials):
        q_pi, _ = agent.infer_policies(qs)
        key, sample_key = jr.split(key)
        action_keys = jr.split(sample_key, agent.batch_size + 1)[1:]
        action = agent.sample_action(q_pi, rng_key=action_keys)
        arm = int(action[0, 0])
        predicted_reward = float(agent.A[0][0, REWARD, arm + 1])
        prior = agent.update_empirical_prior(action, qs)

        p_reward = reward_probability(trial, arm)
        reward = int(rng.random() < p_reward)
        outcome = REWARD if reward else LOSS
        obs = [jnp.array([outcome]), jnp.array([arm + 1])]
        qs = agent.infer_states(obs, empirical_prior=prior)

        if retention < 1.0:
            agent = apply_dirichlet_forgetting(agent, base_pA, retention)
        agent = agent.infer_parameters(
            beliefs_A=qs,
            observations=[jnp.array([[outcome]]),
                          jnp.array([[arm + 1]])],
            actions=None, lr_pA=0.5,
        )
        rows.append({
            "trial": trial + 1, "arm": arm,
            "good_arm": good_arm(trial), "reward": reward,
            "correct_choice": int(arm == good_arm(trial)),
            "predicted_reward": predicted_reward,
            "est_reward_A": float(agent.A[0][0, REWARD, ARM_A]),
            "est_reward_B": float(agent.A[0][0, REWARD, ARM_B]),
        })
    return pd.DataFrame(rows), agent

cumulative_df, _ = run_volatile_learning(retention=1.00)
leaky_df, _ = run_volatile_learning(retention=0.90)

true_prob_A_160 = np.array([reward_probability(t, 0) for t in range(160)])
fig, axes = plt.subplots(2, 2, figsize=(13, 7), sharex=True, sharey="row")
for col, (label, df) in enumerate([
    ("Cumulative counts", cumulative_df),
    ("Forgetting: rho=0.90", leaky_df),
]):
    axes[0, col].plot(df["trial"], df["est_reward_A"], label="estimated A")
    axes[0, col].plot(df["trial"], df["est_reward_B"], label="estimated B")
    axes[0, col].set(title=label, ylabel="P(reward)")
    choice_A = (df["arm"] == 0).rolling(8, min_periods=1).mean()
    axes[1, col].plot(df["trial"], choice_A, color="C3",
                      label="choice A (rolling 8)")
    axes[1, col].step(df["trial"], true_prob_A_160, where="mid",
                      color="black", ls="--", lw=1.8,
                      label="environment P(reward | A)")
    axes[1, col].set(xlabel="trial", ylabel="P(choice = A)",
                     ylim=(-0.05, 1.05))
    for ax in axes[:, col]: ax.axvspan(80.5, 120.5, color="C4", alpha=.06)
axes[0, 0].legend(ncol=2, fontsize=8)
axes[1, 0].legend(ncol=2, fontsize=8)
fig.tight_layout()
plt.show()

ORDINARY, SURPRISING = 0, 1
STABLE, VOLATILE = 0, 1

def build_meta_agent():
    # P(surprising | stable)=.12, P(surprising | volatile)=.55
    A_high = [jnp.array([[0.88, 0.45],
                         [0.12, 0.55]])]
    # 上位状態は高い確率で持続する
    B_high = [jnp.array([[[0.995], [0.020]],
                         [[0.005], [0.980]]])]
    return Agent(
        A=A_high, B=B_high, C=[jnp.zeros(2)],
        D=[jnp.array([0.9, 0.1])],
        policy_len=1, batch_size=1,
        action_selection="deterministic",
    )

higher = build_meta_agent()
print("A_high:", np.asarray(higher.A[0][0]), sep="\n")
print("B_high:", np.asarray(higher.B[0][0, :, :, 0]), sep="\n")

def run_hierarchical_learning(seed=22, n_trials=160):
    lower, base_pA = build_learning_agent()
    higher = build_meta_agent()
    rng, key = np.random.default_rng(seed), jr.PRNGKey(seed)
    q_lower = lower.infer_states(
        [jnp.array([NULL]), jnp.array([UNDECIDED])], lower.D
    )
    prior_high, rows = higher.D, []

    for trial in range(n_trials):
        q_pi, _ = lower.infer_policies(q_lower)
        key, sample_key = jr.split(key)
        action_keys = jr.split(sample_key, lower.batch_size + 1)[1:]
        action = lower.sample_action(q_pi, rng_key=action_keys)
        arm = int(action[0, 0])
        prior_lower = lower.update_empirical_prior(action, q_lower)

        reward = int(rng.random() < reward_probability(trial, arm))
        outcome = REWARD if reward else LOSS
        predicted_p = float(lower.A[0][0, outcome, arm + 1])
        high_obs = SURPRISING if predicted_p < 0.50 else ORDINARY
        q_high = higher.infer_states(
            [jnp.array([high_obs])], empirical_prior=prior_high
        )
        q_volatile = float(q_high[0][0, -1, VOLATILE])
        prior_high = higher.update_empirical_prior(jnp.array([[0]]), q_high)

        retention = 0.995 - 0.15 * q_volatile
        lower = apply_dirichlet_forgetting(lower, base_pA, retention)
        obs = [jnp.array([outcome]), jnp.array([arm + 1])]
        q_lower = lower.infer_states(obs, empirical_prior=prior_lower)
        lower = lower.infer_parameters(
            beliefs_A=q_lower,
            observations=[jnp.array([[outcome]]),
                          jnp.array([[arm + 1]])],
            actions=None, lr_pA=0.5,
        )
        rows.append({
            "trial": trial + 1, "arm": arm, "reward": reward,
            "correct_choice": int(arm == good_arm(trial)),
            "q_volatile": q_volatile, "retention": retention,
            "est_reward_A": float(lower.A[0][0, REWARD, ARM_A]),
            "est_reward_B": float(lower.A[0][0, REWARD, ARM_B]),
        })
    return pd.DataFrame(rows), lower, higher

hierarchical_df, lower_final, higher_final = run_hierarchical_learning()

fig, axes = plt.subplots(3, 1, figsize=(11, 8), sharex=True)
task_volatility = np.where(
    (hierarchical_df["trial"] > 80) &
    (hierarchical_df["trial"] <= 120), 1, 0
)
axes[0].plot(hierarchical_df["trial"], hierarchical_df["q_volatile"],
             label="inferred q(volatile)")
axes[0].step(hierarchical_df["trial"], task_volatility, where="mid",
             color="black", ls=":", label="task volatility")
axes[0].set(ylabel="volatility", ylim=(-.03, 1.03))
axes[0].legend()
axes[1].plot(hierarchical_df["trial"], hierarchical_df["retention"])
axes[1].set_ylabel("retention rho")
axes[2].plot(hierarchical_df["trial"], hierarchical_df["est_reward_A"],
             label="estimated A")
axes[2].plot(hierarchical_df["trial"], hierarchical_df["est_reward_B"],
             label="estimated B")
axes[2].step(hierarchical_df["trial"], hierarchical_df["arm"],
             where="mid", color=".3", alpha=.35, label="choice")
axes[2].step(hierarchical_df["trial"], true_prob_A_160, where="mid",
             color="black", ls="--", lw=1.8,
             label="environment P(reward | A)")
axes[2].set(xlabel="trial", ylabel="probability / choice")
axes[2].legend()
for ax in axes: ax.axvspan(80.5, 120.5, color="C4", alpha=.06)
fig.tight_layout()
plt.show()