import pandas as pd
import numpy as np
import tensorflow as tf
import pickle
import os

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

def save_to_pickle(filename, item):
    with open(filename, 'wb') as handle:
        pickle.dump(item, handle, protocol=pickle.HIGHEST_PROTOCOL)


import tensorflow as tf
import pickle
## this function reads the instances from saved ones
def read_from_pickle_airquality(path):
    """Read dataset from file (pickle)
    """
    objects = []
    with (open(path, "rb")) as openfile:
        while True:
            try:
                objects.append(pickle.load(openfile))
            except EOFError:
                break
    objects = objects[0]
    x, y, t, c = objects
    
    return tf.convert_to_tensor(x, dtype=tf.float32), tf.convert_to_tensor(y, dtype=tf.float32), tf.convert_to_tensor(t, dtype=tf.float32), tf.convert_to_tensor(c, dtype=tf.float32)
    

def read_from_pickle_drugdiffusion(path):
    """Read dataset from file (pickle)
    """
    objects = []
    with (open(path, "rb")) as openfile:
        while True:
            try:
                objects.append(pickle.load(openfile))
            except EOFError:
                break
    objects = objects[0]
    x, t, u = objects
    
    return tf.convert_to_tensor(x, dtype=tf.float32), tf.convert_to_tensor(t, dtype=tf.float32), tf.convert_to_tensor(u, dtype=tf.float32)
    


# Synthetic data for the Air Quality example provided in the paper
np.random.seed(42)
def generate_pollutant_data_random(n_samples=80, noise_level=0.12, save=False, valid=False, file_path='./'):
    U, V = 1.0, 1.0  # Wind velocity components
    D = 0.01  # Diffusion coefficient
    S = 0.1  # Source term
    x = np.random.uniform(0, 1, size=(n_samples, 1))
    y = np.random.uniform(0, 1, size=(n_samples, 1))
    t = np.random.uniform(0, 1, size=(n_samples, 1))
    c = np.exp(-((x - U * t) ** 2 + (y - V * t) ** 2) / (4 * D * t )) / np.sqrt(4 * np.pi * D * t )
    noise = np.random.normal(scale=noise_level, size=c.shape)
    c += noise

    if save:   
        if not valid:
            save_to_pickle(file_path +'ST_AirQ_dataset_{}_{}.pkl'.format(n_samples, noise_level), (x, y, t, c))
        else:
            save_to_pickle(file_path +'SV_AirQ_dataset_{}_{}.pkl'.format(n_samples, noise_level), (x, y, t, c))
        
    return tf.convert_to_tensor(x, dtype=tf.float32), tf.convert_to_tensor(y, dtype=tf.float32), tf.convert_to_tensor(t, dtype=tf.float32), tf.convert_to_tensor(c, dtype=tf.float32)


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

# Synthetic data for the Blood Flow and Stenosis example provided in the paper
np.random.seed(42)
## Reference: Modelling Blood Diffusion;
    ## Mt = (2*pi*H*K*Cs / ln(R0/R1))*t
## Reference: https://math.libretexts.org/Bookshelves/Differential_Equations/Differential_Equations_(Chasnov)/09%3A_Partial_Differential_Equations/9.05%3A_Solution_of_the_Diffusion_Equation#:~:text=The%20governing%20equation%20for%20concentration,%2C4%2C6%2C%E2%80%A6.
    ## Here we implement the Example 9.5.1 with its conditions and assumption
    ## which are: Homogeneous one dimensional diffusion in a pipe of length  L
    ## the concentration of a dye in a pipe of length  L, where the dye has total initial mass  M0 and 
    ## is initially concentrated at the center of the pipe, and the ends of the pipe are held at zero concentration.
    ## u(x,t)\approx\frac{2M_0}{L}\sin (\pi x/L)e^{-\pi ^2Dt/L^2}.\nonumber for t>>L^2/D
def generate_drugdiffusion_data_random(n_samples=80, noise_level=0.12, save=False, valid=False, file_path='./'): 
    M0 = 10
    L = 5 #cm ## length of the cylander
    D = 1e-07 #10^-7 cm^2/s (slow diffusion)
    x = np.random.uniform(0, 1, size=(n_samples, 1))
    t = np.random.uniform(0.2, 1, size=(n_samples, 1))

    u = ((2*M0)/L) * np.sin(np.pi * x / L) * np.e**(-(np.pi)**2 * D * t /L**2)

    noise = np.random.normal(scale=noise_level, size=u.shape)
    u += noise

    if save:
        if not valid:
            save_to_pickle(file_path +'ST_DrugDiffusion_dataset_{}_{}.pkl'.format(n_samples, noise_level), (x, t, u))
        else:
            save_to_pickle(file_path +'SV_DrugDiffusion_dataset_{}_{}.pkl'.format(n_samples, noise_level), (x, t, u))
       
    return tf.convert_to_tensor(x, dtype=tf.float32), tf.convert_to_tensor(t, dtype=tf.float32), tf.convert_to_tensor(u, dtype=tf.float32)



# Saving parameters to a text file
def save_parameters(folder, round_num, iteration, params_before, params_after, losses_before, losses_after, Weight):
    if Weight is not None:
        filename = os.path.join(folder, f"Iteration_{iteration}_Round_{round_num}.txt")
    else:
        filename = os.path.join(folder, f"Iteration_{iteration}_Round_{round_num}_Weight_{Weight}.txt")

        
    with open(filename, "a") as file:
        file.write(f"Round {round_num}, iteration {iteration}\n")
        file.write("Parameters Before Averaging:\n")
        for index in range(len(params_before)):
            file.write(f"{index}:\n{params_before[index]}\n")
        file.write("=======================\n")

        file.write("Parameters After Averaging:\n")
        for index in range(len(params_after)):
            file.write(f"{index}:{params_after[index]}\n")
        file.write("=======================\n")

        file.write("Losses Before Averaging:\n")
        for index in range(len(losses_before)):
            file.write(f"{index}:{losses_before[index]}\n")
        file.write("=======================\n")

        file.write("Losses After Averaging:\n")
        for index in range(len(losses_after)):
            file.write(f"{index}:{losses_after[index]}\n")
        file.write("=======================\n")

