"""PPO (Proximal Policy Optimization) pour Irrational, Peace et 3D.

Chaque « partie » d'entraînement est un seul tir : le bot voit la situation, choisit
les coefficients de sa fonction, et reçoit sa récompense. Le réseau de politique donne
une moyenne, on ajoute du bruit pour explorer, et on renforce ce qui a mieux marché que
prévu par le réseau de valeur (le « critique »).

Le réseau de valeur sert aussi en combat : il estime la chance de toucher chaque ennemi,
et le bot vise celui où il se sent le plus sûr de lui.
"""
from __future__ import annotations

import numpy as np

from ..nets import MLP
from .common import Logger, line, split, task_rollout, validate


def _torch():
    try:
        import torch
        return torch
    except ImportError as e:
        raise SystemExit("PPO a besoin de PyTorch :  pip install torch") from e


def _mlp(torch, sizes):
    layers = []
    for i, (a, b) in enumerate(zip(sizes[:-1], sizes[1:])):
        layers.append(torch.nn.Linear(a, b))
        if i < len(sizes) - 2:
            layers.append(torch.nn.Tanh())
    return torch.nn.Sequential(*layers)


def _to_flat(net):
    parts = []
    for m in net:
        if hasattr(m, "weight"):
            parts += [m.weight.detach().cpu().numpy().T.ravel(), m.bias.detach().cpu().numpy().ravel()]
    return np.concatenate(parts).astype(np.float32)


def _from_flat(torch, net, flat):
    i = 0
    with torch.no_grad():
        for m in net:
            if hasattr(m, "weight"):
                a, b = m.in_features, m.out_features
                m.weight.copy_(torch.tensor(flat[i:i + a * b].reshape(a, b).T)); i += a * b
                m.bias.copy_(torch.tensor(flat[i:i + b])); i += b


def train_ppo(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
              batch=1536, epochs=8, minibatch=256, lr=1e-3, clip=0.2, ent=0.0, log_std0=-1.0):
    torch = _torch()
    torch.manual_seed(int(rng.integers(1 << 30)))
    sizes = bot.policy.sizes
    vsizes = [sizes[0], *bot.hidden, 1]
    pi, vf = _mlp(torch, sizes), _mlp(torch, vsizes)
    log_std = torch.nn.Parameter(torch.full((sizes[-1],), log_std0))
    if bot.trained_iters > 0:                       # reprise
        _from_flat(torch, pi, bot.policy.flat)
        if bot.value is not None:
            _from_flat(torch, vf, bot.value.flat)
    opt = torch.optim.Adam([*pi.parameters(), *vf.parameters(), log_std], lr=lr)
    log = Logger(bot.key, resume)
    best = best0

    def export():
        bot.policy.set_flat(_to_flat(pi))
        bot.value = MLP(vsizes, "linear", flat=_to_flat(vf))

    for it in range(1, iters + 1):
        flat = _to_flat(pi)
        ls = log_std.detach().numpy().copy()
        seeds = rng.integers(0, 2 ** 31 - 1, batch)
        jobs = [(bot.key, flat, ls, chunk, int(rng.integers(1 << 30))) for chunk in split(seeds, pool.workers * 2)]
        res = pool.map(task_rollout, jobs)
        obs = torch.tensor(np.concatenate([r[0] for r in res]), dtype=torch.float32)
        act = torch.tensor(np.concatenate([r[1] for r in res]), dtype=torch.float32)
        rew = torch.tensor(np.concatenate([r[2] for r in res]), dtype=torch.float32)

        with torch.no_grad():
            dist = torch.distributions.Normal(pi(obs), log_std.exp())
            old_logp = dist.log_prob(act).sum(-1)
            adv = rew - vf(obs).squeeze(-1)
            adv = (adv - adv.mean()) / (adv.std() + 1e-8)

        n = len(rew)
        for _ in range(epochs):
            perm = torch.randperm(n)
            for k in range(0, n, minibatch):
                idx = perm[k:k + minibatch]
                dist = torch.distributions.Normal(pi(obs[idx]), log_std.exp())
                logp = dist.log_prob(act[idx]).sum(-1)
                ratio = (logp - old_logp[idx]).exp()
                a = adv[idx]
                loss_pi = -torch.min(ratio * a, ratio.clamp(1 - clip, 1 + clip) * a).mean()
                loss_v = ((vf(obs[idx]).squeeze(-1) - rew[idx]) ** 2).mean()
                loss = loss_pi + 0.5 * loss_v - ent * dist.entropy().sum(-1).mean()
                opt.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_([*pi.parameters(), *vf.parameters(), log_std], 0.5)
                opt.step()
            with torch.no_grad():
                log_std.clamp_(-2.5, 0.5)

        if it % val_every == 0 or it == iters:
            export()
            v = validate(pool, bot.key, bot.policy.flat)
            log.log(it, float(rew.mean()), v)
            if v["touche"] > best:
                best = v["touche"]
                bot.trained_iters = it + log.offset
                save(bot, "best")
            say(line(bot.key, it + log.offset, iters + log.offset, float(rew.mean()), v, best,
                     f"(bruit {float(log_std.detach().exp().mean()):.2f})"))
    export(); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best