from utils import *
import numpy as np
import random
import os


applications = {
                0: 'Air_Dispersion',
                1: 'Drug_Diffusion',
                }
application_id = 0

n_nodes = 100
train_num_samples = {3: 180, 10: 400, 50: 800, 100: 1000}[n_nodes] # For nodes 3 and 10: 180; for nodes 50: 480
valid_num_samples = {3: 20, 10: 100, 50: 200, 100: 400}[n_nodes]
noise = 0.24 ## 0, 0.12, 0.24, 0.36

folder_path = './Datasets/Synthetic'

num_iterations = 5

file_path = f"{folder_path}/IID/{n_nodes} Nodes/"
if not os.path.isdir(file_path):
    os.makedirs(file_path)

if applications[application_id] == 'Air_Dispersion':
    for ind in range(num_iterations):
        file_path_index = file_path + f'IT{ind}' + '_'
        generate_pollutant_data_random(n_samples=train_num_samples, noise_level=noise, save=True, valid=False, file_path=file_path_index)
    generate_pollutant_data_random(n_samples=valid_num_samples, noise_level=noise, save=True, valid=True, file_path=file_path)

elif applications[application_id] == 'Drug_Diffusion':
    for ind in range(num_iterations):
        file_path_index = file_path + f'IT{ind}' + '_'
        generate_drugdiffusion_data_random(n_samples=train_num_samples, noise_level=noise, save=True, valid=False, file_path=file_path_index)
    generate_drugdiffusion_data_random(n_samples=valid_num_samples, noise_level=noise, save=True, valid=True, file_path=file_path)

else:
    print("Application id is not valid.")
