from utils import *
from datetime import datetime
import time, random, os, gc, copy
import tensorflow as tf
from tensorflow import keras
from keras import layers, Model, optimizers
import scipy
from Functions import *


application_id = 11
Application = {0: 'AirQ', ## Synthetic Data
               1: 'DrugDiffusion', ## Synthetic Data
               2: 'OceanAdvDiffusion',
               3: 'Wave',
               4: 'Burger',
               5: 'KDV1',
               6: 'KDV2_PDEFIND', ## PDE Find  (Data-driven discovery of partial differential equations)
               7: 'NLS',
               8: 'AC',
               9: 'KS',
               10: 'Burger_PDEFIND',  ## PDE Find (Data-driven discovery of partial differential equations)
               11: 'KS_PDEFIND',  ## PDE Find (Data-driven discovery of partial differential equations)
               12: 'NLS_PDEFIND', ## PDE Find (Data-driven discovery of partial differential equations)
               13: 'Shrodinger_RealPDE', ## PDE Find (Data-driven discovery of partial differential equations)
               14: 'Shrodinger_EstimatedPDE', ## PDE Find (Data-driven discovery of partial differential equations)
               }[application_id]


## Configs
## Currently, there are limited choices for number of nodes for the synthetic data. 
## Please refer to the main function for more information regarding the number of samples specified for each graph size
## Maximum number of iterations considered for the synthetic datasets is 5. 
## For more iterations, you should first create the data via data creation file.
## The number of itertions for the public datasets depends on the number of samples and the dataset size.


N_NODES = 3 # Number of nodes. 
LEARNING_RATE = 0.01  # Learning rate
EPOCHS = 100  # Number of epochs for local training
COMMUNICATION_ROUNDS = 5 # Number of communication rounds (iterations)
physics_weights = np.array([0, 0.25, 0.5, 0.75, 1]) ## Physical Weights
num_iterations = 5 # Number of Iterations 

from_iteration = 0 # Starting Iteration
to_iteration = num_iterations # Ending Iteration

non_IID = False ## True if you would like to run the models on Non-IID data
save_log = True ## True if you would like to save the execusion logs
shuffle = False ## True if you would like to shuffle data before execusion
Dirichlet_alpha = 0.5 ## Alpha for generating Non-IID data distribution 
noise = 0.24 #0.0, 0.12, 0.24, 0.36 ## The noise added to the data
noise_on_input = True ## True if you would like to add noise on the input data, otherwise the noise will be added to the output

## The fraction of nodes used for FedAvg Aggregation Algorithm
C = 0.5

## The S and R model parameters for the Segmented Gossip algorithm
R = 2 ## The number of Replica of the mixed model should be at most equal to the number of nodes-1 (the model itself is also included in the aggregation process)
S = 2 ## Model Weights are 8; Segments should be determined in a way by that the model weights can be divided

############################
#### Models Definitions ####
############################

