"""Actor--critic somato-dendritic innovation for the oral-B-v2 branch.

This module is deliberately separate from :mod:`sdil.bci`: the failed R2
implementation and its source-bound records remain immutable.
"""
from dataclasses import dataclass

import torch

from sdil.bci import BCIConfig, causal_roles


@dataclass(frozen=True)
class BCIV2Config(BCIConfig):
    gamma: float = 0.9
    critic_eta: float = 0.03
    eligibility_decay: float = 0.8
    velocity_reward_scale: float = 0.25
    terminal_reward: float = 1.0

    def validate(self):
        super().validate()
        if self.feedback != "performance_velocity":
            raise ValueError("oral-B-v2 requires performance_velocity")
        if not 0 <= self.gamma < 1:
            raise ValueError("gamma must lie in [0,1)")
        if self.critic_eta <= 0:
            raise ValueError("critic_eta must be positive")
        if not 0 <= self.eligibility_decay < 1:
            raise ValueError("eligibility_decay must lie in [0,1)")
        if self.velocity_reward_scale < 0 or self.terminal_reward <= 0:
            raise ValueError("reward scales are invalid")


class BCIV2:
    """Manual local actor--critic with a role-vectorized TD innovation."""

    def __init__(self, cfg, model_seed=0, *, device="cpu",
                 dtype=torch.float32):
        cfg.validate()
        self.cfg = cfg
        self.device = device
        self.dtype = dtype
        self.role = causal_roles(cfg, device=device, dtype=dtype)
        generator = torch.Generator(device="cpu").manual_seed(model_seed)
        self.W = (
            0.15
            * torch.randn(
                cfg.n_neurons,
                cfg.context_dim,
                generator=generator,
                dtype=dtype,
            )
            / cfg.context_dim ** 0.5
        ).to(device)
        self.A = (
            0.05
            * torch.randn(cfg.n_neurons, generator=generator, dtype=dtype)
        ).to(device)
        self.coupling = (
            cfg.coupling_scale
            * torch.exp(
                0.2
                * torch.randn(
                    cfg.n_neurons, generator=generator, dtype=dtype
                )
            )
        ).to(device)
        self.P = torch.zeros(cfg.n_neurons, device=device, dtype=dtype)
        self.P_bias = torch.zeros(
            cfg.n_neurons, device=device, dtype=dtype
        )
        # A scalar value circuit reads the surrounding population. Its update
        # is a local delta rule driven by the same scalar TD innovation.
        self.critic = torch.zeros(
            cfg.n_neurons + 1, device=device, dtype=dtype
        )
        self.model_seed = model_seed

    def clone(self):
        other = BCIV2(
            self.cfg,
            self.model_seed,
            device=self.device,
            dtype=self.dtype,
        )
        for name in (
            "W",
            "A",
            "coupling",
            "P",
            "P_bias",
            "critic",
        ):
            setattr(other, name, getattr(self, name).clone())
        return other

    def value_features(self, soma):
        return torch.cat(
            (
                torch.ones(
                    soma.shape[0],
                    1,
                    device=soma.device,
                    dtype=soma.dtype,
                ),
                soma,
            ),
            dim=1,
        )

    def value(self, features):
        return features @ self.critic

    def causal_role_targets(self, cursor_plus, cursor_minus, perturbations):
        sigma = self.cfg.perturb_sigma
        directional_derivative = (
            cursor_plus - cursor_minus
        ) / (2.0 * sigma)
        return directional_derivative.unsqueeze(1) * perturbations

    def update_predictor(self, soma, ordinary, active):
        if not bool(active.any()):
            return
        selected_soma = soma[active]
        selected_target = ordinary[active]
        residual = selected_target - (
            self.P * selected_soma + self.P_bias
        )
        soma_centered = selected_soma - selected_soma.mean(0)
        residual_centered = residual - residual.mean(0)
        variance = soma_centered.square().mean(0)
        self.P += self.cfg.predictor_eta * (
            (residual_centered * soma_centered).mean(0)
            / (variance + 1e-6)
        )
        self.P_bias += self.cfg.predictor_eta * residual.mean(0)

    def update_role(self, target, active):
        if not bool(active.any()):
            return
        self.A += self.cfg.vectorizer_eta * (
            target[active].mean(0) - self.A
        )

    def update_critic(self, previous_features, delta, active):
        if not bool(active.any()):
            return
        self.critic += self.cfg.critic_eta * (
            delta[active].unsqueeze(1) * previous_features[active]
        ).mean(0)

    def initial_episode_state(self, episodes):
        soma = torch.zeros(
            episodes,
            self.cfg.n_neurons,
            device=self.device,
            dtype=self.dtype,
        )
        features = self.value_features(soma)
        return {
            "soma": soma,
            "abs_error": torch.full(
                (episodes,),
                abs(self.cfg.target),
                device=self.device,
                dtype=self.dtype,
            ),
            "value_features": features,
            "value": self.value(features),
            "eligibility": torch.zeros(
                episodes,
                self.cfg.n_neurons,
                self.cfg.context_dim,
                device=self.device,
                dtype=self.dtype,
            ),
            "active": torch.ones(
                episodes, device=self.device, dtype=torch.bool
            ),
            "success": torch.zeros(
                episodes, device=self.device, dtype=torch.bool
            ),
        }

    def step(
        self,
        context,
        process_noise,
        perturbations,
        state,
        global_step,
        *,
        final_step,
        plasticity_gain=1.0,
        learn_role=True,
        probe_role=True,
        learn_predictor=True,
        learn_critic=True,
        critic_enabled=True,
        terminal_outcome_enabled=True,
    ):
        """Advance one BCI transition and apply only manual local updates."""
        with torch.no_grad():
            active = state["active"]
            u = (
                context @ self.W.t()
                + self.cfg.inertia * state["soma"]
                + process_noise
            )
            proposed_soma = torch.tanh(u)
            soma = torch.where(
                active.unsqueeze(1), proposed_soma, state["soma"]
            )
            cursor = soma @ self.role
            error = self.cfg.target - cursor
            abs_error = error.abs()
            velocity = state["abs_error"] - abs_error
            newly_successful = active & (cursor >= self.cfg.target)
            terminal = active & (newly_successful | final_step)
            outcome = newly_successful.to(self.dtype)

            features = self.value_features(soma)
            current_value = (
                self.value(features)
                if critic_enabled
                else torch.zeros_like(error)
            )
            previous_value = (
                state["value"]
                if critic_enabled
                else torch.zeros_like(error)
            )
            terminal_bonus = (
                self.cfg.terminal_reward * outcome
                if terminal_outcome_enabled
                else torch.zeros_like(outcome)
            )
            reward = (
                self.cfg.velocity_reward_scale * velocity + terminal_bonus
            )
            bootstrap = torch.where(
                terminal, torch.zeros_like(current_value), current_value
            )
            delta = (
                reward
                + self.cfg.gamma * bootstrap
                - previous_value
            )
            delta = delta * active

            ordinary = self.coupling * soma
            instructional = delta.unsqueeze(1) * self.A.unsqueeze(0)
            raw = ordinary + instructional
            baseline = self.P * soma + self.P_bias
            innovation = (raw - baseline) * active.unsqueeze(1)

            gain = 1.0 - soma.square()
            local_eligibility = (
                gain.unsqueeze(2) * context.unsqueeze(1)
            )
            eligibility = (
                self.cfg.eligibility_decay * state["eligibility"]
                + local_eligibility
            )
            eligibility = eligibility * active[:, None, None]

            did_perturb = (
                probe_role and global_step % self.cfg.perturb_every == 0
            )
            causal_target = None
            if did_perturb:
                sigma = self.cfg.perturb_sigma
                cursor_plus = (soma + sigma * perturbations) @ self.role
                cursor_minus = (soma - sigma * perturbations) @ self.role
                causal_target = self.causal_role_targets(
                    cursor_plus, cursor_minus, perturbations
                )

            if plasticity_gain and bool(active.any()):
                update = (
                    innovation[:, :, None] * eligibility
                )[active].mean(0)
                self.W += plasticity_gain * self.cfg.forward_eta * update
            if did_perturb and learn_role:
                self.update_role(causal_target, active)
            if learn_critic and critic_enabled:
                self.update_critic(
                    state["value_features"], delta, active
                )
            if learn_predictor:
                self.update_predictor(soma, ordinary, active)

            next_active = active & ~terminal
            next_success = state["success"] | newly_successful
            next_features = torch.where(
                next_active.unsqueeze(1),
                features,
                state["value_features"],
            )
            # Re-evaluate the continuing state's value after the local critic
            # update.  This is the usual online TD(0) convention: the next
            # transition starts from V(s_{t+1}; theta_{t+1}), rather than a
            # stale prediction made with theta_t.
            updated_current_value = (
                self.value(features)
                if critic_enabled
                else torch.zeros_like(error)
            )
            next_value = torch.where(
                next_active,
                updated_current_value,
                state["value"],
            )
            next_state = {
                "soma": soma,
                "abs_error": torch.where(
                    next_active, abs_error, state["abs_error"]
                ),
                "value_features": next_features,
                "value": next_value,
                "eligibility": eligibility,
                "active": next_active,
                "success": next_success,
            }
        record = {
            "soma": soma,
            "base_soma": soma,
            "cursor": cursor,
            "error": error,
            "abs_error": abs_error,
            "previous_abs_error": state["abs_error"],
            "performance_velocity": velocity,
            "value_prediction": previous_value,
            "next_value_prediction": current_value,
            "reward": reward,
            "td_innovation": delta,
            "raw_apical": raw,
            "innovation": innovation,
            "ordinary_apical": ordinary,
            "instructional_apical": instructional,
            "eligibility": eligibility,
            "context": context,
            "causal_target": causal_target,
            "did_perturb": did_perturb,
            "active": active,
            "terminal": terminal,
            "outcome": outcome.bool(),
        }
        return next_state, record


