import torch
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def local_mean(x, kernel_size=5):
    kernel = torch.ones(1,1,kernel_size, device=device) / kernel_size
    return F.conv1d(x, kernel, padding=kernel_size//2)

def functional_batch(x, alpha=1.0, beta=1.0, gamma=1.0):
    mean = local_mean(x)

    G = 0.5 * (x)**2
    H = (x - mean)**2
    I = (torch.gradient(x, dim=2)[0])**2

    Phi = gamma*G - alpha*H + beta*I

    return Phi.mean(), {
        "G": G.mean(dim=[1,2]),
        "H": H.mean(dim=[1,2]),
        "I": I.mean(dim=[1,2])
    }

def evolve_batch(x, steps=100, lr=0.01, k=0.5, sigma=0.01):
    history = []

    for t in range(steps):
        Phi, metrics = functional_batch(x)

        Phi.backward()

        with torch.no_grad():
            E = torch.zeros_like(x)  # pode plugar SIH aqui

            x -= lr * x.grad
            x += k * (E - x)
            x += sigma * torch.randn_like(x)

            x.grad.zero_()

        history.append(metrics)

    return x, history
