"""Un petit réseau de neurones en numpy (perceptron multicouche).

Tous les bots qui apprennent utilisent ce format pour jouer, ce qui permet de lancer
l'arène sans PyTorch. Les poids de départ sont tirés au hasard : le bot ne sait rien.
"""
from __future__ import annotations

import numpy as np


class MLP:
    def __init__(self, sizes, out_act="tanh", rng=None, flat=None):
        self.sizes = list(sizes)
        self.out_act = out_act
        self.shapes = [(a, b) for a, b in zip(self.sizes[:-1], self.sizes[1:])]
        self.n_params = sum(a * b + b for a, b in self.shapes)
        if flat is None:
            rng = rng or np.random.default_rng()
            flat = np.concatenate([
                np.concatenate([rng.normal(0, 1 / np.sqrt(a), a * b), np.zeros(b)]) for a, b in self.shapes
            ])
        self.set_flat(flat)

    def set_flat(self, flat):
        flat = np.asarray(flat, dtype=np.float32)
        assert flat.size == self.n_params, (flat.size, self.n_params)
        self.flat = flat
        self.layers, i = [], 0
        for a, b in self.shapes:
            W = flat[i:i + a * b].reshape(a, b); i += a * b
            bias = flat[i:i + b]; i += b
            self.layers.append((W, bias))

    def __call__(self, x):
        h = np.asarray(x, dtype=np.float32)
        for k, (W, b) in enumerate(self.layers):
            h = h @ W + b
            if k < len(self.layers) - 1:
                h = np.tanh(h)
        if self.out_act == "tanh":
            h = np.tanh(h)
        return h

    # ------------------------------------------------------------- sauvegarde
    def to_dict(self, prefix=""):
        return {prefix + "sizes": np.array(self.sizes), prefix + "flat": self.flat,
                prefix + "out_act": np.array(self.out_act)}

    @staticmethod
    def from_dict(d, prefix=""):
        return MLP(d[prefix + "sizes"].tolist(), str(d[prefix + "out_act"]), flat=d[prefix + "flat"])