import math
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt
import tensorflow as tf
from scipy.integrate import simpson as simps
from sklearn.metrics import mean_squared_error
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
tf.get_logger().setLevel('ERROR')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import time

def lognormal(x, mean, sigma):
    return (1/np.sqrt(2*math.pi*sigma**2)) * np.exp(-(np.log(x) - mean)**2 / (2 * sigma**2))
  
def sum_of_lognormals(x, *weights):
    result = np.zeros(len(x))
    for i in range(num_params):
        result = result + (weights[0][i] * lognormal(x, means[i], sigma))
    return result

def G_prime_fit(w_alpha, g_values):
    sum_over_i = 0.0
    w_squared = w_alpha*w_alpha
    numerator = g_values * tau_values * tau_values
    denom = 1 + (w_squared * tau_values * tau_values)
    to_sum = numerator/denom
    sum_over_i = np.sum(to_sum)
    G_prime_for_alpha = w_squared * sum_over_i
    return G_prime_for_alpha

def G_dub_prime_fit(w_alpha, g_values):
    sum_over_j = 0.0
    w_squared = w_alpha*w_alpha
    numerator = g_values * tau_values
    denom = 1 + (w_squared * tau_values * tau_values)
    to_sum = numerator/denom
    sum_over_j = np.sum(to_sum)
    G_dub_prime_for_alpha = w_alpha * sum_over_j
    return G_dub_prime_for_alpha

def G_concat_fit(w_values, *log_g_values):
    log_g_second_diffs = np.array([])
    for i in range(len(log_g_values)-2):
        g_i_second_diff = log_g_values[i+2] + log_g_values[i] - 2*log_g_values[i+1]
        log_g_second_diffs = np.append(log_g_second_diffs, g_i_second_diff)
    g_values = np.exp(log_g_values)
    G_prime_fit_values = np.array([])
    G_dub_prime_fit_values = np.array([])
    for alpha in w_values:
        G_prime_max = G_prime_fit(alpha, g_values)
        G_prime_fit_values = np.append(G_prime_fit_values, G_prime_max)
        G_dub_prime_max = G_dub_prime_fit(alpha, g_values)
        G_dub_prime_fit_values = np.append(G_dub_prime_fit_values, G_dub_prime_max)
    G_concat = np.concatenate((G_prime_fit_values, G_dub_prime_fit_values), axis=0) 
    G_log_concat_diff = (np.log(G_concat) - G_concat_inp_LN) 
    G_concat_with_diffs = np.concatenate((G_log_concat_diff, lam*log_g_second_diffs), axis = 0)
    return G_concat_with_diffs

## a-PS @180C
M_K=720.0
M_e=12870.0
G0_180=220000.0
tau_e_180=2.20e-4
tau_G_180=1.30e-9
G_G=1.20e9
beta_G=0.390

M_e_PE = 820.0
samples_list=[]
print("All samples available:\n")
rheo_data_dir = "Rheo_Data"
files = os.listdir(rheo_data_dir)
for file in files:
    if os.path.isfile(os.path.join(rheo_data_dir, file)):  
        first_part = file.split('_')[0]  
        print(first_part)
        samples_list.append(first_part)
sample = input("\nEnter sample name from list above: ").upper()
if sample not in samples_list:
    while sample not in samples_list:
        print("Invalid sample name. Please enter a valid sample name from the list.")
        sample = input("Enter sample name from list above: ").upper()
print(f"Making predictions for {sample}")
best_batch, best_model = 2,3 
worst_batch, worst_model = 3,1

# save_figs = True
save_figs = False


file_path = os.path.join("GPC_Data", f"{sample}_GPC.dat")
data = np.loadtxt(file_path, delimiter='\t')
x = np.logspace(2, 7, 200)
x_data = data[:, 0]
y_data = data[:, 1]
y_data = y_data / simps(y_data, x=np.log(x_data))

fig, ax = plt.subplots()
ax.plot(x_data, y_data, label='Original Data', linewidth=2)
plt.xscale('log')
ax.set_xlabel('M (g/mol)',fontsize=14)
ax.set_ylabel('dW/dlogM',fontsize=14)
plt.yticks(fontsize=14)
plt.xticks(fontsize=14)
ax.set_title(f"{sample} Input GPC MWD to Predict")
ax.set_xlim(min(x_data), max(x_data))
plt.tight_layout()
plt.draw()
plt.pause(0.1)

pred_auc = simps(y_data, x=np.log(x_data))
num_params = 28
newlam_value = 1


means = np.linspace(np.log(0.1*M_e_PE), np.log(10000*M_e_PE), num_params) 
known_means = np.linspace(np.log(10*M_e_PE), np.log(1000*M_e_PE), 7) 
means_ratio = np.exp(means[1])/np.exp(means[0])
known_means_ratio = np.exp(known_means[1])/np.exp(known_means[0])
sigma = 0.55 * (means_ratio/known_means_ratio)

