"""Ordered scalar synapse oracle for the canonical bounded CUDA reference."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import TYPE_CHECKING

from src.runtime.determinism import (
    release_uniform as release_uniform,
)
from src.runtime.determinism import (
    ring_slot as ring_slot,
)

if TYPE_CHECKING:
    from ..recurrent.backend import RecurrentInputs


@dataclass(frozen=True)
class SynapseConfig:
    stdp: bool = True
    stp: bool = True
    reward_modulated: bool = True
    seed: int = 12345
    learning_rate: float = 0.2
    weight_decay: float = 0.001
    weight_max: float = 20.0
    eligibility_tau: float = 100.0
    credit_window: int = 64
    td_lambda: float = 0.9
    gamma: float = 0.99

    def parameters(self) -> list[float]:
        return [
            float(self.stdp),
            float(self.stp),
            float(self.reward_modulated),
            float(self.seed),
            self.learning_rate,
            self.weight_decay,
            self.weight_max,
            self.eligibility_tau,
            float(self.credit_window),
            self.td_lambda,
            self.gamma,
        ]

    def validate(self) -> None:
        if any(
            type(value) is not bool
            for value in (self.stdp, self.stp, self.reward_modulated)
        ):
            raise ValueError("synapse mode flags must be boolean")
        if type(self.seed) is not int or not 0 <= self.seed <= 0xFFFFFFFF:
            raise ValueError("synapse seed must be uint32")
        if type(self.credit_window) is not int or not 0 <= self.credit_window <= 2000:
            raise ValueError("credit_window must be in [0,2000]")
        for name, value, low, high in (
            ("learning_rate", self.learning_rate, 0.0, 1.0),
            ("weight_decay", self.weight_decay, 0.0, 1.0),
            ("weight_max", self.weight_max, 0.001, 100.0),
            ("eligibility_tau", self.eligibility_tau, 0.001, 10000.0),
            ("td_lambda", self.td_lambda, 0.0, 1.0),
            ("gamma", self.gamma, 0.0, 1.0),
        ):
            if (
                isinstance(value, bool)
                or not math.isfinite(value)
                or not low <= value <= high
            ):
                raise ValueError(f"invalid synapse {name}")


class SynapseState:
    def __init__(self, inputs: RecurrentInputs, config: SynapseConfig) -> None:
        self.inputs = inputs
        self.config = config
        self.weights = list(inputs.weights)
        self.eligibility = [0.0] * len(inputs.sources)
        self.available = [1.0] * len(inputs.sources)
        self.last_event = [-10_000_000] * len(inputs.sources)
        self.last_spike = [-10_000_000] * inputs.n_neurons
        self.ring_size = max(inputs.delays, default=1) + 1
        self.emitted = [[0.0] * len(inputs.sources) for _ in range(self.ring_size)]

    def incoming(self, tick: int, neuron: int) -> float:
        # Kept for external callers; the membrane oracle sums in edge order.
        return sum(
            self.edge_current(tick, edge)
            for edge in range(
                self.inputs.offsets[neuron], self.inputs.offsets[neuron + 1]
            )
        )

    def edge_current(self, tick: int, edge: int) -> float:
        previous = tick - self.inputs.delays[edge]
        return (
            self.emitted[ring_slot(previous, self.ring_size)][edge]
            if previous >= 0
            else 0.0
        )

    def update(self, tick: int, spikes: list[int], reward: float) -> None:
        c, inputs = self.config, self.inputs
        for target in range(inputs.n_neurons):
            for edge in range(inputs.offsets[target], inputs.offsets[target + 1]):
                source = inputs.sources[edge]
                pre, post = spikes[source], spikes[target]
                weight = self.weights[edge]
                eligibility = self.eligibility[edge] * math.exp(
                    -inputs.dt_ms / c.eligibility_tau
                )
                # Match ascending-neuron pair-rule ordering, including self edges.
                operations = ("pre", "post") if source <= target else ("post", "pre")
                for operation in operations:
                    if operation == "pre" and pre:
                        if c.stdp and 0 < tick - self.last_spike[target] <= 20:
                            weight = max(0.0, weight - 0.08)
                        eligibility -= 0.5
                    if operation == "post" and post:
                        if c.stdp and 0 < tick - self.last_spike[source] <= 20:
                            weight = min(100.0, weight + 0.1)
                        eligibility += 1.0
                if pre or post:
                    self.last_event[edge] = tick
                if (
                    c.reward_modulated
                    and tick - self.last_event[edge] <= c.credit_window
                ):
                    weight = min(
                        100.0,
                        max(
                            0.0,
                            weight
                            + c.learning_rate
                            * c.td_lambda
                            * c.gamma
                            * eligibility
                            * reward,
                        ),
                    )
                amplitude = weight if pre else 0.0
                available = self.available[edge]
                if c.stp:
                    if pre:
                        amplitude *= float(
                            release_uniform(c.seed, tick, edge)
                            < min(0.95, 0.25 + 0.7 * available)
                        )
                        available = max(0.1, available * 0.72)
                    available = min(1.0, available + 0.025)
                self.available[edge] = available
                self.emitted[ring_slot(tick, self.ring_size)][edge] = amplitude
                self.weights[edge] = min(
                    c.weight_max, max(0.0, weight * (1.0 - c.weight_decay))
                )
                self.eligibility[edge] = eligibility
        for neuron, spike in enumerate(spikes):
            if spike:
                self.last_spike[neuron] = tick