"""Independent DeepONet and exponential vRBA implementation from the paper."""
import math
import torch
from torch import nn


class VRBADeepONet(nn.Module):
    def __init__(self, sensors=200, width=100, hidden_layers=4, latent_dim=100):
        super().__init__()
        def network(inputs):
            layers = []
            for _ in range(hidden_layers):
                layers.extend([nn.Linear(inputs, width), nn.GELU()])
                inputs = width
            layers.append(nn.Linear(width, latent_dim))
            return nn.Sequential(*layers)
        self.branch = network(sensors)
        self.trunk = network(1)

    def forward(self, pressure, time):
        """Normalized pressure [B,M], normalized time [N,1] -> radius [B,N]."""
        return self.branch(pressure) @ self.trunk(time).transpose(0, 1)


class ExponentialVRBA:
    """Eqs. 19, 40--44: detached EMA weights and function marginal sampling.

    phi mixes the max-normalized exponential tilt with a uniform floor.
    The exponential divided by its maximum is evaluated as exp((r-rmax)/eps).
    """
    def __init__(self, functions, points, device, eta=0.1, gamma=0.99,
                 phi=0.9, c=1.0, initial_weight=1.0):
        self.weights = torch.full((functions, points), initial_weight, device=device)
        self.eta, self.gamma, self.phi, self.c = eta, gamma, phi, c

    @torch.no_grad()
    def probabilities(self):
        scores = self.weights.sum(dim=1)
        return scores / scores.sum()

    @torch.no_grad()
    def update(self, indices, residual, step):
        r = residual.detach().abs()
        maximum = r.amax(dim=1, keepdim=True)
        epsilon = self.c * maximum / math.log(step + 2.0)
        # Zero residual -> uniform tilt. Numerical floor does not add temperature.
        safe = epsilon.clamp_min(torch.finfo(r.dtype).tiny)
        tilt = torch.exp((r - maximum) / safe)
        # Eq. 14 with q=tilt/sum(tilt), p_u=1/N and eta*=eta/max(q).
        # Thus p_u/max(q)=mean(tilt); using a constant 1 here changes Eq. 14.
        tilt = self.phi * tilt + (1.0 - self.phi) * tilt.mean(dim=1, keepdim=True)
        updated = self.gamma * self.weights[indices] + self.eta * tilt
        self.weights[indices] = updated
        return updated