"""Small differentiable Hodge-dual electrostatic loss on a periodic 2-D grid."""
import torch


def periodic_grad(u, h):
    return torch.stack(((torch.roll(u, -1, -2)-torch.roll(u, 1, -2))/(2*h),
                        (torch.roll(u, -1, -1)-torch.roll(u, 1, -1))/(2*h)), dim=-1)


def curl_2d(A, h):
    """2-D scalar-vector-potential curl: (d_y A_z, -d_x A_z)."""
    az = A[..., 2]
    dy = (torch.roll(az, -1, -2)-torch.roll(az, 1, -2))/(2*h)
    dx = (torch.roll(az, -1, -1)-torch.roll(az, 1, -1))/(2*h)
    return torch.stack((dy, -dx), dim=-1)


def divergence_2d(p, h):
    px, py = p[..., 0], p[..., 1]
    dx = (torch.roll(px, -1, -1)-torch.roll(px, 1, -1))/(2*h)
    dy = (torch.roll(py, -1, -2)-torch.roll(py, 1, -2))/(2*h)
    return dx + dy


def dielectric(n, eps_perp=2.0, eps_a=0.5):
    n = n / (torch.linalg.vector_norm(n, dim=-1, keepdim=True) + 1e-8)
    I = torch.eye(3, device=n.device, dtype=n.dtype)
    return eps_perp*I + eps_a*n[..., :, None]*n[..., None, :]


def dual_electrostatic_loss(p0, A, n, h, eps_perp=2.0, eps_a=0.5):
    """Positive dual density plus optional elastic term is left to the caller."""
    p = p0 + curl_2d(A, h)
    E = dielectric(n, eps_perp, eps_a)
    density = 0.5 * torch.einsum('...i,...ij,...j->...', p, torch.linalg.inv(E), p)
    return density.mean(), p