"""Shared inference-and-teaching pathway for BabyAI imitation learning.

The visual observation is basal.  The language mission is projected into each
hidden population through a fixed apical path.  Manual updates implement exact
BP, clean reciprocal KP, raw shared-path KP, or SDIL residualization.
"""

from dataclasses import asdict, dataclass
import math
import re

import numpy as np
import torch
import torch.nn.functional as F


CONDITIONS = ("bp", "clean_kp", "raw_shared", "sdil")
TOKEN_PATTERN = re.compile(r"[a-z0-9]+")


def tokenize_mission(mission):
    return TOKEN_PATTERN.findall(str(mission).lower())


def build_vocabulary(missions):
    tokens = sorted({token for mission in missions
                     for token in tokenize_mission(mission)})
    return ("<unk>", *tokens)


def missions_to_bow(missions, vocabulary):
    lookup = {token: index for index, token in enumerate(vocabulary)}
    result = np.zeros((len(missions), len(vocabulary)), dtype=np.float32)
    for row, mission in enumerate(missions):
        tokens = tokenize_mission(mission)
        for token in tokens:
            result[row, lookup.get(token, 0)] += 1.0
        norm = np.linalg.norm(result[row])
        if norm > 0:
            result[row] /= norm
    return result


@dataclass(frozen=True)
class BabyAISharedConfig:
    input_dim: int
    mission_dim: int
    action_dim: int = 7
    width: int = 256
    hidden_layers: int = 2
    learning_rate: float = 0.03
    reciprocal_learning_rate: float = 0.03
    momentum: float = 0.9
    weight_decay: float = 1e-4
    context_gain: float = 1.0

    def to_dict(self):
        return asdict(self)


class BabyAISharedNet:
    """Tanh policy with fixed mission projections and local manual updates."""

    def __init__(self, config, seed, device="cpu", dtype=torch.float32):
        if config.hidden_layers < 1:
            raise ValueError("BabyAI shared policy needs a hidden population")
        self.config = config
        self.device = torch.device(device)
        self.dtype = dtype
        sizes = ([config.input_dim]
                 + [config.width] * config.hidden_layers
                 + [config.action_dim])
        forward_generator = torch.Generator(device="cpu").manual_seed(seed)
        reciprocal_generator = torch.Generator(device="cpu").manual_seed(seed + 1)
        context_generator = torch.Generator(device="cpu").manual_seed(seed + 2)

        def normal(shape, scale, generator):
            value = torch.randn(*shape, generator=generator) * scale
            return value.to(device=self.device, dtype=dtype)

        self.W = [normal((sizes[i + 1], sizes[i]),
                         1.0 / math.sqrt(sizes[i]), forward_generator)
                  for i in range(len(sizes) - 1)]
        self.b = [torch.zeros(size, device=self.device, dtype=dtype)
                  for size in sizes[1:]]
        # Q[i] has the same storage orientation as W[i].  Q[0] is not needed
        # because no teaching signal is transported into the visual input.
        self.Q = [None] + [normal(tuple(self.W[i].shape),
                                  1.0 / math.sqrt(sizes[i]),
                                  reciprocal_generator)
                           for i in range(1, len(self.W))]
        self.C = [normal((config.width, config.mission_dim),
                         config.context_gain
                         / math.sqrt(config.mission_dim),
                         context_generator)
                  for _ in range(config.hidden_layers)]

        self.P = [torch.zeros(config.width, device=self.device, dtype=dtype)
                  for _ in range(config.hidden_layers)]
        self.P_bias = [torch.zeros_like(value) for value in self.P]
        self.mW = [torch.zeros_like(value) for value in self.W]
        self.mb = [torch.zeros_like(value) for value in self.b]
        self.mQ = [None] + [torch.zeros_like(value) for value in self.Q[1:]]

    def clone(self):
        copied = BabyAISharedNet(
            self.config, seed=0, device=self.device, dtype=self.dtype)
        for name in ("W", "b", "Q", "C", "P", "P_bias", "mW", "mb", "mQ"):
            values = getattr(self, name)
            setattr(copied, name, [
                None if value is None else value.clone() for value in values])
        return copied

    def context_fields(self, missions, enabled=True):
        if not enabled:
            return [torch.zeros(
                (missions.shape[0], self.config.width),
                device=self.device, dtype=self.dtype) for _ in self.C]
        return [missions @ projection.t() for projection in self.C]

    def forward_features(self, features, missions, context_enabled=True):
        contexts = self.context_fields(missions, enabled=context_enabled)
        h = [features]
        u = []
        for layer in range(self.config.hidden_layers):
            value = (h[-1] @ self.W[layer].t() + self.b[layer]
                     + contexts[layer])
            u.append(value)
            h.append(torch.tanh(value))
        logits = h[-1] @ self.W[-1].t() + self.b[-1]
        h.append(logits)
        return {"h": h, "u": u, "context": contexts, "logits": logits}

    def predictor(self, layer, soma):
        return self.P[layer] * soma + self.P_bias[layer]

    @torch.no_grad()
    def fit_neutral_predictor(self, features, missions):
        """Fit the per-cell neutral relation without actions or teaching."""
        state = self.forward_features(features, missions)
        reports = []
        for layer, target in enumerate(state["context"]):
            soma = state["h"][layer + 1]
            soma_centered = soma - soma.mean(0)
            target_centered = target - target.mean(0)
            variance = soma_centered.square().mean(0)
            covariance = (soma_centered * target_centered).mean(0)
            slope = covariance / variance.clamp_min(1e-8)
            intercept = target.mean(0) - slope * soma.mean(0)
            self.P[layer].copy_(slope)
            self.P_bias[layer].copy_(intercept)
            prediction = slope * soma + intercept
            residual = target - prediction
            target_ss = target_centered.square().sum(0)
            residual_ss = residual.square().sum(0)
            valid = target_ss > 1e-12
            r2 = 1.0 - residual_ss[valid] / target_ss[valid]
            context_rms = target.square().mean().sqrt()
            reports.append({
                "mean_per_cell_r2": float(r2.mean()) if r2.numel() else 0.0,
                "context_rms": float(context_rms),
                "residual_context_rms_ratio": float(
                    residual.square().mean().sqrt()
                    / context_rms.clamp_min(1e-12)),
                "neutral_observations": int(features.shape[0]),
                "action_observations": 0,
                "teaching_observations": 0,
            })
        return reports


