"""Continuous BCI dynamics for the preregistered Harnett-signature study.

Learning is manual and local. Autograd is never enabled. The environment's
scalar loss is observed only through antithetic neural perturbations; the
resulting causal targets calibrate a cell-specific apical vectorizer.
"""
from dataclasses import dataclass

import torch


@dataclass(frozen=True)
class BCIConfig:
    n_plus: int = 5
    n_minus: int = 5
    n_background: int = 30
    context_dim: int = 16
    steps_per_episode: int = 28
    episodes_per_day: int = 64
    days: int = 14
    target: float = 0.8
    inertia: float = 0.65
    process_noise: float = 0.12
    context_ar: float = 0.8
    coupling_scale: float = 1.0
    predictor_eta: float = 0.2
    vectorizer_eta: float = 0.03
    forward_eta: float = 0.01
    perturb_sigma: float = 0.03
    perturb_every: int = 4
    kappa: float = 0.0
    feedback: str = "error"

    @property
    def n_neurons(self):
        return self.n_plus + self.n_minus + self.n_background

    @property
    def feedback_dim(self):
        return 1 if self.feedback in ("error", "performance_velocity") else 2

    def validate(self):
        if self.feedback not in (
                "error", "error_velocity", "performance_velocity"):
            raise ValueError(f"unknown BCI feedback: {self.feedback}")
        if self.n_plus < 1 or self.n_minus < 1 or self.n_background < 1:
            raise ValueError("all BCI populations must be nonempty")
        if self.context_dim < 2:
            raise ValueError("context_dim includes a bias and must be at least two")
        if self.steps_per_episode < 2 or self.episodes_per_day < 2 or self.days < 2:
            raise ValueError("BCI trajectories require multiple steps, episodes, and days")
        if not 0 <= self.context_ar < 1:
            raise ValueError("context_ar must lie in [0,1)")
        if self.perturb_every < 1 or self.perturb_sigma <= 0:
            raise ValueError("perturbation schedule must be positive")


@dataclass
class BCITrajectories:
    context: torch.Tensor
    process_noise: torch.Tensor
    perturbations: torch.Tensor
    task_seed: int

    @property
    def shape(self):
        return self.context.shape[:3]


def causal_roles(cfg, *, device="cpu", dtype=torch.float32):
    """Experimenter-defined P+/P-/P0 contribution to the scalar cursor."""
    role = torch.zeros(cfg.n_neurons, device=device, dtype=dtype)
    role[:cfg.n_plus] = 1.0 / cfg.n_plus
    role[cfg.n_plus:cfg.n_plus + cfg.n_minus] = -1.0 / cfg.n_minus
    return role


def generate_trajectories(cfg, task_seed, *, days=None, episodes=None,
                          device="cpu", dtype=torch.float32):
    """Generate paired exogenous context, process noise, and causal probes.

    Model variants receive this same object, which makes all environmental and
    perturbation randomness exactly paired. The first context component is a
    constant bias; the others follow independent AR(1) dynamics within each
    episode and reset at episode boundaries.
    """
    cfg.validate()
    days = cfg.days if days is None else days
    episodes = cfg.episodes_per_day if episodes is None else episodes
    generator = torch.Generator(device="cpu").manual_seed(task_seed)
    shape = (days, episodes, cfg.steps_per_episode)
    innovations = torch.randn(
        *shape, cfg.context_dim - 1, generator=generator, dtype=dtype)
    context = torch.empty(*shape, cfg.context_dim, dtype=dtype)
    context[..., 0] = 1.0
    state = torch.zeros(days, episodes, cfg.context_dim - 1, dtype=dtype)
    innovation_scale = (1.0 - cfg.context_ar ** 2) ** 0.5
    for step in range(cfg.steps_per_episode):
        state = cfg.context_ar * state + innovation_scale * innovations[:, :, step]
        context[:, :, step, 1:] = state
    noise = cfg.process_noise * torch.randn(
        *shape, cfg.n_neurons, generator=generator, dtype=dtype)
    perturbations = torch.empty(
        *shape, cfg.n_neurons, dtype=dtype).bernoulli_(0.5, generator=generator)
    perturbations.mul_(2).sub_(1)
    return BCITrajectories(
        context=context.to(device), process_noise=noise.to(device),
        perturbations=perturbations.to(device), task_seed=task_seed)


