"""The DQN training loop (provided: you don't need to edit this file).

It is written to be read. Each numbered comment matches a line of the DQN
pseudocode in the chapter. It calls the pieces you write in dqn.py:

    ReplayBuffer, QNetwork, epsilon_by_step, select_action,
    update_target, train_step

`train_dqn(m, ...)` takes the module `m` (your dqn.py or solutions/dqn.py) as an
argument, so the same loop runs your code and the reference code.
"""

import copy
from dataclasses import dataclass

import numpy as np
import torch

from cartpole import CartPole

torch.set_num_threads(1)


@dataclass
class Config:
    total_steps: int = 50_000  # environment steps
    gamma: float = 0.99
    lr: float = 5e-4
    batch_size: int = 128
    buffer_size: int = 50_000  # replay capacity (set small to "turn replay off")
    learning_starts: int = 1_000  # collect this many random-ish steps before learning
    train_every: int = 1  # one gradient step per environment step
    target_update_every: int = 500  # C: hard-copy theta -> theta_minus every C steps
    use_target: bool = True  # False: targets use the online net itself (naive deep Q)
    double: bool = True  # Double DQN targets
    eps_start: float = 1.0
    eps_end: float = 0.01
    eps_decay_steps: int = 10_000
    hidden: int = 128
    seed: int = 0


def evaluate(m, q_net, n_episodes=10, seed=12345):
    """Average return of the greedy policy (epsilon = 0) over n_episodes."""
    env = CartPole(seed=seed)
    rng = np.random.default_rng(seed)
    returns = []
    for _ in range(n_episodes):
        s, done, ret = env.reset(), False, 0.0
        while not done:
            a = m.select_action(q_net, s, 0.0, rng)
            s, r, done = env.step(a)
            ret += r
        returns.append(ret)
    return float(np.mean(returns))


def train_dqn(m, cfg=None, log=None, eval_every=2_500, eval_episodes=5):
    """Train DQN on CartPole. Returns a dict with the trained network and a history.

    log: optional function called with a short progress string.
    """
    cfg = cfg or Config()
    torch.manual_seed(cfg.seed)
    rng = np.random.default_rng(cfg.seed)
    env = CartPole(seed=cfg.seed)

    # 1. Initialise the online network Q(s, a; theta), the target network
    #    theta_minus <- theta, and an empty replay buffer D.
    q_net = m.QNetwork(env.obs_dim, env.n_actions, hidden=cfg.hidden)
    target_net = copy.deepcopy(q_net) if cfg.use_target else q_net
    optimizer = torch.optim.Adam(q_net.parameters(), lr=cfg.lr)
    buffer = m.ReplayBuffer(cfg.buffer_size, env.obs_dim, seed=cfg.seed)

    hist = {"episode_end_step": [], "episode_return": [], "eval_step": [], "eval_return": [],
            "loss_step": [], "loss": []}
    s, ep_ret, losses = env.reset(), 0.0, []
    for step in range(cfg.total_steps):
        # 2. Act epsilon-greedily with the online network.
        eps = m.epsilon_by_step(step, cfg.eps_start, cfg.eps_end, cfg.eps_decay_steps)
        a = m.select_action(q_net, s, eps, rng)
        s_next, r, done = env.step(a)
        ep_ret += r

        # 3. Store the transition. A time-limit cut-off (500 steps) is NOT a
        #    terminal state: the pole was still up, the future still had value.
        #    So we only mark `terminal` when the pole actually fell.
        terminal = done and env.t < env.max_steps
        buffer.push(s, a, r, s_next, terminal)

        s = s_next
        if done:
            hist["episode_end_step"].append(step + 1)
            hist["episode_return"].append(ep_ret)
            s, ep_ret = env.reset(), 0.0

        # 4. Learn: one minibatch gradient step on the DQN loss.
        if step >= cfg.learning_starts and step % cfg.train_every == 0:
            losses.append(m.train_step(q_net, target_net, optimizer, buffer,
                                       cfg.batch_size, cfg.gamma, double=cfg.double))

        # 5. Every C steps, copy the online weights into the target network.
        if cfg.use_target and step % cfg.target_update_every == 0:
            m.update_target(target_net, q_net, tau=1.0)

        if (step + 1) % eval_every == 0:
            ev = evaluate(m, q_net, n_episodes=eval_episodes, seed=10_000 + step)
            hist["eval_step"].append(step + 1)
            hist["eval_return"].append(ev)
            if losses:
                hist["loss_step"].append(step + 1)
                hist["loss"].append(float(np.mean(losses)))
            if log:
                recent = hist["episode_return"][-10:]
                log(f"step {step + 1:6d}  eps {eps:.2f}  "
                    f"train return (last 10) {np.mean(recent) if recent else 0:6.1f}  "
                    f"greedy eval {ev:6.1f}  loss {np.mean(losses) if losses else 0:.4f}")
            losses = []

    return {"q_net": q_net, "target_net": target_net, "history": hist, "config": cfg}