"""Differential readout: delta_out = U (r(S_final, z) - r(S0, z)); U initialised at zero."""
from __future__ import annotations

import math

import torch
import torch.nn as nn

from .config import PrismeConfig


class DifferentialReadout(nn.Module):
    def __init__(self, cfg: PrismeConfig):
        super().__init__()
        self.cfg = cfg
        self.q_proj = nn.Linear(cfg.w, cfg.w, bias=False)
        self.k_proj = nn.Linear(cfg.w, cfg.w, bias=False)
        self.U = nn.Linear(cfg.w, cfg.d_model, bias=False)
        if cfg.readout_zero_init:
            nn.init.zeros_(self.U.weight)
        else:
            nn.init.normal_(self.U.weight, std=cfg.readout_init_std)

    def read(self, S: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        """Attention read over slots with a query that depends only on z. S: [...,M,w]; z: [...,w] -> [...,w]."""
        q = self.q_proj(z).unsqueeze(-2)                                    # [...,1,w]
        k = self.k_proj(S)                                                  # [...,M,w]
        att = torch.softmax((q * k).sum(-1) / math.sqrt(self.cfg.w), dim=-1)  # [...,M]
        return (att.unsqueeze(-1) * S).sum(-2)

    def forward(self, S_final: torch.Tensor, S0: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        return self.U(self.read(S_final, z) - self.read(S0, z)) * self.cfg.readout_scale