class BCISDIL:
    """Manual local learner for the continuous synthetic BCI."""

    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, cfg.feedback_dim, generator=generator, dtype=dtype)
                  / cfg.feedback_dim ** 0.5).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)
        self.model_seed = model_seed

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

    def feedback_features(self, error, previous_abs_error):
        if self.cfg.feedback == "error":
            return error.unsqueeze(1)
        improvement_velocity = (torch.zeros_like(error) if previous_abs_error is None
                                else previous_abs_error - error.abs())
        if self.cfg.feedback == "performance_velocity":
            return improvement_velocity.unsqueeze(1)
        return torch.stack((error, improvement_velocity), dim=1)

    def apical_components(self, soma, feedback):
        ordinary = self.coupling * soma
        instructional = feedback @ self.A.t()
        raw = ordinary + instructional
        baseline = self.P * soma + self.P_bias
        return raw, raw - baseline, ordinary, instructional

    def causal_targets(self, soma, target, perturbations):
        """One simultaneous antithetic forward-only causal target per cell."""
        sigma = self.cfg.perturb_sigma
        plus = soma + sigma * perturbations
        minus = soma - sigma * perturbations
        error_plus = target - plus @ self.role
        error_minus = target - minus @ self.role
        loss_plus = 0.5 * error_plus.square()
        loss_minus = 0.5 * error_minus.square()
        directional_derivative = (loss_plus - loss_minus) / (2.0 * sigma)
        return -directional_derivative.unsqueeze(1) * perturbations

    def causal_role_targets(self, cursor_plus, cursor_minus, perturbations):
        """Estimate each cell's signed causal effect on the BCI cursor.

        Unlike ``causal_targets``, this target contains no instantaneous error
        magnitude.  A scalar antithetic cursor difference is tagged by the
        locally available perturbation at each cell; its expectation is the
        experimenter-unknown causal role vector.  The learner receives the two
        scalar cursor observations and cannot read the environment's role map.
        """
        sigma = self.cfg.perturb_sigma
        directional_derivative = (
            cursor_plus - cursor_minus) / (2.0 * sigma)
        return directional_derivative.unsqueeze(1) * perturbations

    def exact_causal_direction(self, soma, target):
        """Analytic diagnostic only; never used by a learning update."""
        error = target - soma @ self.role
        return error.unsqueeze(1) * self.role.unsqueeze(0)

    def update_predictor(self, soma, ordinary, active):
        """Per-cell centered normalized-LMS fit of neutral apical traffic."""
        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_vectorizer(self, feedback, innovation, causal_target, active):
        if not bool(active.any()):
            return
        if self.cfg.feedback == "performance_velocity":
            # Role identification and performance modulation are deliberately
            # separated.  The former is an amortized node-perturbation estimate
            # of dz/dh_i; the latter enters apical activity only through the
            # within-episode performance innovation in ``feedback``.
            target = causal_target[active].mean(0)
            self.A[:, 0] += self.cfg.vectorizer_eta * (
                target - self.A[:, 0])
            return
        c = feedback[active]
        calibration_error = causal_target[active] - innovation[active]
        self.A += self.cfg.vectorizer_eta * (
            calibration_error.t() @ c / c.shape[0])

    def update_forward(self, context, soma, innovation, active, plasticity_gain):
        if plasticity_gain == 0 or not bool(active.any()):
            return
        eligibility_gain = 1.0 - soma.square()
        local_delta = plasticity_gain * innovation * eligibility_gain
        self.W += self.cfg.forward_eta * (
            local_delta[active].t() @ context[active] / active.sum())

    def step(self, context, process_noise, perturbations, previous_soma,
             previous_abs_error, active, global_step, *, control_gain=1.0,
             plasticity_gain=1.0, learn_vectorizer=True,
             learn_predictor=True):
        """Advance one contiguous BCI step and optionally apply local updates.

        ``control_gain`` affects state but not the plasticity signal.
        ``plasticity_gain`` affects W but not online state. This separation is
        the preregistered phase lesion. All tensors are minibatches of episodes.
        """
        with torch.no_grad():
            u = context @ self.W.t() + self.cfg.inertia * previous_soma + process_noise
            base_soma = torch.tanh(u)
            base_cursor = base_soma @ self.role
            base_error = self.cfg.target - base_cursor
            feedback = self.feedback_features(base_error, previous_abs_error)
            raw, innovation, ordinary, instructional = self.apical_components(
                base_soma, feedback)
            innovation = innovation * active.unsqueeze(1)
            controlled = torch.clamp(
                base_soma + control_gain * self.cfg.kappa * innovation, -1.0, 1.0)
            soma = torch.where(active.unsqueeze(1), controlled, previous_soma)
            cursor = soma @ self.role
            error = self.cfg.target - cursor
            loss = 0.5 * error.square()

            did_perturb = (learn_vectorizer
                           and global_step % self.cfg.perturb_every == 0)
            causal_target = None
            if did_perturb:
                if self.cfg.feedback == "performance_velocity":
                    sigma = self.cfg.perturb_sigma
                    # These two scalar values are environment observations.
                    # ``causal_role_targets`` has no access to ``self.role``.
                    cursor_plus = (
                        base_soma + sigma * perturbations) @ self.role
                    cursor_minus = (
                        base_soma - sigma * perturbations) @ self.role
                    causal_target = self.causal_role_targets(
                        cursor_plus, cursor_minus, perturbations)
                else:
                    causal_target = self.causal_targets(
                        base_soma, self.cfg.target, perturbations)

            # All updates see the same pre-update state. The order below only
            # mutates parameters after every required local quantity is stored.
            self.update_forward(
                context, base_soma, innovation, active, plasticity_gain)
            if did_perturb and learn_vectorizer:
                self.update_vectorizer(
                    feedback, innovation, causal_target, active)
            if learn_predictor:
                self.update_predictor(base_soma, ordinary, active)

        return {
            "soma": soma,
            "base_soma": base_soma,
            "cursor": cursor,
            "error": error,
            "abs_error": error.abs(),
            "loss": loss,
            "feedback": feedback,
            "raw_apical": raw,
            "innovation": innovation,
            "ordinary_apical": ordinary,
            "instructional_apical": instructional,
            "causal_target": causal_target,
            "did_perturb": did_perturb,
        }