longtau = 4
smalltau = -6
numoftau = 1 + (longtau-smalltau) * 8
tau_values = np.logspace(longtau,smalltau,numoftau,base=10)
tau_values_Rept = tau_values

unsorted_Gp_data_Rept = np.array([])
unsorted_Gpp_data_Rept = np.array([])
unsorted_w_values_Rept = np.array([])

with open(os.path.join("Rheo_Data", f"{sample}_rheo.dat"), 'r') as file:
    for i,line in enumerate(file):
        values = line.split()
        if i == 0:
            T = float(values[0])
            continue
        value1 = float(values[0]) 
        value2 = float(values[1])
        value3 = float(values[2])
        unsorted_w_values_Rept = np.append(unsorted_w_values_Rept,value1)
        unsorted_Gp_data_Rept = np.append(unsorted_Gp_data_Rept,value2) 
        unsorted_Gpp_data_Rept = np.append(unsorted_Gpp_data_Rept,value3) 

print(f"Rheology temperature: {T}C")
B1 = 651.9
B2 = -52.24
log10alpha = -3.161
alpha = 10**log10alpha
T_r = 180

log10alpha_T = (-B1*(T-T_r))/((B2+T_r)*(B2+T))
alpha_T = 10**log10alpha_T
b_T = ((1+alpha*T)*(T_r+273.15))/((1+alpha*T_r)*(T+273.15))

sorted_indices = np.argsort(unsorted_w_values_Rept)
w_values_Rept = unsorted_w_values_Rept[sorted_indices] * alpha_T 
Gp_data_Rept = unsorted_Gp_data_Rept[sorted_indices] / b_T
Gpp_data_Rept = unsorted_Gpp_data_Rept[sorted_indices] / b_T

nat_log_w_values_Rept = np.log(w_values_Rept)
log10_w_values_Rept = np.log10(w_values_Rept)
G_concat_inp_nat_Rept = np.concatenate((Gp_data_Rept, Gpp_data_Rept))

G_len_zeros = np.zeros(2*np.shape(w_values_Rept)[0] + len(tau_values_Rept) - 2)

N = len(w_values_Rept)
p_w = N / (max(np.log10(w_values_Rept))-min(np.log10(w_values_Rept)))
p_w_base = 95 / (14)
lam = newlam_value * (p_w/p_w_base)**0.5

numomega = np.shape(w_values_Rept)[0]
bounds = ([-40] * numoftau, [50]*numoftau)


    
G_concat_inp_Rept = np.log10(G_concat_inp_nat_Rept)
G_concat_inp_LN_Rept = np.log(G_concat_inp_nat_Rept)
G_concat_inp_LN = G_concat_inp_LN_Rept 

if numoftau < numomega:
    G_split = np.array_split(G_concat_inp_LN_Rept[0:int(numomega)], numoftau)
    G_dub_split = np.array_split(G_concat_inp_LN_Rept[int(numomega):], numoftau)
    log_prime_means = np.array([np.mean(subarray) for subarray in G_split])
    log_dub_means = np.array([np.mean(subarray) for subarray in G_dub_split])
    initial_guess = (log_prime_means + log_dub_means) / 2  - 2
elif numoftau >= numomega:
    desired_size = numoftau        
    interpolated_indices = np.linspace(0, G_concat_inp_Rept.shape[0] / 2 - 1, desired_size)
    interpolated_prime = np.interp(interpolated_indices, np.arange(G_concat_inp_Rept.shape[0]/2), G_concat_inp_LN_Rept[0:int(numomega)])
    interpolated_dub = np.interp(interpolated_indices, np.arange(G_concat_inp_Rept.shape[0]/2), G_concat_inp_LN_Rept[int(numomega):])
    initial_guess = (interpolated_prime + interpolated_dub) / 2 - 2

optimiseresult, pcov = opt.curve_fit(G_concat_fit, w_values_Rept, G_len_zeros, p0=initial_guess, bounds=bounds, method='trf')
optimiseresult_Rept = optimiseresult    

pred_G_prime = np.array([])
pred_G_dub_prime = np.array([])

for alpha in w_values_Rept:
    predicted_Gprime = G_prime_fit(alpha, np.exp(optimiseresult))
    pred_G_prime = np.append(pred_G_prime, predicted_Gprime)
    predicted_Gdubprime = G_dub_prime_fit(alpha, np.exp(optimiseresult))
    pred_G_dub_prime = np.append(pred_G_dub_prime, predicted_Gdubprime)