def run_day_v2(
    model,
    trajectories,
    day,
    *,
    plasticity_gain=1.0,
    learn_role=True,
    probe_role=True,
    learn_predictor=True,
    learn_critic=True,
    critic_enabled=True,
    terminal_outcome_enabled=True,
    collect=False,
    global_step=0,
    horizon=None,
):
    """Run a paired episode batch with an explicit terminal outcome event."""
    episodes = trajectories.context.shape[1]
    available_steps = trajectories.context.shape[2]
    horizon = model.cfg.steps_per_episode if horizon is None else horizon
    if not 1 <= horizon <= min(model.cfg.steps_per_episode, available_steps):
        raise ValueError("v2 horizon exceeds the available trajectory")
    state = model.initial_episode_state(episodes)
    events = []
    active_transitions = 0
    role_probe_examples = 0
    for step_index in range(horizon):
        previous_soma = state["soma"]
        state, record = model.step(
            trajectories.context[day, :, step_index],
            trajectories.process_noise[day, :, step_index],
            trajectories.perturbations[day, :, step_index],
            state,
            global_step,
            final_step=step_index == horizon - 1,
            plasticity_gain=plasticity_gain,
            learn_role=learn_role,
            probe_role=probe_role,
            learn_predictor=learn_predictor,
            learn_critic=learn_critic,
            critic_enabled=critic_enabled,
            terminal_outcome_enabled=terminal_outcome_enabled,
        )
        if collect:
            event = {
                key: value.detach().cpu().clone()
                for key, value in record.items()
                if isinstance(value, torch.Tensor)
            }
            event["previous_soma"] = (
                previous_soma.detach().cpu().clone()
            )
            event["step"] = step_index
            events.append(event)
        active_count = int(record["active"].sum().item())
        active_transitions += active_count
        if record["did_perturb"]:
            role_probe_examples += active_count
        global_step += 1
    if bool(state["active"].any()):
        raise RuntimeError("every v2 episode must terminate")
    return {
        "success": state["success"].detach().cpu(),
        "success_rate": state["success"].float().mean().item(),
        "events": events,
        "global_step": global_step,
        "active_transitions": active_transitions,
        "role_probe_examples": role_probe_examples,
    }