def run_day(model, trajectories, day, *, control_gain=1.0,
            plasticity_gain=1.0, learn_vectorizer=True,
            learn_predictor=True, collect=False, global_step=0):
    """Run one parallel batch of episodes and return success plus event data."""
    cfg = model.cfg
    episodes = trajectories.context.shape[1]
    soma = torch.zeros(
        episodes, cfg.n_neurons, device=model.device, dtype=model.dtype)
    previous_abs_error = None
    active = torch.ones(episodes, device=model.device, dtype=torch.bool)
    success = torch.zeros_like(active)
    events = []
    for step_index in range(cfg.steps_per_episode):
        previous_soma_for_record = soma
        previous_error_for_record = (torch.full(
            (episodes,), abs(cfg.target), device=model.device, dtype=model.dtype)
                                     if previous_abs_error is None
                                     else previous_abs_error)
        record = model.step(
            trajectories.context[day, :, step_index],
            trajectories.process_noise[day, :, step_index],
            trajectories.perturbations[day, :, step_index],
            soma, previous_abs_error, active, global_step,
            control_gain=control_gain, plasticity_gain=plasticity_gain,
            learn_vectorizer=learn_vectorizer,
            learn_predictor=learn_predictor)
        newly_successful = active & (record["cursor"] >= cfg.target)
        success |= newly_successful
        if collect:
            events.append({
                key: value.detach().cpu().clone()
                for key, value in record.items()
                if isinstance(value, torch.Tensor)
            } | {
                "active": active.detach().cpu().clone(),
                "success": success.detach().cpu().clone(),
                "previous_soma": previous_soma_for_record.detach().cpu().clone(),
                "previous_abs_error": previous_error_for_record.detach().cpu().clone(),
                "step": step_index,
            })
        soma = record["soma"]
        previous_abs_error = record["abs_error"]
        active = active & ~newly_successful
        global_step += 1
    return {
        "success": success.detach().cpu(),
        "success_rate": success.float().mean().item(),
        "events": events,
        "global_step": global_step,
    }