import torch
import torch.nn as nn
import torch.optim as optim
import torch.autograd as ag

from tqdm import tqdm

import numpy as np
from matplotlib import pyplot as plt

# y'+y=x, y(0)=1
# Analytical solution: y(x) = -1 + 2 * exp(-x) + x

x_min = 0
x_max = 2
xs = torch.linspace(x_min, x_max, 100_000).reshape(-1, 1)

model = nn.Sequential(
    nn.Linear(1, 256),
    nn.Sigmoid(),
    nn.Linear(256, 1),
)

x_ds = torch.utils.data.TensorDataset(xs)
x_dl = torch.utils.data.DataLoader(x_ds, batch_size=128, shuffle=True)
zero = torch.tensor([0.0]).reshape(1, 1)

losses = []
opt = optim.Adam(model.parameters(), lr=1e-3)
for x in (bar := tqdm(x_dl)):
    x = x[0].requires_grad_(True)
    y = model(x)
    dy_dx = torch.autograd.grad(
        outputs=y,
        inputs=x,
        grad_outputs=torch.ones_like(x),
        create_graph=True,
    )[0]

    y_0 = model(zero)

    l = (dy_dx + y - x)**2 + (y_0 - 1.0)**2
    l = l.mean()

    bar.set_description(f'Loss: {l.item():.2f}')
    losses.append(l.item())

    opt.zero_grad()
    l.backward()
    opt.step()

xs = torch.linspace(x_min - 1.5, x_max + 1.5, 1000).reshape(-1, 1)
ys = model(xs).detach().numpy()
xs = xs.detach().numpy()
ys_true = -1 + 2 * np.exp(-xs) + xs

plt.plot(xs, ys, label='NN')
plt.plot(xs, ys_true, label='Analytical')
plt.plot([x_min, x_min], [0, 1], color='black')
plt.plot([x_max, x_max], [0, 1], color='black')
plt.legend()
plt.show()
plt.plot(losses)