import math
import random
import logging
from typing import Optional

from mathema.actors.actor import Actor

log = logging.getLogger(__name__)


def tanh(x): return math.tanh(x)


class Neuron(Actor):
    """
    Neuron actor implementing a weighted-sum neuron with an activation function
    and optional recurrent inputs.

    The Neuron receives input vectors from upstream neurons, accumulates them
    according to its weight configuration, applies an activation function,
    and forwards the resulting output to downstream actors.

    It supports:
        - feed-forward and recurrent connections
        - bias handling
        - asynchronous message-based execution
        - weight backup, restoration, and stochastic perturbation (mutation)
        - cycle-based updates for recurrent networks

    This actor is designed to be used inside a cortex/agent actor network,
    where synchronization and evolution are coordinated externally.
    """
    def __init__(self, nid, cx_pid, af_name, input_idps, output_pids, bias: Optional[float] = None):
        super().__init__(f"Neuron-{nid}")
        self.nid = nid
        self.cx_pid = cx_pid
        self.af = tanh if af_name == "tanh" else tanh

        self.inputs = {}
        self.order = []
        self._has_recurrent = False

        """
        for (inp_id, weights) in input_idps:
            self.order.append(inp_id)
            self.inputs[inp_id] = {"weights": list(weights), "got": False, "val": None}
        """

        self.bias = float(bias) if bias is not None else 0.0
        for inp_id, weights, recurrent in input_idps:
            recurrent = bool(recurrent)
            if inp_id == "bias":
                self.bias = float(weights[0])
            else:
                self.order.append(inp_id)
                self.inputs[inp_id] = {
                    "weights": list(weights),
                    "got": False,
                    "val": [],
                    "recurrent": recurrent,
                    "next_val": []
                }
                if recurrent:
                    self._has_recurrent = True

        self._backup_inputs = None
        self._backup_bias = None

        self.outputs = output_pids

        log.debug(f"Neuron {nid}: inputs={list(self.inputs.keys())}, bias={self.bias}")

    async def run(self):
        while True:
            msg = await self.inbox.get()
            tag = msg[0]

            if tag == "forward":
                _, from_id, data = msg
                if from_id not in self.inputs:
                    continue
                slot = self.inputs[from_id]
                if not isinstance(data, list):
                    data = [float(data)]

                if slot["recurrent"]:
                    slot["next_val"] = data
                else:
                    slot["got"] = True
                    slot["val"] = data

                if all(self.inputs[i]["got"] for i in self.order):
                    acc = 0.0
                    for i in self.order:
                        w = self.inputs[i]["weights"]
                        v = self.inputs[i]["val"]
                        if len(w) != len(v):
                            raise ValueError(f"Lengths of weights and values must be equal")
                        acc += sum(wj * vj for wj, vj in zip(w, v))
                    out = self.af(acc + self.bias)
                    for pid in self.outputs:
                        await pid.send(("forward", self.nid, [out]))
                    for i in self.order:
                        self.inputs[i]["got"] = False
                        self.inputs[i]["val"] = []
                    log.debug(f"Neuron {self.nid}: input_sum={acc + self.bias:.3f}, output={out:.3f}")

            elif tag == "tick":

                if self.order and all(self.inputs[i]["got"] for i in self.order):
                    acc = 0.0
                    for i in self.order:
                        w = self.inputs[i]["weights"]
                        v = self.inputs[i]["val"]
                        if len(w) != len(v):
                            raise ValueError("Lengths of weights and values must be equal")
                        acc += sum(wj * vj for wj, vj in zip(w, v))
                    out = self.af(acc + self.bias)
                    for pid in self.outputs:
                        await pid.send(("forward", self.nid, [out]))
                    for i in self.order:
                        self.inputs[i]["got"] = False
                        self.inputs[i]["val"] = []
                    log.debug(f"Neuron {self.nid}: input_sum={acc + self.bias:.3f}, output={out:.3f}")

            elif tag == "get_backup":
                idps = [(i, self.inputs[i]["weights"]) for i in self.order]
                idps.append(("bias", [self.bias]))
                await self.cx_pid.send(("backup_from_neuron", self.nid, idps))

            elif tag == "weight_backup":
                log.debug(f"Neuron {self.nid}: backing up weights")
                self._backup_inputs = {k: {"weights": v["weights"][:]} for k, v in self.inputs.items()}
                self._backup_bias = self.bias

            elif tag == "weight_restore":
                if self._backup_inputs is not None:
                    for k in self.inputs:
                        self.inputs[k]["weights"] = self._backup_inputs[k]["weights"][:]
                    self.bias = self._backup_bias

            elif tag == "weight_perturb":
                log.debug(
                    f"Neuron {self.nid}: perturbing {len([w for i in self.order for w in self.inputs[i]['weights']])}"
                    f"weights")
                tot_w = sum(len(self.inputs[i]["weights"]) for i in self.order) + 1
                mp = 1 / math.sqrt(tot_w)
                delta_mag = 2.0 * math.pi
                sat_lim = 2.0 * math.pi

                for i in self.order:
                    ws = self.inputs[i]["weights"]
                    for j in range(len(ws)):
                        if random.random() < mp:
                            ws[j] = _sat(ws[j] + (random.random() - 0.5) * delta_mag, -sat_lim, sat_lim)
                if random.random() < mp:
                    self.bias = _sat(self.bias + (random.random() - 0.5) * delta_mag, -sat_lim, sat_lim)

            elif tag == "cycle_start":
                for i in self.order:
                    slot = self.inputs[i]
                    if slot["recurrent"]:
                        nv = slot["next_val"]
                        if not nv:
                            w = slot["weights"]
                            nv = [0.0] * max(1, len(w))
                        slot["val"] = nv
                        slot["got"] = True
                        slot["next_val"] = []
                    else:

                        slot["got"] = False
                        slot["val"] = []

            elif tag == "terminate":
                return


def _sat(val, lo, hi):
    return lo if val < lo else (hi if val > hi else val)