# Importing Libraries

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils
import torch.utils.data
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from math import pi
from PINN import NeuralNet
from PINN import init_xavier
import time
import argparse
from sklearn.metrics import mean_squared_error
import csv
import os
import pylab as p
import matplotlib.cm as cm

device = torch.device('cpu')
dtype = torch.float32

# Load .npy file
train_interior = np.load('X_int_train.npy')
train_boundary = np.load('X_b_train.npy')
test_interior = np.load('X_test.npy')
test_boundary = np.load('X_b_test.npy')
test_tensor = np.concatenate((test_interior, test_boundary))

print("train_interior", train_interior.shape)
print("train_boundary", train_boundary.shape)
print("test_boundary", test_boundary.shape)
print("test_interior", test_interior.shape)

quit()
# Define the exact solution
def exact_solution(x, y, z):
    return 5 + torch.sin(pi*x)*torch.sin(pi*y)*torch.sin(4*pi*z)
    
# Define BC
def BC(x, y, z):
    return 5 + torch.sin(pi*x)*torch.sin(pi*y)*torch.sin(4*pi*z)
    
# Type of optimizer (ADAM or LBFGS)
opt_type = "LBFGS"


train_interior_tensor = torch.tensor(train_interior, dtype = dtype)
train_boundary_tensor = torch.tensor(train_boundary, dtype = dtype)
test_interior_tensor = torch.tensor(test_interior, dtype = dtype)
test_boundary_tensor = torch.tensor(test_boundary, dtype = dtype)
test_tensor = torch.tensor(test_tensor, dtype = dtype)

train_boundary_value_tensor = BC(train_boundary_tensor[:,0], train_boundary_tensor[:,1], train_boundary_tensor[:,2]).reshape(-1,1)

#################################################


training_set = DataLoader(torch.utils.data.TensorDataset(train_boundary_tensor.to(device), train_boundary_value_tensor.to(device)), batch_size=888, shuffle=False)
    
my_network = NeuralNet(input_dimension = 3, output_dimension = 1, n_hidden_layers=4, neurons=20)

# Random Seed for weight initialization
retrain = 1
# Xavier weight initialization
init_xavier(my_network, retrain)

if opt_type == "ADAM":
    optimizer_ = optim.Adam(my_network.parameters(), lr=0.001)
elif opt_type == "LBFGS":
    optimizer_ = optim.LBFGS(my_network.parameters(), lr=0.1, max_iter=1, max_eval=50000, tolerance_change=1.0 * np.finfo(float).eps)
else:
    raise ValueError("Optimizer not recognized")


def fit(model, training_set, interior, num_epochs, optimizer, p, verbose=True):
    history = list()

    # Loop over epochs
    for epoch in range(num_epochs):
        if verbose: print("################################ ", epoch, " ################################")

        running_loss = list([0])

        # Loop over batches
        for j, (bd, bd_val) in enumerate(training_set):
            def closure():
                # zero the parameter gradients
                optimizer.zero_grad()
                # for bottom boundary
                u_bd_pred_ = model(bd)

                # residual calculation
                interior.requires_grad = True
                u_hat = model(train_interior_tensor)
                inputs = torch.ones(6000, 1).to(device)
                grad_u_hat = torch.autograd.grad(u_hat, interior, grad_outputs=inputs, create_graph=True)[0]

                u_x = grad_u_hat[:, 0]
                u_y = grad_u_hat[:, 1]
                u_z = grad_u_hat[:, 2]

                grad_grad_u_x = \
                torch.autograd.grad(u_x, interior, grad_outputs=torch.ones(interior.shape[0]), create_graph=True)[0]
                u_xx = grad_grad_u_x[:, 0]

                grad_grad_u_y = \
                torch.autograd.grad(u_y, interior, grad_outputs=torch.ones(interior.shape[0]), create_graph=True)[0]
                u_yy = grad_grad_u_y[:, 1]
                
                grad_grad_u_z = \
                torch.autograd.grad(u_z, interior, grad_outputs=torch.ones(interior.shape[0]), create_graph=True)[0]
                u_zz = grad_grad_u_z[:, 2]


                # Item 1. below
                loss = (torch.mean((u_bd_pred_.reshape(-1, ) - bd_val.reshape(-1, )) ** p) +
                        0.1*torch.mean((u_xx.reshape(-1, ) + u_yy.reshape(-1, ) + u_zz.reshape(-1, ) - u_hat.reshape(-1, ) + 
                        19*pi**2*torch.sin(pi*train_interior_tensor[:,0])*torch.sin(pi*train_interior_tensor[:,1])*torch.sin(4*pi*train_interior_tensor[:,2]).reshape(-1, ) ) ** p))

                # Item 2. below
                loss.backward()
                # Compute average training loss over batches for the current epoch
                running_loss[0] += loss.item()
                return loss

            # Item 3. below
            optimizer.step(closure=closure)

        print('Loss: ', (running_loss[0] / len(training_set)))
        history.append(running_loss[0])

    return history