fig, ax = plt.subplots()
plt.title("Input Rheology Data for PS Sample")
ax.scatter(w_values_Rept, Gp_data_Rept, facecolors='none', edgecolors='orange', linewidth = 0.5, label = "G' input", s = 10)
ax.plot(w_values_Rept, pred_G_prime, '-', color='orange', label = "G' Maxwell Fit", linewidth=1)
ax.scatter(w_values_Rept, Gpp_data_Rept, facecolors='none', edgecolors='blue', linewidth = 0.5, label = "G'' input", s = 10)
ax.plot(w_values_Rept, pred_G_dub_prime, '-', color='blue', label = "G'' Maxwell Fit", linewidth=1)
ax.plot((1/tau_values_Rept), np.exp(optimiseresult), ".", c='green', label="Rept Maxwell modes")
plt.xlabel(r'$\omega$ (rad/s)', fontsize = 14)
plt.ylabel("G', G'' (Pa)", fontsize = 14)
plt.yticks(fontsize = 14)
plt.xticks(fontsize = 14)
plt.xscale('log')
plt.yscale('log')
plt.legend(frameon=False)
plt.tight_layout()
plt.draw()
plt.pause(0.1)

X_val = np.log10(np.exp(np.stack((optimiseresult, optimiseresult))))
X_val = np.concatenate((X_val, X_val), axis=1)
X_val = np.reshape(X_val, (2,2,len(optimiseresult)))

RMSE_list = []
models=[]
model_count = 0

num_of_models = 9

preds = np.zeros((num_of_models,len(x)))

for i in range(1,num_of_models+1):
    
    model = tf.keras.models.load_model(os.path.join("NN_Models", f"PS_Model_{i}.keras"))
    
    prediction = abs(model.predict(X_val))
    
    pred_MWD = sum_of_lognormals(x, prediction[0])
    pred_MWD = pred_MWD / simps(pred_MWD, x = np.log(x))
    preds[i-1] = pred_MWD
    pred_auc = simps(pred_MWD, x=np.log(x))
        
    fig, ax = plt.subplots()
    plt.xlabel('M (g/mol)',fontsize=14)
    plt.ylabel('dW/dlogM',fontsize=14)
    plt.yticks(fontsize=14)
    plt.xticks(fontsize=14)
    ax.plot(x_data, y_data, c='blue', label="GPC Data")
    ax.plot(x, pred_MWD, '--', c='orange', label="Predicted MWD", linewidth=3)
    ax.set_xscale('log')
    ax.set_xlim(min(x), max(x))
    ax.legend(bbox_to_anchor=(0.38, 1), frameon=False)
    plt.title(f"{sample} Predicted MWD for model {i}")
    
    for k in range(num_params):
        individual_pred_curve = prediction[0, k] * lognormal(x, means[k], sigma)
        ax.plot(x, individual_pred_curve, label=None, linestyle='--', color=plt.cm.viridis(k / len(means)))
    if save_figs == True:
        if i == best_model:
            plt.savefig(f"{sample}_best_{i}.png", bbox_inches='tight', dpi=5e+2)
        if i == worst_model:
            plt.savefig(f"{sample}_worst_{i}.png", bbox_inches='tight', dpi=5e+2)
    
    plt.tight_layout()
    plt.draw()
    plt.pause(0.1)
    y_interpolated = np.interp(x, x_data, y_data, left=0, right=0)
    
    RMSE_MWD= np.sqrt(mean_squared_error(y_interpolated, pred_MWD))
    print(f"RMSE for model {i}: {RMSE_MWD}")
    RMSE_list.append(RMSE_MWD)
    time.sleep(0.1)
        
print(f"\nMean RMSE = {np.mean(RMSE_list):.4f}\n")
        
min_error, min_model = min(RMSE_list), np.argmin(RMSE_list)+1
max_error, max_model = max(RMSE_list), np.argmax(RMSE_list)+1

print(f"Model {min_model} had lowest error = {min_error:.4f}")
print(f"Model {max_model} had largest error = {max_error:.4f}")

fig, ax = plt.subplots()
plt.xlabel('M (g/mol)',fontsize=14)
plt.ylabel('dW/dlogM',fontsize=14)
plt.yticks(fontsize=14)
plt.xticks(fontsize=14)
ax.set_xscale('log')
ax.set_xlim(min(x), max(x))
ax.set_xlim(1000, 1e+7)
ax.plot(x_data, y_data, c='blue', label="GPC Data")

for i in range(1,num_of_models+1):        
    pred_MWD = preds[i-1]
    if i==1:
        ax.plot(x, pred_MWD, '--', c='orange',alpha=0.8, linewidth=1, label="All Predicted MWDs")
    else:
        ax.plot(x, pred_MWD, '--', c='orange',alpha=0.8, linewidth=1)

plt.title(f"{sample} Predicted MWDs for all models")
plt.tight_layout()        

ax.legend(bbox_to_anchor=(0.44, 1), frameon=False)
plt.show()        
stddevs = []
for i in range(len(preds[0])):
    stddev = np.sqrt(np.sum((preds[:,i]-np.mean(preds[:,i]))**2)/9)
    stddevs.append(stddev)
mean_std_dev = np.mean(stddevs)  
  
print(f"Mean standard deviation = {np.mean(stddevs):.5f}")