# 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 matplotlib.ticker as ticker

# Create the parser
parser = argparse.ArgumentParser(description='Process beta value')

# Add the argument
parser.add_argument('beta', type=np.float32, help='beta value to be used in the code')

# Parse the argument
args = parser.parse_args()

# Use the argument in your code
beta = args.beta
print(f'Input beta is: {beta}')

# initial condition

def initial_condition(x):
    return torch.sin(x)

def exact_solution(x, t):
    return torch.sin(x-beta*t)


# assigning number of points
initial_pts = 200
left_boundary_pts = 200
right_boundary_pts = 200
residual_pts = 2000

# Type of optimizer (ADAM or LBFGS)
opt_type = "LBFGS"

#first condition
x_init = 2*pi*torch.rand((initial_pts,1))  # initial pts 
t_init =  0*x_init                            #t=0
init =  torch.cat([x_init, t_init],1)                   #concatenate the variable x and t
u_init = initial_condition(init[:,0]).reshape(-1, 1)     #(initial condition)

#second condition
xb_left = 0*torch.ones((left_boundary_pts, 1)) # left spatial boundary 
tb_left = torch.rand((left_boundary_pts, 1)) #randomly generate the value of t
b_left = torch.cat([xb_left, tb_left ],1)    #concatenate
u_b_l = exact_solution(xb_left, tb_left)            # [boundary condition]

#third condition
xb_right = 2*pi*torch.ones((right_boundary_pts, 1)) # right spatial boundary
tb_right = torch.rand((right_boundary_pts, 1)) # right boundary pts
b_right = torch.cat([xb_right, tb_right ],1)   #concatinate
u_b_r =  exact_solution(xb_right, tb_right)    # [boundary condition]

#fourth conditon for interior points
x_interior = 2*pi*torch.rand((residual_pts, 1))   #randomy generate the value of x in domain
t_interior = torch.rand((residual_pts, 1))        #randomly generate the value of t
interior = torch.cat([x_interior, t_interior],1)   #concatenate

training_set = DataLoader(torch.utils.data.TensorDataset(init, u_init, b_left, b_right), batch_size=200, shuffle=False)
    
my_network = NeuralNet(input_dimension = init.shape[1], output_dimension = u_init.shape[1], n_hidden_layers=4, neurons=30)

# Random Seed for weight initialization
retrain = 128
# 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()
    start_total_time = time.time()  # Start timing for the entire training process

    # Loop over epochs
    for epoch in range(num_epochs):
        if verbose: 
            print("################################ ", epoch, " ################################")
        
        start_epoch_time = time.time()  # Start timing for the epoch

        running_loss = list([0])

        # Loop over batches
        for j, (initial, u_initial, bd_left, bd_right) in enumerate(training_set):

            def closure():
                # zero the parameter gradients
                optimizer.zero_grad()
                # for initial
                u_initial_pred_ = model(initial)
                # for left boundary
                u_bd_left_pred_ = model(bd_left)
                # for right boundary
                u_bd_right_pred_ = model(bd_right)

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

                u_x = grad_u_hat[:, 0]
                u_t = grad_u_hat[:, 1]

                # Item 1. below
                loss = torch.mean((u_initial_pred_.reshape(-1, ) - u_initial.reshape(-1, ))**p) + torch.mean((u_t.reshape(-1, ) + beta*(u_x.reshape(-1, )).reshape(-1, ))**p) +torch.mean((u_bd_left_pred_.reshape(-1,) - u_b_l.reshape(-1, ))**p) + torch.mean((u_bd_right_pred_.reshape(-1,) - u_b_r.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)

        #epoch_time = time.time() - start_epoch_time  # Calculate time taken for the epoch
        print('Loss: ', (running_loss[0] / len(training_set)))

        history.append(running_loss[0])

    
    return history

#n_epochs = 5000
#start_time = time.time()
#history = fit(my_network, training_set, interior, 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))

# Number of points for x and t
num_points_x = 256
num_points_t = 100

# Create uniformly spaced values for x and t
x_test = torch.linspace(0, 2 * pi, num_points_x).reshape(-1, 1)
t_test = torch.linspace(0, 1, num_points_t).reshape(-1, 1)

# Create a grid of x and t values
x_grid, t_grid = torch.meshgrid(x_test.squeeze(), t_test.squeeze(), indexing='ij')

# Reshape the grids to create the test tensor
x_flat = x_grid.reshape(-1, 1)
t_flat = t_grid.reshape(-1, 1)
test = torch.cat([x_flat, t_flat], dim=1)

my_network.load_state_dict(torch.load('Trained_Models/Beta40/model3.pth'))
my_network = my_network.cpu()
u_test_pred = my_network(test).reshape(-1,1)

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

u_exact = exact_solution(x_flat, t_flat )

# 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_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/Beta100/RMSE_results.csv'
Rel_err_file = 'Results/Beta100/Rel_err_results.csv'
Time_file = 'Results/Beta100/Time_results.csv'

# Detach and convert tensors to numpy arrays
u_test_pred = u_test_pred.detach().numpy()
x_test = x_test.detach().numpy()
t_test = t_test.detach().numpy()
u_exact = u_exact.detach().numpy()

# Compute the absolute error
abs_error = np.abs(u_test_pred - u_exact)

abs_error = abs_error.reshape(256, 100)


fontsize = 20
aspect = 0.07

# visualize the solution
fig, ax = plt.subplots(1, 1, figsize=(6, 5), constrained_layout=True)
extent = [t_test.min(), t_test.max(), x_test.min(), x_test.max()]

sol_img1 = ax.imshow(abs_error, extent=extent, origin='lower', aspect=aspect, cmap='jet')

# Colorbar
cb = fig.colorbar(sol_img1, ax=ax, location='bottom', aspect=20)

# Set the formatter for the colorbar to scientific notation
cb.formatter = ticker.FuncFormatter(lambda x, pos: f'{x:.1e}')
cb.update_ticks()

# Set only three ticks on the colorbar
cb.set_ticks([abs_error.min(), (abs_error.min() + abs_error.max()) / 2, abs_error.max()])

# Set the font size of the colorbar labels
cb.ax.tick_params(labelsize=fontsize)

# Set axis labels
ax.set_xlabel('t', fontsize=fontsize)
ax.set_ylabel('x', fontsize=fontsize)

# Set specific ticks on the x-axis
ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1])

# Set tick parameters for the axes
plt.tick_params(axis='both', labelsize=fontsize)

# Save the figure
plt.savefig('Results/Beta40/abs_err_adv40.pdf')
plt.show()

# 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)