@torch.no_grad()
def fit_population_predictor(soma, target, ridge=1e-3):
    """Fit target ~= soma @ coefficient + intercept for a capacity audit."""
    soma_mean = soma.mean(0)
    target_mean = target.mean(0)
    centered_soma = soma - soma_mean
    centered_target = target - target_mean
    gram = centered_soma.t() @ centered_soma / soma.shape[0]
    cross = centered_soma.t() @ centered_target / soma.shape[0]
    scale = gram.diagonal().mean().clamp_min(1e-8)
    regularized = gram + ridge * scale * torch.eye(
        gram.shape[0], device=gram.device, dtype=gram.dtype)
    coefficient = torch.linalg.solve(regularized, cross)
    intercept = target_mean - soma_mean @ coefficient
    return coefficient, intercept


@torch.no_grad()
def population_predictor_metrics(soma, target, coefficient, intercept):
    prediction = soma @ coefficient + intercept
    residual = target - prediction
    centered_target = target - target.mean(0)
    target_ss = centered_target.square().sum(0)
    residual_ss = residual.square().sum(0)
    valid = target_ss > 1e-12
    r2 = 1.0 - residual_ss[valid] / target_ss[valid]
    context_rms = target.square().mean().sqrt()
    return {
        "mean_per_cell_r2": float(r2.mean()) if r2.numel() else 0.0,
        "context_rms": float(context_rms),
        "residual_context_rms_ratio": float(
            residual.square().mean().sqrt()
            / context_rms.clamp_min(1e-12)),
        "observations": int(soma.shape[0]),
    }


def visual_input_dim(object_cardinality, color_cardinality,
                     state_cardinality, view_size=7):
    return (view_size * view_size
            * (object_cardinality + color_cardinality + state_cardinality) + 4)