if Application == 'AirQ':
    U, V = 1.0, 1.0  # Wind velocity components
    D = 0.01  # Diffusion coefficient
    S = 0.1  # Source term

    ###############################
    ######## Air Quality ##########
    ###############################

    # Creating a PINN model
    def create_pollutant_dispersion_pinn_model():
        inputs = layers.Input(shape=(3,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    # Physics-informed loss function
    def pollutant_dispersion_loss(model, x, y, t, c_observed, weight_physics=physics_weights):
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, y, t])
            X = tf.concat([x, y, t], axis=1)
            c_pred = model(X)
            c_t = tape.gradient(c_pred, t)
            c_x = tape.gradient(c_pred, x)
            c_y = tape.gradient(c_pred, y)
            c_xx = tape.gradient(c_x, x)
            c_yy = tape.gradient(c_y, y)

            advection_diffusion = c_t + U * c_x + V * c_y - D * (c_xx + c_yy) - S
            physics_loss = tf.reduce_mean(tf.square(advection_diffusion))
            data_loss = tf.reduce_mean(tf.square(c_observed - c_pred))
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, y, t, c_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = pollutant_dispersion_loss(model, x, y, t, c_observed, physics_weight)
                else:
                    X = tf.concat([x, y, t], axis=1)
                    c_pred = model(X)
                    loss = tf.reduce_mean(tf.square(c_observed - c_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:3], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == "DrugDiffusion":

    M0 = 10
    L = 5 #cm ## length of the cylander
    D = 1e-07 #10^-7 cm^2/s (slow diffusion)

    ###################################
    ######## Drug Diffusion ###########
    ###################################

    # Creating a PINN model
    def create_drugdiffusion_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    # Physics-informed loss function
    ## Reference: 
    #   @misc{chasnov2019differential,
    #   title={Differential Equations for Engineers},
    #   author={Chasnov, Jeffrey R},
    #   year={2019},
    #   publisher={The Hong Kong University of Science and Technology Department of Mathematics}
    # }
    def drugdiffusion_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            u_xx = tape.gradient(u_x, x)

            advection_diffusion = D * (u_xx) - u_t
            physics_loss = tf.reduce_mean(tf.square(advection_diffusion))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = drugdiffusion_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'OceanAdvDiffusion':
    U, V = 1.0, 1.0  # Wind velocity components
    D = 0.01  # Diffusion coefficient
    S = 0.1  # Source term

    #######################################
    ######## Advection Diffusion ##########
    #######################################

    # Creating a PINN model
    def create_oceanAD_pinn_model():
        inputs = layers.Input(shape=(3,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

        # Physics-informed loss function
    # Physics-informed loss function
    def OceanAD_dispersion_loss(model, X_, Y, weight_physics=physics_weights):

        with tf.GradientTape(persistent=True) as tape:
            x = tf.reshape(tf.convert_to_tensor(X_[:,0]), (X_.shape[0], 1))
            y = tf.reshape(tf.convert_to_tensor(X_[:,1]), (X_.shape[0], 1))
            t = tf.reshape(tf.convert_to_tensor(X_[:,2]), (X_.shape[0], 1))

            tape.watch([x, y, t])
            X = tf.concat([x, y, t], axis=1)
            c_pred = model(X)

            y_true = tf.reshape(tf.convert_to_tensor(Y, dtype=c_pred.dtype), (Y.shape[0], 1)) 

            c_t = tape.gradient(c_pred, t)
            c_x = tape.gradient(c_pred, x)
            c_y = tape.gradient(c_pred, y)

            c_xx = tape.gradient(c_x, x)
            c_yy = tape.gradient(c_y, y)

            advection_diffusion = c_t + U * c_x + V * c_y - D * (c_xx + c_yy) - S
        
            physics_loss = tf.reduce_mean(tf.square(advection_diffusion))
            data_loss = tf.reduce_mean(tf.square(y_true - c_pred))
            data_loss = tf.cast(data_loss, dtype=tf.float64)

            loss = data_loss + weight_physics * physics_loss

        return  loss
    
    #####################################
    ######## General Functions ##########
    #####################################
    # Physics-informed loss function
    def OceanAD_loss(model, x, y, t, c_observed, weight_physics=physics_weights):
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, y, t])
            X = tf.concat([x, y, t], axis=1)
            c_pred = model(X)
            c_t = tape.gradient(c_pred, t)
            c_x = tape.gradient(c_pred, x)
            c_y = tape.gradient(c_pred, y)
            c_xx = tape.gradient(c_x, x)
            c_yy = tape.gradient(c_y, y)

            advection_diffusion = c_t + U * c_x + V * c_y - D * (c_xx + c_yy) - S
            physics_loss = tf.reduce_mean(tf.square(advection_diffusion))
            data_loss = tf.reduce_mean(tf.square(c_observed - c_pred))
        return  data_loss + weight_physics * physics_loss

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, y, t, c_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = OceanAD_loss(model, x, y, t, c_observed, physics_weight)
                else:
                    X = tf.concat([x, y, t], axis=1)
                    c_pred = model(X)
                    loss = tf.reduce_mean(tf.square(c_observed - c_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:3], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses
    
elif Application == 'Wave': 

    #################################
    ######## Wave Equation ##########
    #################################

    # Creating a PINN model
    def create_wave_pinn_model():
        inputs = layers.Input(shape=(3,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(2)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    # Physics-informed loss function
    ## Reference: Ocean Modelling
    def wave_loss(model, x, y, t, u, v, weight_physics=physics_weights):
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, y, t])
            X = tf.concat([x, y, t], axis=1)
            preds = model(X)
            u_pred = preds[:,0]
            v_pred = preds[:,1]

            u_x  = tape.gradient(u_pred, x)
            u_xx  = tape.gradient(u_x, x)

            u_y = tape.gradient(u_pred, y)
            u_yy = tape.gradient(u_y, y)

            u_t = tape.gradient(u_pred, t)
            u_tt = tape.gradient(u_t, t)

            v_x  = tape.gradient(v_pred, x)
            v_y  = tape.gradient(v_pred, y)

            advection_diffusion = u_tt - v*(u_xx + u_yy) - (v_x*u_x + v_y*u_y)
            physics_loss = tf.reduce_mean(tf.square(advection_diffusion))

            observed = tf.convert_to_tensor([u, v], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            data_loss = tf.reduce_mean(tf.square(observed - pred))

        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, y, t, u, v = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = wave_loss(model, x, y, t, u, v, physics_weight)
                else:
                    X = tf.concat([x, y, t], axis=1)
                    preds = model(X)
                    u_pred = preds[:,0]
                    v_pred = preds[:,1]
                    observed = tf.convert_to_tensor([u, v], tf.float32)
                    pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
                    pred = tf.expand_dims(pred, axis=2)
                    loss = tf.reduce_mean(tf.square(observed - pred))
            
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, _, u_val, v_val = validation_data
            X_val = tf.concat(validation_data[:3], axis=1)
            preds = model(X_val)
            u_pred = preds[:,0]
            v_pred = preds[:,1]
            observed = tf.convert_to_tensor([u_val, v_val], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            loss = tf.reduce_mean(tf.square(observed - pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'Burger':
    lambda1 = 1.0
    lambda2 = 0.01/np.pi
    
    ###################################
    ######## Burger Equation ###########
    ###################################

    # Creating a PINN model
    def create_burger_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: https://en.wikipedia.org/wiki/Burgers%27_equation
    ## Reference: Ocean Modelling
    ## du/dt + u*du/dx - v*d^2u/dx2 = 0
    def burger_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            u_xx = tape.gradient(u_x, x)

            pde = u_t + lambda1 * u_pred* u_x - lambda2*u_xx
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = burger_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'KDV1':
    lambda1 = 6.0
    lambda2 = 1
    
    ###################################
    ######## KDV Equation ###########
    ###################################

    # Creating a PINN model
    def create_KDV_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: Pinnstf Paper
    ## $$ u_t + \lambda_1 uu_x + \lambda_2 u_{xxx} = 0,$$
    def KDV_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)
            
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)
            u_xx = tape.gradient(u_x, x)
            u_xxx = tape.gradient(u_xx, x)

            pde = u_t + lambda1* u_pred * u_x - lambda2* u_xxx
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = KDV_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses
 
elif Application == 'KDV2_PDEFIND': ## PDE Find  (Data-driven discovery of partial differential equations)
    lambda1 = 6.0
    lambda2 = 1
    
    ###################################
    ######## KDV Equation ###########
    ###################################

    # Creating a PINN model
    def create_KDV2_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: PDE Find Paper
    ## ut + 6uux + uxxx = 0
    def KDV2_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)
            
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)
            u_xx = tape.gradient(u_x, x)
            u_xxx = tape.gradient(u_xx, x)

            pde = u_t + lambda1* u_pred * u_x + lambda2* u_xxx 
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = KDV2_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'NLS':

    ###################################
    ######## NLS Equation ###########
    ###################################

    # Creating a PINN model
    def create_NLS_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(2)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: Pinnstf Paper
    ## Continuous Forward Schrodinger Equation
    ## we partition $h(t, x)$ into its real part $u$ and imaginary part $v$. 
    ## Thus, our complex-valued neural network representation is $[u(t, x), v(t, x)]$.
    def NLS_loss(model, x, t, u_observed, v_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            pred = model(X)
            u_pred = pred[:,0]
            v_pred = pred[:,1]
            
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            v_t = tape.gradient(v_pred, t)
            v_x = tape.gradient(v_pred, x)

            u_xx = tape.gradient(u_x, x)
            v_xx = tape.gradient(v_x, x)
    
            pde_u = u_t + 0.5 * v_xx + (u_pred ** 2 + v_pred ** 2) * v_pred
            pde_v = v_t - 0.5 * u_xx - (u_pred** 2 + v_pred ** 2) * u_pred

            physics_loss = tf.reduce_mean(tf.square([pde_u, pde_v]))

            observed = tf.convert_to_tensor([u_observed, v_observed], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            data_loss = tf.reduce_mean(tf.square(observed - pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed, v_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = NLS_loss(model, x,t, u_observed, v_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, u_val, v_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            preds = model(X_val)
            u_pred = preds[:,0]
            v_pred = preds[:,1]
            observed = tf.convert_to_tensor([u_val, v_val], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            loss = tf.reduce_mean(tf.square(observed - pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'AC':
    ###################################
    ######## AC Equation ###########
    ###################################

    # Creating a PINN model
    def create_AC_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: Pinnstf Paper
    ## Allen-Cahn
    def AC_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)

            u_x = tape.gradient(u_pred, x)
            u_xx = tape.gradient(u_x, x)

            pde = 5.0 * u_pred - 5.0 * (u_pred**3) + 0.0001 * u_xx
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = AC_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'KS':
    ###################################
    ######## KS Equation ###########
    ###################################

    # Creating a PINN model
    def create_KS_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: Sparsistent model discovery Paper
    ## Kuramoto–Sivashinsky
    ## Wiki Pedia: {\displaystyle u_{t}+u_{xx}+u_{xxxx}+{\frac {1}{2}}u_{x}^{2}=0}
    def KS_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)

            u_x = tape.gradient(u_pred, x)
            u_t = tape.gradient(u_pred, t)
            u_xx = tape.gradient(u_x, x)
            u_xxx = tape.gradient(u_xx, x)

            pde = u_t + u_xx + u_xxx + 0.5 * u_x**2
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = KS_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'Burger_PDEFIND': ## Burger_PDEFIND (Data-Driven discovery of partial differential equations)
    lambda1 = 1.0
    lambda2 = 0.01/np.pi
    
    ###################################
    ######## Burger_PDEFIND Equation ###########
    ###################################

    # Creating a PINN model
    def create_burger_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: PDE_FIND
    ## ut + uux − epsilon*uxx = 0
    def burger_loss(model, x, t, u_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            u_xx = tape.gradient(u_x, x)

            pde = u_t + lambda1 * u_pred* u_x - lambda2*u_xx
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = burger_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'KS_PDEFIND':

    lambda1 = 0.48
    lambda2 = 0.52
    lambda3 = 0.53
    ###################################
    ######## KS_PDEFIND Equation ######
    ###################################

    # Creating a PINN model
    def create_KS_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(1)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: PDE Find Data-Driven discovery of PDEs
    ## Kuramoto–Sivashinsky
    ## ut + uux + uxx + uxxxx = 0
    def KS_loss(model, x, t, u_observed, weight_physics=physics_weights):
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            u_pred = model(X)

            u_x = tape.gradient(u_pred, x)
            u_t = tape.gradient(u_pred, t)
            u_xx = tape.gradient(u_x, x)
            u_xxx = tape.gradient(u_xx, x)

            pde = u_t + lambda1 * u_pred * u_x + lambda2 * u_xx + lambda3 * u_xxx
            physics_loss = tf.reduce_mean(tf.square(pde))
            data_loss = tf.reduce_mean(tf.square(u_observed - u_pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = KS_loss(model, x,t, u_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, c_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            c_pred = model(X_val)
            loss = tf.reduce_mean(tf.square(c_val - c_pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'NLS_PDEFIND': ## PDE Find  (Data-driven discovery of partial differential equations)
    
    ###################################
    ######## NLS Equation ###########
    ###################################

    # Creating a PINN model
    def create_NLS_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(2)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: PDE FIND Paper
    ## iut + 1/2uxx + |u|**2 u = 0
    # u_t = (0.000170 +0.484056i)u_{xx} + (0.000160 +0.985729i)u|u|^2
    def NLS_loss(model, x, t, u_observed, v_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            pred = model(X)
            u_pred = pred[:,0]
            v_pred = pred[:,1]
            
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            v_t = tape.gradient(v_pred, t)
            v_x = tape.gradient(v_pred, x)

            u_xx = tape.gradient(u_x, x)
            v_xx = tape.gradient(v_x, x)
    
            ## iut + 1/2uxx + |u|**2 u = 0
            # u_t = (0.000170 +0.484056i)u_{xx} + (0.000160 +0.985729i)u|u|^2

            # u_t - (0.000170) u_xx - (0.000160) u|u|^2
            # v_t - (0.484056) u_xx - (0.985729) u|u|^2

            pde_u = u_t - 0.000170 * u_xx - 0.000160 * (u_pred ** 2) * u_pred
            pde_v = v_t - 0.484056 * v_xx + 0.985729 * (v_pred ** 2) * v_pred

            physics_loss = tf.reduce_mean(tf.square([pde_u, pde_v]))

            observed = tf.convert_to_tensor([u_observed, v_observed], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            data_loss = tf.reduce_mean(tf.square(observed - pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed, v_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = NLS_loss(model, x,t, u_observed, v_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, u_val, v_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            preds = model(X_val)
            u_pred = preds[:,0]
            v_pred = preds[:,1]
            observed = tf.convert_to_tensor([u_val, v_val], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            loss = tf.reduce_mean(tf.square(observed - pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'Shrodinger_RealPDE': ## PDE Find  (Data-driven discovery of partial differential equations)
    ########################################
    ######## Shrodinger Equation ###########
    ########################################

    # Creating a PINN model
    def create_Shrodinger_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(2)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: PDE FIND Paper
    ## iu_t = \frac{-1}{2}u_{xx} + \frac{x^2}{2}u
    # PDE derived using STRidge
    # u_t = (0.000011 +0.498831i)u_{xx} + (0.000013 -0.997372i)uV
    def Shrodinger_loss(model, x, t, u_observed, v_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            pred = model(X)
            u_pred = pred[:,0]
            v_pred = pred[:,1]
            
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            v_t = tape.gradient(v_pred, t)
            v_x = tape.gradient(v_pred, x)

            u_xx = tape.gradient(u_x, x)
            v_xx = tape.gradient(v_x, x)
   
            ## iu_t = \frac{-1}{2}u_{xx} + \frac{x^2}{2}u
            pde_u = u_t + 0.5 * u_xx - (x**2)/2 * u_pred
            pde_v = v_t + 0.5 * v_xx - (x**2)/2 * v_pred

            physics_loss = tf.reduce_mean(tf.square([pde_u, pde_v]))

            observed = tf.convert_to_tensor([u_observed, v_observed], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            data_loss = tf.reduce_mean(tf.square(observed - pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed, v_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = Shrodinger_loss(model, x,t, u_observed, v_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, u_val, v_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            preds = model(X_val)
            u_pred = preds[:,0]
            v_pred = preds[:,1]
            observed = tf.convert_to_tensor([u_val, v_val], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            loss = tf.reduce_mean(tf.square(observed - pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses

elif Application == 'Shrodinger_EstimatedPDE': ## PDE Find  (Data-driven discovery of partial differential equations)
    ########################################
    ######## Shrodinger Equation ###########
    ########################################

    # Creating a PINN model
    def create_Shrodinger_pinn_model():
        inputs = layers.Input(shape=(2,))
        x = layers.Dense(10, activation="tanh")(inputs)
        x = layers.Dense(10, activation="tanh")(x)
        x = layers.Dense(10, activation="tanh")(x)
        outputs = layers.Dense(2)(x)
        model = Model(inputs=inputs, outputs=outputs)
        return model

    ## Reference: PDE FIND Paper
    ## iu_t = \frac{-1}{2}u_{xx} + \frac{x^2}{2}u
    # PDE derived using STRidge
    # u_t = (0.000011 +0.498831i)u_{xx} + (0.000013 -0.997372i)uV
    def Shrodinger_loss(model, x, t, u_observed, v_observed, weight_physics=physics_weights):
        ## ut=Duxx, 0≤x≤L,t>0. => 0 = D(∂²u/∂x²) - ut
        with tf.GradientTape(persistent=True) as tape:
            tape.watch([x, t])
            X = tf.concat([x, t], axis=1)
            pred = model(X)
            u_pred = pred[:,0]
            v_pred = pred[:,1]
            
            u_t = tape.gradient(u_pred, t)
            u_x = tape.gradient(u_pred, x)

            v_t = tape.gradient(v_pred, t)
            v_x = tape.gradient(v_pred, x)

            u_xx = tape.gradient(u_x, x)
            v_xx = tape.gradient(v_x, x)
   
            # u_t = (0.000011 +0.498831i)u_{xx} + (0.000013 -0.997372i)uV
            pde_u = u_t - 0.000011 * u_xx - 0.000013 * (u_pred) * v_pred
            pde_v = v_t - 0.498831 * v_xx + 0.997372 * (v_pred) * u_pred


            physics_loss = tf.reduce_mean(tf.square([pde_u, pde_v]))

            observed = tf.convert_to_tensor([u_observed, v_observed], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            data_loss = tf.reduce_mean(tf.square(observed - pred))

            # print(f"data loss: {data_loss}, physical loss: {physics_loss}")
        return  data_loss + weight_physics * physics_loss

    #####################################
    ######## General Functions ##########
    #####################################

    # Local training 
    def train_local_model(model, local_data, epochs, use_physics=True, physics_weight=physics_weights):
        optimizer = optimizers.Adam(learning_rate=LEARNING_RATE)
        x, t, u_observed, v_observed = local_data
        for _ in range(epochs):
            with tf.GradientTape() as tape:
                if use_physics:
                    loss = Shrodinger_loss(model, x,t, u_observed, v_observed, physics_weight)
                else:
                    X = tf.concat([x, t], axis=1)
                    u_pred = model(X)
                    loss = tf.reduce_mean(tf.square(u_observed - u_pred))
            gradients = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # Evaluation Function
    def evaluate_models(models, validation_data):
        losses = []
        for model in models:
            _, _, u_val, v_val = validation_data
            X_val = tf.concat(validation_data[:2], axis=1)
            preds = model(X_val)
            u_pred = preds[:,0]
            v_pred = preds[:,1]
            observed = tf.convert_to_tensor([u_val, v_val], tf.float32)
            pred = tf.convert_to_tensor([u_pred, v_pred], tf.float32)
            pred = tf.expand_dims(pred, axis=2)
            loss = tf.reduce_mean(tf.square(observed - pred))
            losses.append(loss.numpy())
        # return np.mean(losses)
        return losses


############################
####### Main Function ######
############################

if __name__ == "__main__":
    print()

    start_time = time.time()
    
    Models_dict = dict()
    print("Selected Application:", Application)
    
    nst_keys = {3: 180, 10: 400, 30: 800, 50: 800, 100:1000}
    nsv_keys = {3: 20, 10: 100, 30: 200, 50: 200, 100:400}

    if N_NODES in list(nst_keys.keys()):
        nst = nst_keys[N_NODES] ## Number of samples for the training data
        nsv = nsv_keys[N_NODES] ## Number of samples for the validation data

    elif Application in ['AirQ', 'DrugDiffusion'] or num_iterations > 5:
        print("There is no synthetic data available for this number of nodes.")
        print("Please first generate your desired dataset.")
        x = input("Do you want to create a random dataset for your settings? Please enter Y as yes, and N as No: \t")
        if x == 'Y':
            nst = int(input("Please Enter the Number of Tranining Samples. The Minimum Number is 20.\t"))
            nsv = int(input("Please Enter the Number of Validation Samples. The Minimum Number is 10.\t"))

            if nst < 20 or nsv < 10:
                print("The Entered Number Is Not Valid. Please Enter A Valid Number")
                exit()

            file_path = f"./Datasets/Synthetic/IID/{N_NODES} Nodes/"
            if not os.path.isdir(file_path):
                os.makedirs(file_path)

            if Application == 'AirQ':
                for ind in range(num_iterations):
                    file_path_index = file_path + f'IT{ind}' + '_'
                    generate_pollutant_data_random(n_samples=nst, noise_level=noise, save=True, valid=False, file_path=file_path_index)
                generate_pollutant_data_random(n_samples=nsv, noise_level=noise, save=True, valid=True, file_path=file_path)

            elif Application == 'Drug_Diffusion':
                for ind in range(num_iterations):
                    file_path_index = file_path + f'IT{ind}' + '_'
                    generate_drugdiffusion_data_random(n_samples=nst, noise_level=noise, save=True, valid=False, file_path=file_path_index)
                generate_drugdiffusion_data_random(n_samples=nsv, noise_level=noise, save=True, valid=True, file_path=file_path)
        
        else:
            print("Bye until next time ...")
            exit()

    else:
        print("There Is No Specified Number of Training/Validation Samples For Your Graph Size.")
        nst = int(input("Please Enter the Number of Tranining Samples. The Minimum Number is 20."))
        nsv = int(input("Please Enter the Number of Validation Samples. The Minimum Number is 10."))
        if nst < 20 or nsv < 10:
            print("The Entered Number Is Not Valid. Please Enter A Valid Number")
            exit()

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

    if Application == 'AirQ':
        folder_path = f"./Datasets/Synthetic/IID/{N_NODES} Nodes/"
        training_file_name = f"ST_{Application}_dataset_{nst}_{noise}.pkl" ## 3 nodes: 60; 10 nodes: 18; 15 nodes: 3
        validation_file_name = f"SV_{Application}_dataset_{nsv}_{noise}.pkl"

        Models_dict = {'AirQ':create_pollutant_dispersion_pinn_model}

        if os.path.isfile(training_file_name):
            ST_node = dict()
            for ind in range(num_iterations):
                ST_node[ind] = read_from_pickle_airquality(f"{folder_path}IT{ind}_{training_file_name}")
            SV_Data = read_from_pickle_airquality(folder_path + validation_file_name)
            num_true_values = 0
        else:
            file_path = f"./Datasets/Synthetic/IID/{N_NODES} Nodes/"
            if not os.path.isdir(file_path):
                os.makedirs(file_path)
            
            for ind in range(num_iterations):
                file_path_index = file_path + f'IT{ind}' + '_'
                generate_pollutant_data_random(n_samples=nst, noise_level=noise, save=True, valid=False, file_path=file_path_index)
            generate_pollutant_data_random(n_samples=nsv, noise_level=noise, save=True, valid=True, file_path=file_path)

            ST_node = dict()
            for ind in range(num_iterations):
                ST_node[ind] = read_from_pickle_airquality(f"{folder_path}IT{ind}_{training_file_name}")
            SV_Data = read_from_pickle_airquality(folder_path + validation_file_name)
            num_true_values = 0

    elif Application == "DrugDiffusion":
        folder_path = f"./Datasets/Synthetic/IID/{N_NODES} Nodes/"
        training_file_name = f"ST_{Application}_dataset_{nst}_{noise}.pkl" ## 3 nodes: 60; 10 nodes: 18; 15 nodes: 3
        validation_file_name = f"SV_{Application}_dataset_{nsv}_{noise}.pkl"

        Models_dict = {'DrugDiffusion': create_drugdiffusion_pinn_model}
        
        if os.path.isfile(training_file_name):
            ST_node = dict()
            for ind in range(num_iterations):
                ST_node[ind] = read_from_pickle_drugdiffusion(f"{folder_path}IT{ind}_{training_file_name}")
            SV_Data = read_from_pickle_drugdiffusion(folder_path + validation_file_name)
            num_true_values = 0
        
        else:
            file_path = f"./Datasets/Synthetic/IID/{N_NODES} Nodes/"
            for ind in range(num_iterations):
                file_path_index = file_path + f'IT{ind}' + '_'
                generate_drugdiffusion_data_random(n_samples=nst, noise_level=noise, save=True, valid=False, file_path=file_path_index)
            generate_drugdiffusion_data_random(n_samples=nsv, noise_level=noise, save=True, valid=True, file_path=file_path)

            ST_node = dict()
            for ind in range(num_iterations):
                ST_node[ind] = read_from_pickle_drugdiffusion(f"{folder_path}IT{ind}_{training_file_name}")
            SV_Data = read_from_pickle_drugdiffusion(folder_path + validation_file_name)
            num_true_values = 0

    elif Application == 'OceanAdvDiffusion':
        Models_dict = {'OceanAdvDiffusion': create_oceanAD_pinn_model}

        file_path = "./Datasets/Public/AdvDif"

        x_data = np.loadtxt(f"{file_path}/X_star.txt") # N x 2  
        t_data = np.loadtxt(f"{file_path}/T_star.txt") # T
        T_data = np.loadtxt(f"{file_path}/Temp.txt")   # N x T
        xy = np.tile(x_data, (t_data.shape[0],1)) # TN x 2
        t = np.repeat(t_data.reshape(-1,1), x_data.shape[0], axis=0) # TN x 1
        Y = T_data.T.reshape(-1,1) # TN x 1
        X = np.concatenate([xy,t, Y], axis=1) # TN x 4

        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)

        print("==============================")
        print("x.shape:", x_data.shape)
        print("t.shape:", t_data.shape)
        print("T.shape:", T_data.shape)
        print("X.shape:", X.shape)
        print("Y.shape:", Y.shape)
        print("X: %s ± %s" % (X.mean(axis=0), X.std(axis=0)))
        print("Y: %s ± %s" % (Y.mean(axis=0), Y.std(axis=0)))

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(list(X[:, i])[ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(list(X[:, i])[(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)
        print("==============================")

    elif Application == 'Wave':
        Models_dict = {'Wave': create_wave_pinn_model}

        file_path = "./Datasets/Public/Waves_Square"
        X = np.loadtxt(f"{file_path}/X_star.txt") # 2 x N
        T = np.loadtxt(f"{file_path}/T_star.txt") # T
        U = np.loadtxt(f"{file_path}/U_star.txt") # N x T  
        V = np.loadtxt(f"{file_path}/H_0.txt")    # N

        # print("samples:", T.shape[0]*X.shape[1])

        xx = np.tile(X.T, (T.shape[0],1)) # TN x 2
        tt = np.repeat(T.reshape(-1,1), X.shape[1], axis=0) # TN x 1

        uu = U.T.reshape(-1,1) # TN x 1
        vv = np.tile(V.reshape(-1,1), (T.shape[0],1)) # TN x 1

        X = np.concatenate([xx, tt, uu, vv], axis=1) # TN x 4
        num_true_values = 2
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(list(X[:, i])[ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(list(X[:, i])[(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'Burger':
        Models_dict = {'Burger': create_burger_pinn_model}

        # x     = np.load('./Datasets/Public/Burger/burgers_x_2000.npy')
        # t     = np.load('./Datasets/Public/Burger/burgers_t_2000.npy')
        # u     = np.array(np.load('./Datasets/Public/Burger/burgers_u_2000.npy'),dtype=np.float32).reshape(len(x),len(t))

        data = scipy.io.loadmat('./Datasets/Public/Burger/burgers_shock.mat')
        x = np.tile(data['x'], (data['t'].shape[0],1)) # TN x 1
        t = np.repeat(data['t'], data['x'].shape[0], axis=0) # TN x 1
        u = data['usol'].T.reshape(-1,1) # TN x 1

        print("x", x.shape)
        print("t", t.shape)
        print("u", u.shape)

        X = np.concatenate([x, t, u], axis=1) # TN x 2
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(list(X[:, i])[ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(list(X[:, i])[(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'KDV1': ##ToDo
        Models_dict = {'KDV1': create_KDV_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/Pinnstf Data/KdV.mat')
        x_data_ = data['x'][0]
        t_data_ = data['tt'][0]
        u_data = data['uu'].T.reshape(-1, 1)

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)

        print(x_data.shape)
        print(t_data.shape)
        print(u_data.shape)

        X = np.concatenate([x_data, t_data, u_data], axis=1) # TN x 3
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(X[:, i][ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(X[:, i][(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'KDV2_PDEFIND': ## PDE Find Repo (Data-driven discovery of partial differential equations)
        Models_dict = {'KDV2_PDEFIND': create_KDV2_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/PDE_Find/kdv.mat')
        x_data_ = data['x'][0]
        t_data_ = data['t']
        u_data = tf.math.real(data['usol'].T.reshape(-1,1))

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)
        
        print(x_data.shape)
        print(t_data.shape)
        print(u_data.shape)

        X = np.concatenate([x_data, t_data, u_data], axis=1) # TN x 3
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(list(X[:, i])[ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(list(X[:, i])[(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'NLS':
        Models_dict = {'NLS': create_NLS_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/Pinnstf Data/NLS.mat')
        x_data_ = data['x'][0]
        t_data_ = data['tt'][0]
        u_data = data['uu'].T.reshape(-1,1)

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)
        
        print(x_data.shape)
        print(t_data.shape)
        print(u_data.shape)

        X = np.concatenate([x_data, t_data, tf.math.real(u_data), tf.math.imag(u_data)], axis=1) # TN x 4
        num_true_values = 2
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(X[:, i][ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(X[:, i][(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("========================")

    elif Application == 'AC': 
        Models_dict = {'AC': create_AC_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/Pinnstf Data/AC.mat')

        x_data_ = data['x'][0]
        t_data_ = data['tt']
        u_data = tf.math.real(data['uu'].T.reshape(-1,1))

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)

        print(x_data.shape)
        print(t_data.shape)
        print(u_data.shape)

        X = np.concatenate([x_data, t_data, u_data], axis=1) # TN x 3
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(X[:, i][ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(X[:, i][(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("========================")

    elif Application == 'KS':
        Models_dict = {'KS': create_KS_pinn_model}
        x     = np.load('./Datasets/Public/KS/KS_x.npy')
        t     = np.load('./Datasets/Public/KS/KS_t.npy')
        u     = np.array(np.load('./Datasets/Public/KS/KS_u.npy'),dtype=np.float32).reshape(len(x),len(t))

        x_data = tf.repeat(x, t.shape[0])
        t_data = tf.repeat(t, x.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)
        u_data = tf.reshape(u, (u.shape[0]*u.shape[1], 1))

        print(x_data.shape, t_data.shape, u_data.shape)
    
        X = np.concatenate([x_data, t_data, u_data], axis=1) # TN x 3
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)
        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(X[:, i][ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(X[:, i][(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'Burger_PDEFIND': 
        Models_dict = {'Burger_PDEFIND': create_burger_pinn_model}

        data = scipy.io.loadmat('./Datasets/Public/PDE_Find/burgers.mat')

        x_data_ = data['x'][0]
        t_data_ = data['t']
        u_data = tf.math.real(data['usol'].T.reshape(-1,1)) ## the imaginary parts are discarded in the paper's codes

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)

        print("x", x_data.shape)
        print("t", t_data.shape)
        print("u", u_data.shape)

        X = np.concatenate([x_data, t_data, u_data], axis=1) # TN x 3
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(list(X[:, i])[ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(list(X[:, i])[(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'KS_PDEFIND':
        Models_dict = {'KS_PDEFIND': create_KS_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/PDE_Find/kuramoto_sivishinky.mat')

        x_data_ = data['x'][:,0]
        t_data_ = data['tt'][0,:]
        u_data = data['uu'].T.reshape(-1,1)

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)

        print("x", x_data.shape)
        print("t", t_data.shape)
        print("u", u_data.shape)

        X = np.concatenate([x_data, t_data, u_data], axis=1) # TN x 3
        num_true_values = 1
        if shuffle:
            np.random.shuffle(X)
        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(list(X[:, i])[ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(list(X[:, i])[(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("==============================")

    elif Application == 'NLS_PDEFIND': ## PDE Find Repo (Data-driven discovery of partial differential equations)
        Models_dict = {'NLS_PDEFIND': create_NLS_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/PDE_Find/nls.mat')

        x_data_ = data['x'][0]
        t_data_ = data['t'][:,0]
        u_data = data['usol'].T.reshape(-1,1)

        print(x_data_.shape)
        print(t_data_.shape)
        print(u_data.shape)

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)
        
        print(x_data.shape)
        print(t_data.shape)
        print(u_data.shape)

        X = np.concatenate([x_data, t_data, tf.math.real(u_data), tf.math.imag(u_data)], axis=1) # TN x 4
        num_true_values = 2
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(X[:, i][ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(X[:, i][(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("========================")

    elif Application in ['Shrodinger_RealPDE', 'Shrodinger_EstimatedPDE']:
        Models_dict = {'Shrodinger_RealPDE': create_Shrodinger_pinn_model,
                       'Shrodinger_EstimatedPDE': create_Shrodinger_pinn_model}
        data = scipy.io.loadmat('./Datasets/Public/PDE_Find/harmonic_osc.mat')
        x_data_ = data['x'][0]
        t_data_ = data['t'][:,0]
        u_data = data['usol'].T.reshape(-1,1)

        print(x_data_.shape)
        print(t_data_.shape)
        print(u_data.shape)

        x_data = np.tile(x_data_, t_data_.shape[0])
        t_data = tf.repeat(t_data_, x_data_.shape[0])

        x_data = tf.expand_dims(x_data, axis=1)
        t_data = tf.expand_dims(t_data, axis=1)
        
        print(x_data.shape)
        print(t_data.shape)
        print(u_data.shape)

        X = np.concatenate([x_data, t_data, tf.math.real(u_data), tf.math.imag(u_data)], axis=1) # TN x 4
        num_true_values = 2
        if shuffle:
            np.random.shuffle(X)

        print("X.shape:", X.shape)

        ST_node = dict()
        for ind in range(num_iterations):
            data_list = [tf.convert_to_tensor(X[:, i][ind*nst: (ind+1)*nst], dtype=tf.float32) for i in range(X.shape[1])]
            data_list = [tf.reshape(x, (x.shape[0],1)) for x in data_list]

            ST_node[ind] = tuple(data_list)

        val_data = [tf.convert_to_tensor(X[:, i][(ind+1)*nst: (ind+1)*nst + nsv], dtype=tf.float32) for i in range(X.shape[1])]
        val_data = [tf.reshape(x, (x.shape[0],1)) for x in val_data]

        SV_Data = tuple(val_data)

        del X
        gc.collect()
        print("========================")


    # Creating timestamp for the directory name
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_dir = f"Output_{Application}_{N_NODES}_{num_iterations}_{COMMUNICATION_ROUNDS}_{noise}_{timestamp}"
    if save_log:
        os.makedirs(output_dir, exist_ok=True)


    # Adjacency matrix A (Uniform)
    A = np.ones((N_NODES, N_NODES)) ## All the nodes are connected to each other
    # A = np.random.randint(2, size=(N_NODES, N_NODES)) ## Random connection between pairs of nodes

    configs = {'N_NODES' : N_NODES,
            'LEARNING_RATE' : LEARNING_RATE,
            'EPOCHS' : EPOCHS,
            'COMMUNICATION_ROUNDS' : COMMUNICATION_ROUNDS,
            'physics_weights' : physics_weights,
            'num_iterations' : num_iterations,
            'from_iteration' : from_iteration,
            'to_iteration' : to_iteration,
            'non_IID' : non_IID,
            'save_log' : save_log,
            'shuffle' : shuffle,
            'Dirichlet_alpha' : Dirichlet_alpha,
            'noise' : noise,
            'noise_on_input': noise_on_input,
            'num_true_values' : num_true_values,
            'nst' : nst,
            'A' : A,
            'C' : C,
            'S' : S,
            'R' : R,
            'root_path': output_dir,
            }
    
    print("PIDFL Algorithm Running ...")
    ## PIDFL Algorithm 
    PIDF_main(Application, configs, ST_node, SV_Data, Models_dict[Application], train_local_model, evaluate_models)

    print("\nCL Base Model Running ...")
    ## Cental Learning Base Model; No FL
    CL_Base_model(Application, configs, ST_node, SV_Data, Models_dict[Application], train_local_model, evaluate_models)

    print("\nCFL FedAvg Running ...")
    ## CFL FedAvg Aggregation Algorithm
    Fed_Avg(Application, configs, ST_node, SV_Data, Models_dict[Application], train_local_model, evaluate_models)
    
    print("\nDFL Segmented Gossip Running ...")
    ## DFL Segmented Gossip Algorithm
    Segmented_Gossip(Application, configs, ST_node, SV_Data, Models_dict[Application], train_local_model, evaluate_models)


    finish_time = time.time()
    print(f"Time taken by this process: {(finish_time-start_time)/60} Minutes",)