import torch

import torchdeq
from torchdeq import get_deq
from torchdeq.norm import apply_norm, reset_norm

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer = torch.nn.Linear(10, 10)

        # deq
        self.deq = get_deq()
        apply_norm(self.layer, 'weight_norm')

    def implicit_layer(self, x):
        return self.layer(x)
    
    def forward(self, x, pos):

        z = torch.zeros_like(x)

        reset_norm(self.layer)

        f = lambda z: self.f(z, pos)

        z_pred, info = self.deq(self.implicit_layer, z)
        
        # if model.eval() -> z_pred[-1].requires_grad is False!
        energy = z_pred[-1]
        forces = -1 * (
            torch.autograd.grad(
                energy,
                # diff with respect to pos
                # if you get 'One of the differentiated Tensors appears to not have been used in the graph'
                # then because pos is not 'used' to calculate the energy
                pos, 
                grad_outputs=torch.ones_like(energy),
                create_graph=True,
                # allow_unused=True, 
            )[0]
        )

        return energy, forces


def run(model, eval=False):

    if eval:
        model.eval()
    else:
        model.train()

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for step in range(10):
        x = torch.randn(10, 10)
        pos = torch.randn(10, 3)
        energy, forces = model(x, pos)
        
        # loss
        optimizer.zero_grad()
        energy_target = torch.randn(10, 1)
        energy_loss = torch.nn.functional.mse_loss(energy, energy_target)
        force_target = torch.randn(10, 3)
        force_loss = torch.nn.functional.mse_loss(forces, force_target)
        loss = energy_loss + force_loss

        if not eval:
            loss.backward()
            optimizer.step()
    
    return True

if __name__ == '__main__':
    model = MyModel()
    success = run(model, eval=False)
    print(f'train success: {success}')
    success = run(model, eval=True)
    print(f'eval success: {success}')

import torch
import torch.nn as nn

from torchdeq import get_deq
from torchdeq.norm import apply_norm, reset_norm
import os


import matplotlib.pyplot as plt
import numpy as np


import torch
import torch.nn as nn
import torch.nn.functional as F


class FCNN(nn.Module):
    def __init__(self, input_dim, neurons_per_layer, output_dim):
        super(FCNN, self).__init__()

        self.linear1 = nn.Linear(input_dim, neurons_per_layer)
        self.linear2 = nn.Linear(neurons_per_layer,neurons_per_layer)
        self.out = nn.Linear(neurons_per_layer, output_dim)

    def forward(self, x, z):
        # z: links' flow  x: context

        zx = torch.cat((z, x))

        z_processed = F.relu(self.linear1(zx))
        z_processed = F.relu(self.linear2(z_processed))

        return F.relu(self.out(z_processed))


if __name__ == '__main__':

    seed = 1
    torch.manual_seed(seed)

    fcnn = FCNN(input_dim=15, neurons_per_layer=20, output_dim=10)
    print(fcnn)

    x_, z_ = torch.ones(5), torch.ones(10)  # z0
    for i in range(10):
        z_ = fcnn(x_, z_)
        print('z:', z_)

    # Let's try a multi-variable DEQ!
    deq = get_deq(f_solver='broyden', f_max_iter=20, f_tol=1e-6)

    x_, z_ = torch.ones(5), torch.zeros(10)  # z0
    # f = lambda z: fcnn(x, z)
    z_out, info = deq(fcnn, (x_, z_))