#n_epochs = 10000
#start_time = time.time()
#history = fit(my_network, training_set, train_interior_tensor, n_epochs, optimizer_, p=2, verbose=True )
#end_time = time.time()
#total_time = end_time - start_time
#print("Training time: {:.2f} seconds".format(total_time))

u_exact = exact_solution(test_tensor[:,0], test_tensor[:,1], test_tensor[:,2]).reshape(-1,1)

my_network.load_state_dict(torch.load('Trained_Models/sine/model3.pth'))

my_network = my_network.cpu()
u_test_pred = my_network(test_tensor).reshape(-1,1)

#torch.save(my_network.state_dict(), 'Trained_Models/sine/model3.pth')

# Compute the relative L2 error norm (generalization error)

RMSE = np.sqrt(mean_squared_error(u_test_pred.detach().numpy(), u_exact.detach().numpy()))
print("RMSE Test: ", RMSE)

u_test_pred = u_test_pred.reshape(-1, )
u_exact = u_exact.reshape(-1, )

relative = torch.abs((u_test_pred - u_exact) / u_exact)

relative_error_test = torch.sqrt(torch.mean((u_test_pred - u_exact)**2) / torch.mean(u_exact**2))
print("Relative Error Test: ", relative_error_test.detach().numpy())

# Function to append result to CSV file
def append_to_csv(file_name, result):
    file_exists = os.path.isfile(file_name)
    
    with open(file_name, mode='a', newline='') as file:
        writer = csv.writer(file)
        
        # Write header only if file does not exist
        if not file_exists:
            writer.writerow(['Result'])
        
        # Append the result
        writer.writerow([result])
        
# File names for storing results
RMSE_file = 'Results/sine/RMSE_results.csv'
Rel_err_file = 'Results/sine/Rel_err_results.csv'
Time_file = 'Results/sine/Time_results.csv'

# Append results to respective CSV files
#append_to_csv(RMSE_file, RMSE)
#append_to_csv(Rel_err_file, relative_error_test.detach().numpy())
#append_to_csv(Time_file, total_time)

u_test_pred = u_test_pred.reshape(-1,)
relative = relative.reshape(-1, )

abs_err = torch.abs(u_test_pred - u_exact)

### contour plotting

fontsize = 12
cmap =cm.jet


# Plot of ground truth
fontsize = 12
cmap =cm.jet

# Plot of absolute error
fig = p.figure(figsize=(4, 3))
ax =fig.add_subplot(projection='3d')#
rel_err_plot = ax.scatter(test_tensor[:,0].detach().numpy(), test_tensor[:,1].detach().numpy(), test_tensor[:,2].detach().numpy(), c=abs_err.detach().numpy(), 
                          marker='o', cmap=cmap) 
cb = fig.colorbar(rel_err_plot , ax=ax , location='bottom',fraction=0.046, format='%.0e')
cb.ax.tick_params(labelsize=fontsize)
plt.tick_params(axis='both', labelsize=fontsize)
plt.savefig('Results/Figures/sine/sine_abs_err.pdf', bbox_inches="tight")
plt.show()


# 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