def history_input_dim(object_cardinality, color_cardinality,
                      state_cardinality, history_steps, action_dim=7,
                      view_size=7):
    if history_steps < 1:
        raise ValueError("history_steps must be positive")
    frame_dim = visual_input_dim(
        object_cardinality, color_cardinality, state_cardinality, view_size)
    return history_steps * (frame_dim + action_dim)


def encode_visual(images, directions, object_cardinality,
                  color_cardinality, state_cardinality, device,
                  dtype=torch.float32):
    """One-hot encode the compact MiniGrid symbolic observation."""
    images = torch.as_tensor(images, device=device, dtype=torch.long)
    directions = torch.as_tensor(directions, device=device, dtype=torch.long)
    object_features = F.one_hot(
        images[..., 0], num_classes=object_cardinality)
    color_features = F.one_hot(
        images[..., 1], num_classes=color_cardinality)
    state_features = F.one_hot(
        images[..., 2], num_classes=state_cardinality)
    grid = torch.cat(
        (object_features, color_features, state_features), dim=-1)
    grid = grid.reshape(grid.shape[0], -1).to(dtype)
    direction = F.one_hot(directions, num_classes=4).to(dtype)
    return torch.cat((grid, direction), dim=1)


def build_history_index(episode_offsets, actions, history_steps):
    """Build right-aligned history indices without crossing episodes."""
    if history_steps < 1:
        raise ValueError("history_steps must be positive")
    episode_offsets = np.asarray(episode_offsets, dtype=np.int64)
    actions = np.asarray(actions, dtype=np.int64)
    total = len(actions)
    episode_start = np.empty(total, dtype=np.int64)
    for start, stop in zip(episode_offsets[:-1], episode_offsets[1:]):
        episode_start[start:stop] = start
    positions = np.arange(total, dtype=np.int64)
    indices = np.zeros((total, history_steps), dtype=np.int64)
    previous_actions = np.full((total, history_steps), -1, dtype=np.int64)
    mask = np.zeros((total, history_steps), dtype=bool)
    for lag in range(history_steps):
        column = history_steps - 1 - lag
        source = positions - lag
        valid = source >= episode_start
        indices[valid, column] = source[valid]
        mask[valid, column] = True
        has_previous = valid & (source > episode_start)
        previous_actions[has_previous, column] = actions[
            source[has_previous] - 1]
    return indices, previous_actions, mask


def encode_history_visual(images, directions, previous_actions, mask,
                          object_cardinality, color_cardinality,
                          state_cardinality, action_dim, device,
                          dtype=torch.float32):
    """Encode a padded sequence of symbolic frames and preceding actions."""
    images = np.asarray(images)
    directions = np.asarray(directions)
    previous_actions = np.asarray(previous_actions)
    mask = np.asarray(mask)
    if images.ndim != 5:
        raise ValueError("history images must have shape [batch, time, H, W, 3]")
    batch, history_steps = images.shape[:2]
    frames = encode_visual(
        images.reshape(batch * history_steps, *images.shape[2:]),
        directions.reshape(batch * history_steps), object_cardinality,
        color_cardinality, state_cardinality, device=device, dtype=dtype)
    frame_dim = frames.shape[1]
    mask_tensor = torch.as_tensor(mask, device=device, dtype=dtype)
    frames = frames.reshape(batch, history_steps, frame_dim)
    frames = frames * mask_tensor[:, :, None]
    action_values = torch.as_tensor(
        np.maximum(previous_actions, 0), device=device, dtype=torch.long)
    action_valid = torch.as_tensor(
        (previous_actions >= 0) & mask, device=device, dtype=dtype)
    action_features = F.one_hot(
        action_values, num_classes=action_dim).to(dtype)
    action_features = action_features * action_valid[:, :, None]
    return torch.cat((frames, action_features), dim=2).reshape(batch, -1)


def select_teaching_signal(net, layer, instruction, context, soma, condition):
    raw = instruction + context
    innovation = raw - net.predictor(layer, soma)
    if condition == "clean_kp":
        used = instruction
    elif condition == "raw_shared":
        used = raw
    elif condition == "sdil":
        used = innovation
    else:
        raise ValueError(f"no shared teaching rule for condition {condition}")
    return used, raw, innovation


