"""Train DQN on CartPole and plot the learning curves.

    python run.py                    # Double DQN with replay + target network (~1 min)
    python run.py --no-target        # targets from the online net itself
    python run.py --no-replay        # replay buffer of only 128 (the latest) transitions
    python run.py --no-double        # plain DQN max-targets
    python run.py --compare          # full DQN vs no-target vs no-replay, same seed (~3 min)

Saves dqn_cartpole.png next to this file. Uses your dqn.py, or the reference
solution if RL_SOLUTIONS=1 is set.
"""

import argparse
import importlib
import os
import time
from dataclasses import replace

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

from train import Config, evaluate, train_dqn  # noqa: E402

m = importlib.import_module("solutions.dqn" if os.environ.get("RL_SOLUTIONS") else "dqn")
HERE = os.path.dirname(os.path.abspath(__file__))


def smooth(x, k=20):
    x = np.asarray(x, dtype=float)
    if len(x) < k:
        return x
    return np.convolve(x, np.ones(k) / k, mode="valid")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--steps", type=int, default=Config.total_steps)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--no-target", action="store_true")
    ap.add_argument("--no-replay", action="store_true")
    ap.add_argument("--no-double", action="store_true")
    ap.add_argument("--compare", action="store_true")
    args = ap.parse_args()

    base = Config(total_steps=args.steps, seed=args.seed, double=not args.no_double)
    if args.compare:
        runs = {
            "DQN (replay + target)": base,
            "no target network": replace(base, use_target=False),
            "no replay (buffer = 128)": replace(base, buffer_size=base.batch_size),
        }
    else:
        cfg = base
        label = "DQN (replay + target)"
        if args.no_target:
            cfg, label = replace(cfg, use_target=False), "no target network"
        if args.no_replay:
            cfg = replace(cfg, buffer_size=cfg.batch_size)
            label = "no replay (buffer = 128)" if not args.no_target else "no target, no replay"
        runs = {label: cfg}

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
    summary = []
    for i, (label, cfg) in enumerate(runs.items()):
        print(f"--- {label}")
        t0 = time.time()
        out = train_dqn(m, cfg, log=print)
        final = evaluate(m, out["q_net"], n_episodes=20, seed=999)
        h = out["history"]
        summary.append(f"{label}: final greedy return {final:.1f} (20 episodes), "
                       f"best eval {max(h['eval_return']):.0f}, {time.time() - t0:.0f}s")
        c = f"C{i}"
        ends, rets = np.array(h["episode_end_step"]), np.array(h["episode_return"])
        axes[0].plot(ends, rets, color=c, alpha=0.15, lw=0.8)
        sm = smooth(rets)
        axes[0].plot(ends[len(ends) - len(sm):], sm, color=c, label=label)
        axes[1].plot(h["eval_step"], h["eval_return"], "o-", color=c, ms=3, label=label)
        axes[2].plot(h["loss_step"], h["loss"], color=c, label=label)

    axes[0].set(title="Training episodes (epsilon-greedy)", xlabel="environment step", ylabel="return")
    axes[1].set(title="Greedy evaluation (5 episodes)", xlabel="environment step", ylabel="return")
    axes[2].set(title="Mean Huber loss", xlabel="environment step", yscale="log")
    for ax in axes:
        ax.axhline(500, color="gray", lw=0.6, ls=":") if ax is not axes[2] else None
        ax.legend(fontsize=8)
    fig.tight_layout()
    path = os.path.join(HERE, "dqn_cartpole.png")
    fig.savefig(path, dpi=110)
    for s in summary:
        print(s)
    print(f"saved {path}")


if __name__ == "__main__":
    main()