"""Trois façons d'apprendre sans gradient, en ne regardant que le score :

- algorithme génétique (Evolution) : population, sélection, croisement, mutation
- stratégies d'évolution OpenAI-ES (Pole-Zero) : un seul cerveau, on estime dans quelle
  direction bouger ses poids en essayant des petites secousses au hasard
- entropie croisée CEM (Fly) : une moyenne et un écart-type par poids, on recentre sur l'élite
"""
from __future__ import annotations

import numpy as np

from ..nets import MLP
from .common import Logger, eval_flats, line, validate


INIT_OUT = 0.1          # dernière couche petite (comme d'habitude en RL) : au départ le bot tire presque droit


def _start(bot, rng, pop):
    """Population de départ : des cerveaux tirés au hasard (ou repris d'une sauvegarde)."""
    sizes = bot.policy.sizes
    a, b = sizes[-2], sizes[-1]
    flats = []
    for _ in range(pop):
        f = MLP(sizes, rng=rng).flat.copy()
        f[-(a * b + b):] *= INIT_OUT
        flats.append(f)
    if bot.trained_iters > 0:
        flats[0] = bot.policy.flat.copy()
    return np.array(flats)


def _finish(bot, flat, it, v, best, log, save):
    """À chaque contrôle : sauvegarde le record s'il est battu, et toujours l'état courant (_last)
    pour pouvoir reprendre avec --resume, même si l'entraînement est arrêté avec Ctrl+C."""
    bot.policy.set_flat(flat)
    bot.trained_iters = it + log.offset
    if v["touche"] > best:
        best = v["touche"]
        save(bot, "best")
    save(bot, "last")
    return best


def train_ga(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
             pop=48, n_seeds=128, elite=6, mut_sigma=0.02, cross_p=0.2):
    """Algorithme génétique.

    Les tirs sont très bruités (une carte facile ou difficile change tout) : pour ne pas garder
    un cerveau juste parce qu'il a eu de la chance, chaque élite garde la moyenne de tous ses
    scores depuis sa naissance. Plus il survit, plus son score est fiable.
    """
    log = Logger(bot.key, resume)
    P = _start(bot, rng, pop)
    h_sum, h_n = np.zeros(pop), np.zeros(pop)
    best, champion = best0, P[0]
    for it in range(1, iters + 1):
        seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
        F, _ = eval_flats(pool, bot.key, P, seeds)
        h_sum += F.mean(axis=1)
        h_n += 1
        fit = h_sum / h_n
        order = np.argsort(-fit)
        champion = P[order[0]].copy()

        # ---- génération suivante : les élites survivent (avec leur historique), les autres sont remplacés
        keep = order[:elite]
        children, cs, cn = [P[i].copy() for i in keep], list(h_sum[keep]), list(h_n[keep])
        parents = order[: pop // 4]                                          # le meilleur quart se reproduit
        while len(children) < pop:
            a = P[rng.choice(parents)]
            if rng.random() < cross_p:                                       # croisement : moitié des poids de chacun
                b = P[rng.choice(parents)]
                child = np.where(rng.random(a.size) < 0.5, a, b)
            else:
                child = a.copy()
            child = child + rng.normal(0, mut_sigma, child.size).astype(np.float32)
            children.append(child); cs.append(0.0); cn.append(0.0)
        P, h_sum, h_n = np.array(children, dtype=np.float32), np.array(cs), np.array(cn)

        if it % val_every == 0 or it == iters:
            v = validate(pool, bot.key, champion)
            log.log(it, float(fit[order[0]]), v)
            best = _finish(bot, champion, it, v, best, log, save)
            say(line(bot.key, it + log.offset, iters + log.offset, float(fit[order[0]]), v, best,
                     f"(champion vieux de {int(h_n[0])} gén.)"))
    bot.policy.set_flat(champion); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best


class _Adam:
    def __init__(self, n, lr):
        self.m, self.v, self.t, self.lr = np.zeros(n), np.zeros(n), 0, lr

    def step(self, g):
        self.t += 1
        self.m = 0.9 * self.m + 0.1 * g
        self.v = 0.999 * self.v + 0.001 * g * g
        mh, vh = self.m / (1 - 0.9 ** self.t), self.v / (1 - 0.999 ** self.t)
        return self.lr * mh / (np.sqrt(vh) + 1e-8)


def _ranks(x):
    r = np.empty(len(x)); r[np.argsort(x)] = np.arange(len(x))
    return r / (len(x) - 1) - 0.5


def train_es(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
             pairs=24, n_seeds=64, sigma=0.02, lr=0.01, decay=0.005):
    log = Logger(bot.key, resume)
    theta = _start(bot, rng, 1)[0].astype(np.float64)
    opt, best = _Adam(theta.size, lr), best0
    for it in range(1, iters + 1):
        seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
        eps = rng.standard_normal((pairs, theta.size))
        cand = np.vstack([theta + sigma * eps, theta - sigma * eps, theta[None]])   # secousses + / -, + le centre
        F, _ = eval_flats(pool, bot.key, cand.astype(np.float32), seeds)
        fit = F.mean(axis=1)
        r = _ranks(fit[:-1])
        grad = (r[:pairs] - r[pairs:]) @ eps / (2 * pairs * sigma)
        theta += opt.step(grad - decay * theta)
        if it % val_every == 0 or it == iters:
            v = validate(pool, bot.key, theta.astype(np.float32))
            log.log(it, float(fit[-1]), v)
            best = _finish(bot, theta.astype(np.float32), it, v, best, log, save)
            say(line(bot.key, it + log.offset, iters + log.offset, float(fit[-1]), v, best))
    bot.policy.set_flat(theta.astype(np.float32)); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best


def train_cem(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
              pop=64, n_seeds=96, elite_frac=0.25, sigma0=0.04, extra_noise=0.005):
    log = Logger(bot.key, resume)
    mu = _start(bot, rng, 1)[0].astype(np.float64)
    std = np.full(mu.size, sigma0)
    n_el, best = max(2, int(pop * elite_frac)), best0
    for it in range(1, iters + 1):
        seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
        cand = mu + std * rng.standard_normal((pop, mu.size))
        cand[0] = mu
        F, _ = eval_flats(pool, bot.key, cand.astype(np.float32), seeds)
        fit = F.mean(axis=1)
        el = cand[np.argsort(-fit)[:n_el]]
        noise = extra_noise * max(0.0, 1 - it / (0.8 * iters))
        mu = 0.5 * el.mean(axis=0) + 0.5 * mu
        std = np.sqrt(0.5 * el.var(axis=0) + 0.5 * std ** 2) + noise
        if it % val_every == 0 or it == iters:
            v = validate(pool, bot.key, mu.astype(np.float32))
            log.log(it, float(fit[0]), v)
            best = _finish(bot, mu.astype(np.float32), it, v, best, log, save)
            say(line(bot.key, it + log.offset, iters + log.offset, float(fit[0]), v, best, f"(σ moyen {std.mean():.3f})"))
    bot.policy.set_flat(mu.astype(np.float32)); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best