"""Train A2C on CartPole with several values of the GAE parameter lambda and plot the
learning curves (mean over seeds, shaded min..max). Takes about 1-2 minutes on a laptop CPU.

    python run.py                      # your actor_critic.py
    RL_SOLUTIONS=1 python run.py       # the reference solution
"""

import importlib
import os
import time

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

ac = importlib.import_module("solutions.actor_critic" if os.environ.get("RL_SOLUTIONS") else "actor_critic")

LAMBDAS = [0.0, 0.5, 0.95, 1.0]
SEEDS = [0, 1, 2]
STEPS = 200_000

if __name__ == "__main__":
    t0 = time.time()
    fig, ax = plt.subplots(figsize=(7, 4.2))
    summary = []
    for lam in LAMBDAS:
        curves = []
        for seed in SEEDS:
            h = ac.train_a2c(total_steps=STEPS, lam=lam, seed=seed, log_every=10_000)
            curves.append(h["avg_return"])
            steps = h["steps"]
        curves = np.array(curves)
        mean = curves.mean(0)
        ax.plot(steps, mean, label=f"λ = {lam}")
        ax.fill_between(steps, curves.min(0), curves.max(0), alpha=0.15)
        summary.append(f"λ={lam}: {mean[len(mean) // 4]:.0f} at {steps[len(steps) // 4] // 1000}k, {mean[-1]:.0f} at end")
    ax.set_xlabel("environment steps")
    ax.set_ylabel("average return (last 20 episodes)")
    ax.set_title(f"A2C on CartPole: effect of GAE λ ({len(SEEDS)} seeds, γ = 0.99)")
    ax.set_ylim(0, 520)
    ax.legend()
    fig.tight_layout()
    out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "a2c_lambda.png")
    fig.savefig(out, dpi=120)
    print(f"saved {out} in {time.time() - t0:.0f}s | " + " | ".join(summary))