import torch
from torch import nn


def _grad(y, value):
    return torch.autograd.grad(value.sum(), y, create_graph=True, retain_graph=True)[0]


class SymplecticBlock(nn.Module):
    """Learned separable Hamiltonian block using velocity-Verlet/leapfrog."""
    def __init__(self, width, input_dim, substeps=2, hidden=32, step=0.1):
        super().__init__()
        self.width, self.substeps, self.step = width, substeps, step
        self.potential = nn.Sequential(nn.Linear(width + input_dim, hidden), nn.Tanh(),
                                       nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
        self.log_mass = nn.Parameter(torch.zeros(width))

    def hamiltonian(self, q, p, x):
        u = self.potential(torch.cat([q, x], dim=-1)).squeeze(-1)
        mass = torch.nn.functional.softplus(self.log_mass) + 1e-3
        kinetic = 0.5 * (p * p / mass).sum(-1)
        return u + kinetic

    def forward(self, q, p, x):
        # x stays fixed over all internal substeps, as in the proposed block.
        for _ in range(self.substeps):
            q = q.requires_grad_(True); p = p.requires_grad_(True)
            h = self.hamiltonian(q, p, x)
            gq, gp = _grad(q, h), _grad(p, h)
            p = p - 0.5 * self.step * gq
            q = q + self.step * gp
            q = q.requires_grad_(True); p = p.requires_grad_(True)
            h = self.hamiltonian(q, p, x)
            gq = _grad(q, h)
            p = p - 0.5 * self.step * gq
        return q, p


def leapfrog(q, p, step, n=1):
    """Analytic separable H=(q^2+p^2)/2 update for cheap verification."""
    for _ in range(n):
        p = p - 0.5 * step * q
        q = q + step * p
        p = p - 0.5 * step * q
    return q, p


def euler(q, p, step, n=1):
    for _ in range(n):
        q, p = q + step * p, p - step * q
    return q, p