
# Importing Libraries

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from PINN import NeuralNet, init_xavier

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dtype = torch.float32

# Define the forcing function f(x, y, t) generalized to 5D
def forcing_function(X):
    # Assume X has shape (N, 5): [x1, x2, x3, x4, t]
    x1 = X[:, 0:1]
    x2 = X[:, 1:2]  # treat as y
    x3 = X[:, 2:3]
    x4 = X[:, 3:4]
    t = X[:, 4:5]
    
    pi = torch.acos(torch.zeros(1)).item() * 2  # pi = 2 * arccos(0)
    
    sin_pix = torch.sin(pi * x1)
    y = x2
    exp_term = torch.exp(-t)
    
    f = 5 * exp_term * sin_pix * y**(-3) * (-1 + exp_term * sin_pix * y**(-5) * (-12 + pi**2 * y**2))
    return f

# Define the PDE residual: u_t - Δu - f(x, y, t)
def pde_residual(model, X):
    X.requires_grad_(True)
    u = model(X)

    grads = torch.autograd.grad(u, X, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_t = grads[:, 4:5]

    # Compute second derivatives for Laplacian
    laplacian = 0
    for i in range(4):  # x1, x2, x3, x4
        grad_i = torch.autograd.grad(grads[:, i:i+1], X, grad_outputs=torch.ones_like(grads[:, i:i+1]), create_graph=True)[0][:, i:i+1]
        laplacian += grad_i

    f_val = forcing_function(X)
    residual = u_t - laplacian - f_val
    return residual

# You can now use `pde_residual(model, X)` during training.
