"""Detailed-balance graph transport primitives and a small learnable layer."""
import numpy as np
import torch
from torch import nn


def logarithmic_mean(a, b):
    a, b = np.asarray(a, float), np.asarray(b, float)
    aa, bb = np.broadcast_arrays(a, b)
    out = np.empty_like(aa)
    close = np.isclose(aa, bb, rtol=1e-10, atol=1e-14)
    out[close] = (aa[close] + bb[close]) / 2
    out[~close] = (aa[~close] - bb[~close]) / (np.log(aa[~close]) - np.log(bb[~close]))
    return out


def energy(rho, pi):
    rho, pi = np.asarray(rho), np.asarray(pi)
    return float(np.sum(rho * np.log(rho / pi)))


def rhs(rho, pi, conductance):
    """Master-equation RHS; rho/pi may be [N] or [N,K]."""
    q = np.asarray(rho) / np.asarray(pi)
    if q.ndim == 1:
        return np.sum(conductance * (q[None, :] - q[:, None]), axis=1)
    return np.sum(conductance[:, :, None] * (q[None, :, :] - q[:, None, :]), axis=1)


def explicit_step(rho, pi, conductance, dt):
    return np.asarray(rho) + dt * rhs(rho, pi, conductance)


def positivity_dt_bound(rho, pi, conductance):
    out = np.sum(conductance, axis=1)
    pi = np.asarray(pi); rho = np.asarray(rho)
    if pi.ndim == 1:
        out = out / pi
    else:
        out = out[:, None] / pi
    valid = out > 0
    return float(np.min(rho[valid] / out[valid]))


def dissipation(rho, pi, conductance):
    q = np.asarray(rho) / np.asarray(pi)
    mu = np.log(q)
    if q.ndim == 1:
        lm = logarithmic_mean(q[:, None], q[None, :])
        return float(.5 * np.sum(conductance * lm * (mu[:, None]-mu[None, :])**2))
    lm = logarithmic_mean(q[:, None, :], q[None, :, :])
    return float(.5 * np.sum(conductance[:, :, None] * lm * (mu[None, :, :] - mu[:, None, :])**2))


class DetailedBalanceTransport(nn.Module):
    """K-channel positive graph transport with conservative explicit updates."""
    def __init__(self, in_dim, channels=2, dt=0.1, eps=1e-6):
        super().__init__()
        self.rho_head = nn.Linear(in_dim, channels)
        self.pi_head = nn.Linear(in_dim, channels)
        self.edge_head = nn.Linear(3 * in_dim, channels)
        self.dt, self.eps = dt, eps

    def forward(self, h, edge_index):
        n = h.shape[0]; src, dst = edge_index
        rho = torch.nn.functional.softplus(self.rho_head(h)) + self.eps
        pi = torch.nn.functional.softplus(self.pi_head(h)) + self.eps
        pi = pi / pi.sum(0, keepdim=True) * rho.sum(0, keepdim=True).detach()
        z = torch.cat([h[src], h[dst], (h[src]-h[dst]).abs()], dim=1)
        c0 = torch.nn.functional.softplus(self.edge_head(z)) + self.eps
        key = torch.minimum(src, dst) * n + torch.maximum(src, dst)
        cs = torch.zeros((n*n, c0.shape[1]), device=h.device)
        counts = torch.zeros((n*n, 1), device=h.device)
        cs.index_add_(0, key, c0); counts.index_add_(0, key, torch.ones_like(c0[:, :1]))
        c = cs[key] / counts[key].clamp_min(1)
        delta = rho[src] / pi[src] - rho[dst] / pi[dst]
        dr = torch.zeros_like(rho)
        dr.index_add_(0, src, -c * delta); dr.index_add_(0, dst, c * delta)
        return rho + self.dt * dr, pi, c, rho