"""Outils partagés par les méthodes d'apprentissage : évaluation en parallèle sur plusieurs cœurs."""
from __future__ import annotations

import csv
import os
import time
from collections import Counter

import numpy as np

from .. import bots, scenarios, shot
from ..vocab import ctx_of

_BOT = {}


def _bot(key):
    if key not in _BOT:
        _BOT[key] = bots.make(key, path="__aucun__")      # cerveau vierge, on remplace les poids
    return _BOT[key]


def _scenes(bot, seeds):
    scenes = [scenarios.make(int(s)) for s in seeds]
    obs = np.stack([bot.obs(w, s, t) for w, s, t in scenes])
    return scenes, obs


def run_outs(bot, scenes, outs):
    """Tire avec les sorties de cerveau données ; renvoie récompenses et étiquettes."""
    R, L = np.zeros(len(scenes)), []
    for i, ((w, s, t), o) in enumerate(zip(scenes, outs)):
        sh = bot.shot_from(o, ctx_of(w, s, t), t)
        res = shot.simulate(w, s, sh.g, sh.z, t)
        R[i], lab = scenarios.reward(res, w, s, t, bot.ally_penalty)
        L.append(lab)
    return R, L


# ---------------------------------------------------------------- tâches envoyées aux cœurs
def task_eval_flats(args):
    """Évalue plusieurs cerveaux (vecteurs de poids) sur les mêmes situations."""
    key, flats, seeds = args
    bot = _bot(key)
    scenes, obs = _scenes(bot, seeds)
    out = np.zeros((len(flats), len(seeds)))
    labels = Counter()
    for k, f in enumerate(flats):
        bot.policy.set_flat(f)
        R, L = run_outs(bot, scenes, bot.brain_out(obs))
        out[k] = R
        labels.update(L)
    return out, labels


def task_rollout(args):
    """Pour PPO : joue des tirs avec une politique gaussienne (exploration)."""
    key, flat, log_std, seeds, rseed = args
    bot = _bot(key)
    bot.policy.set_flat(flat)
    rng = np.random.default_rng(rseed)
    scenes, obs = _scenes(bot, seeds)
    mean = bot.policy(obs)
    act = mean + np.exp(log_std) * rng.standard_normal(mean.shape)
    R, L = run_outs(bot, scenes, np.tanh(act))
    return obs, act.astype(np.float32), R, L


SANS_FIN = 10 ** 9      # nombre d'itérations utilisé par --sans-limite


class Pool:
    """Petit enrobage de multiprocessing (fonctionne aussi sous Windows)."""

    def __init__(self, workers: int):
        self.workers = max(1, workers)
        self.pool = None
        if self.workers > 1:
            import multiprocessing as mp
            self.pool = mp.get_context("spawn").Pool(self.workers)

    def map(self, fn, jobs):
        return self.pool.map(fn, jobs) if self.pool else list(map(fn, jobs))

    def close(self, force=False):
        if self.pool:
            if force:
                self.pool.terminate()
            else:
                self.pool.close()
            self.pool.join()


def split(arr, n):
    n = max(1, min(n, len(arr)))
    return [a for a in np.array_split(np.asarray(arr), n) if len(a)]


def eval_flats(pool, key, flats, seeds):
    """Matrice des récompenses (cerveaux x situations)."""
    jobs = [(key, chunk, seeds) for chunk in split(flats, pool.workers * 2)]
    res = pool.map(task_eval_flats, jobs)
    labels = Counter()
    for _, l in res:
        labels.update(l)
    return np.vstack([r for r, _ in res]), labels


def validate(pool, key, flat, n=len(scenarios.VAL_SEEDS)):
    """Mesure un cerveau sur les 300 cartes de contrôle (toujours les mêmes)."""
    seeds = scenarios.VAL_SEEDS[:n]
    jobs = [(key, [flat], chunk) for chunk in split(seeds, pool.workers)]
    res = pool.map(task_eval_flats, jobs)
    R = np.concatenate([r[0] for r, _ in res])
    labels = Counter()
    for _, l in res:
        labels.update(l)
    return {
        "recompense": float(R.mean()),
        "cible": labels["cible"] / n,
        "touche": (labels["cible"] + labels["ennemi"]) / n,
        "allie": labels["allié"] / n,
    }


class Logger:
    COLS = ["iter", "secondes", "recompense_entrainement", "recompense_controle",
            "precision_cible", "precision_ennemi", "tirs_allies"]

    def __init__(self, key, resume):
        root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        os.makedirs(os.path.join(root, "logs"), exist_ok=True)
        self.path = os.path.join(root, "logs", f"{key}.csv")
        self.offset = 0                      # en reprise, on continue la numérotation des itérations
        if resume and os.path.exists(self.path):
            with open(self.path) as f:
                rows = list(csv.reader(f))[1:]
            self.offset = int(rows[-1][0]) if rows else 0
        if not resume or not os.path.exists(self.path):
            with open(self.path, "w", newline="") as f:
                csv.writer(f).writerow(self.COLS)
        self.t0 = time.time()

    def log(self, it, train_r, v):
        with open(self.path, "a", newline="") as f:
            csv.writer(f).writerow([it + self.offset, round(time.time() - self.t0, 1), round(train_r, 4), round(v["recompense"], 4),
                                    round(v["cible"], 4), round(v["touche"], 4), round(v["allie"], 4)])


def line(key, it, total, train_r, v, best, extra=""):
    bar = "█" * int(20 * v["touche"]) + "·" * (20 - int(20 * v["touche"]))
    total = "∞" if total >= SANS_FIN else total
    return (f"[{key:10s}] it {it:4d}/{total}  entraîn. {train_r:+.3f}  │ contrôle : {bar} "
            f"{100 * v['touche']:5.1f}% touché, {100 * v['allie']:4.1f}% alliés  │ record {100 * best:5.1f}% {extra}")