@torch.no_grad()
def manual_step(net, features, missions, actions, condition):
    """Apply one exact-BP or reciprocal local-learning update."""
    if condition not in CONDITIONS:
        raise ValueError(f"unknown BabyAI condition: {condition}")
    state = net.forward_features(features, missions)
    h, u, contexts = state["h"], state["u"], state["context"]
    probabilities = torch.softmax(state["logits"], dim=1)
    output_instruction = (
        F.one_hot(actions, num_classes=net.config.action_dim).to(net.dtype)
        - probabilities)
    batch = features.shape[0]

    hidden_deltas = [None] * net.config.hidden_layers
    used_fields = [None] * net.config.hidden_layers
    raw_fields = [None] * net.config.hidden_layers
    innovation_fields = [None] * net.config.hidden_layers
    child_delta = output_instruction
    for layer in reversed(range(net.config.hidden_layers)):
        if condition == "bp":
            used = child_delta @ net.W[layer + 1]
            raw = used + contexts[layer]
            innovation = raw - net.predictor(layer, h[layer + 1])
        else:
            instruction = child_delta @ net.Q[layer + 1]
            used, raw, innovation = select_teaching_signal(
                net, layer, instruction, contexts[layer], h[layer + 1],
                condition)
        delta = used * (1.0 - torch.tanh(u[layer]).square())
        hidden_deltas[layer] = delta
        used_fields[layer] = used
        raw_fields[layer] = raw
        innovation_fields[layer] = innovation
        child_delta = delta

    directions = [hidden_deltas[0].t() @ h[0] / batch]
    for layer in range(1, net.config.hidden_layers):
        directions.append(hidden_deltas[layer].t() @ h[layer] / batch)
    directions.append(output_instruction.t() @ h[-2] / batch)
    bias_directions = [delta.mean(0) for delta in hidden_deltas]
    bias_directions.append(output_instruction.mean(0))

    for layer, (direction, bias_direction) in enumerate(
            zip(directions, bias_directions)):
        net.mW[layer].mul_(net.config.momentum).add_(
            direction - net.config.weight_decay * net.W[layer])
        net.mb[layer].mul_(net.config.momentum).add_(bias_direction)
        net.W[layer].add_(net.mW[layer], alpha=net.config.learning_rate)
        net.b[layer].add_(net.mb[layer], alpha=net.config.learning_rate)
        if condition != "bp" and layer > 0:
            net.mQ[layer].mul_(net.config.momentum).add_(
                direction - net.config.weight_decay * net.Q[layer])
            net.Q[layer].add_(
                net.mQ[layer], alpha=net.config.reciprocal_learning_rate)

    loss = F.cross_entropy(state["logits"], actions)
    return float(loss), {
        "directions": directions,
        "bias_directions": bias_directions,
        "used": used_fields,
        "raw": raw_fields,
        "innovation": innovation_fields,
        "context": contexts,
    }


@torch.no_grad()
def evaluate_actions(net, images, directions, missions, actions,
                     cardinalities, batch_size=1024, context_enabled=True,
                     history=None):
    device = net.device
    total_loss = 0.0
    correct = 0
    total = len(actions)
    for start in range(0, total, batch_size):
        stop = min(start + batch_size, total)
        if history is None:
            features = encode_visual(
                images[start:stop], directions[start:stop], *cardinalities,
                device=device, dtype=net.dtype)
        else:
            indices, previous_actions, mask = history
            batch_indices = indices[start:stop]
            features = encode_history_visual(
                images[batch_indices], directions[batch_indices],
                previous_actions[start:stop], mask[start:stop],
                *cardinalities, net.config.action_dim,
                device=device, dtype=net.dtype)
        mission_batch = torch.as_tensor(
            missions[start:stop], device=device, dtype=net.dtype)
        action_batch = torch.as_tensor(
            actions[start:stop], device=device, dtype=torch.long)
        logits = net.forward_features(
            features, mission_batch,
            context_enabled=context_enabled)["logits"]
        total_loss += float(F.cross_entropy(
            logits, action_batch, reduction="sum"))
        correct += int((logits.argmax(1) == action_batch).sum())
    return {"accuracy": correct / total, "loss": total_loss / total}