interface_GDA_fitting.py

import os
from datetime import datetime

import tkinter as tk
from tkinter import filedialog, messagebox
import numpy as np
from scipy.optimize import brentq, minimize
import matplotlib.pyplot as plt

from pltstyle import create_plots

def format_value(value):
    return f"{value:.0f}" if value > 10 else f"{value:.2f}"
    
# Function to run the fitting process
def run_fitting(file_path, results_file_path, Kd_in_M, h0_in_M, g0_in_M, number_of_fit_trials, rmse_threshold_factor, r2_threshold, save_plots, display_plots, plots_dir, save_results, results_save_dir):
    try:
        # Load data from input file
        data_lines = load_data(file_path)
        if data_lines is None:
            raise ValueError("Data loading failed.")
        replicas = split_replicas(data_lines)
        if replicas is None:
            raise ValueError("Replica splitting failed.")
        print(f"Number of replicas detected: {len(replicas)}")

        # Initialize parameter ranges for fitting
        I0_range = (0, None)
        Id_range = (None, None)
        Ihd_range = (None, None)

        # Check and load boundaries from results file, if available
        if results_file_path and os.path.exists(results_file_path):
            try:
                with open(results_file_path, 'r') as f:
                    lines = f.readlines()

                # Extract Id prediction interval
                id_prediction_line = next((line for line in lines if 'Id prediction interval' in line), None)
                if id_prediction_line and 'not applicable' not in id_prediction_line:
                    Id_lower = float(id_prediction_line.split('[')[-1].split(',')[0].strip())
                    Id_upper = float(id_prediction_line.split(',')[-1].split(']')[0].strip())
                else:
                    average_Id = float(next(line for line in lines if 'Average Id' in line).split('\t')[-1].strip())
                    Id_lower = 0.5 * average_Id
                    Id_upper = 2.0 * average_Id

                # Extract I0 prediction interval
                i0_prediction_line = next((line for line in lines if 'I0 prediction interval' in line), None)
                if i0_prediction_line and 'not applicable' not in i0_prediction_line:
                    I0_lower = float(i0_prediction_line.split('[')[-1].split(',')[0].strip())
                    I0_upper = float(i0_prediction_line.split(',')[-1].split(']')[0].strip())
                else:
                    average_I0 = float(next(line for line in lines if 'Average I0' in line).split('\t')[-1].strip())
                    I0_lower = 0.5 * average_I0
                    I0_upper = 2.0 * average_I0

            except Exception as e:
                print(f"Error parsing boundaries from the results file: {e}")
                Id_lower, Id_upper = 1e3, 1e18
                I0_lower, I0_upper = 0, None
        else:
            Id_lower, Id_upper = 1e3, 1e18
            I0_lower, I0_upper = 0, None

        # Convert bounds to µM⁻¹ for fitting
        Id_lower /= 1e6
        Id_upper /= 1e6
        Ihd_lower = Ihd_range[0] / 1e6 if Ihd_range[0] is not None else 0.001
        Ihd_upper = Ihd_range[1] / 1e6 if Ihd_range[1] is not None else 1e12

        # Convert constants to µM
        Kd = Kd_in_M / 1e6
        h0 = h0_in_M * 1e6
        g0 = g0_in_M * 1e6

        figures = []  # List to store figures

        # Process each replica for fitting
        for replica_index, replica_data in enumerate(replicas, start=1):
            d0_values = replica_data[:, 0] * 1e6
            Signal_observed = replica_data[:, 1]

            if len(d0_values) < 2:
                print(f"Replica {replica_index} has insufficient data. Skipping.")
                continue

            # Update I0_upper if needed
            I0_upper = np.min(Signal_observed) if I0_upper is None or np.isinf(I0_upper) else I0_upper

            # Generate initial guesses for parameters
            Ihd_guess_smaller = Signal_observed[0] < Signal_observed[-1]
            initial_params_list = []
            for _ in range(number_of_fit_trials):
                I0_guess = np.random.uniform(I0_lower, I0_upper)
                Kg_guess = 10 ** np.random.uniform(np.log10(Kd) - 5, np.log10(Kd) + 5)
                if Ihd_guess_smaller:
                    Id_guess = 10 ** np.random.uniform(np.log10(Id_lower), np.log10(Id_upper))
                    Ihd_guess = Id_guess * np.random.uniform(0.1, 0.5)
                else:
                    Ihd_guess = 10 ** np.random.uniform(np.log10(Ihd_lower), np.log10(Ihd_upper))
                    Id_guess = Ihd_guess * np.random.uniform(0.1, 0.5)
                initial_params_list.append([I0_guess, Kg_guess, Id_guess, Ihd_guess])

            # Fit parameters for replica using least-squares minimization
            best_result, best_cost = None, np.inf
            fit_results = []
            for initial_params in initial_params_list:
                result = minimize(lambda params: np.sum(residuals(params, d0_values, Signal_observed, Kd, h0, g0) ** 2),
                                  initial_params, method='L-BFGS-B',
                                  bounds=[(I0_lower, I0_upper), (1e-8, 1e8), (Id_lower, Id_upper), (Ihd_lower, Ihd_upper)])
                Signal_computed = compute_signal(result.x, d0_values, Kd, h0, g0)
                rmse, r_squared = calculate_fit_metrics(Signal_observed, Signal_computed)
                fit_results.append((result.x, result.fun, rmse, r_squared))

                if result.fun < best_cost:
                    best_cost = result.fun
                    best_result = result

            # Filter fits by RMSE and R² thresholds
            best_rmse = min(fit_rmse for _, _, fit_rmse, _ in fit_results)
            rmse_threshold = best_rmse * rmse_threshold_factor

            filtered_results = [
                (params, fit_rmse, fit_r2) for params, _, fit_rmse, fit_r2 in fit_results
                if fit_rmse <= rmse_threshold and fit_r2 >= r2_threshold
            ]

            # Calculate median parameters if valid results are found
            if filtered_results:
                median_params = np.median(np.array([result[0] for result in filtered_results]), axis=0)
            else:
                print("Warning: No fits meet the filtering criteria.")
                continue

            # Compute metrics for median fit
            Signal_computed = compute_signal(median_params, d0_values, Kd, h0, g0)
            rmse, r_squared = calculate_fit_metrics(Signal_observed, Signal_computed)

            # Plot observed vs. simulated fitting curve
            fitting_curve_x, fitting_curve_y = [], []
            for i in range(len(d0_values) - 1):
                extra_points = np.linspace(d0_values[i], d0_values[i + 1], 21)
                fitting_curve_x.extend(extra_points)
                fitting_curve_y.extend(compute_signal(median_params, extra_points, Kd, h0, g0))
            # TODO: implement a "dynamic" unit and scale for the x-axis
            fig, ax = create_plots(x_label=r'$D_0$ $\rm{[\mu M]}$', y_label=r'Signal $\rm{[AU]}$')

            ax.plot(d0_values, Signal_observed, 'o', label='Observed Signal')
            ax.plot(fitting_curve_x, fitting_curve_y, '--', color='blue', alpha=0.6, label='Simulated Fitting Curve')
            ax.set_title(f'Observed vs. Simulated Fitting Curve for Replica {replica_index}')
            ax.legend(loc='upper left', bbox_to_anchor=(0.02, 0.98))
            
            # Annotate plot with median parameter values and fit metrics
            param_text = (f"$K_g$: {median_params[1] * 1e6:.2e} $M^{{-1}}$\n"
                        f"$I_0$: {median_params[0]:.2e}\n"
                        f"$I_d$: {median_params[2] * 1e6:.2e} $M^{{-1}}$\n"
                        f"$I_{{hd}}$: {median_params[3] * 1e6:.2e} $M^{{-1}}$\n"
                        f"$RMSE$: {format_value(rmse)}\n"
                        f"$R^2$: {r_squared:.3f}")
            
            ax.annotate(param_text, xy=(0.8, 0.04), xycoords='axes fraction', fontsize=10,
                ha='left', va='bottom', bbox=dict(boxstyle="round,pad=0.3", edgecolor="black", facecolor="lightgrey", alpha=0.5))

            if save_plots:
                plot_file = os.path.join(plots_dir, f"fit_plot_replica_{replica_index}.png")
                fig.savefig(plot_file, bbox_inches='tight')
                print(f"Plot saved to {plot_file}")

            figures.append(fig)  # Store the figure

            # Save results to a file if save_results is True
            if save_results:
                replica_file = os.path.join(results_save_dir, f"fit_results_replica_{replica_index}.txt")
                with open(replica_file, 'w') as f:
                    # Input section
                    f.write("Input:\n")
                    f.write(f"g0 (M): {g0_in_M:.6e}\n")
                    f.write(f"h0 (M): {h0_in_M:.6e}\n")
                    f.write(f"Kd (M^-1): {Kd_in_M:.6e}\n")
                    f.write(f"Id lower bound (signal/M): {Id_lower * 1e6:.3e}\n")
                    f.write(f"Id upper bound (signal/M): {Id_upper * 1e6:.3e}\n")
                    f.write(f"I0 lower bound: {I0_lower:.3e}\n")
                    f.write(f"I0 upper bound: {I0_upper:.3e}\n")
                
                    # Output - Median parameters
                    f.write("\nOutput:\nMedian parameters:\n")
                    f.write(f"Kg (M^-1): {median_params[1] * 1e6:.2e}\n")
                    f.write(f"I0: {median_params[0]:.2e}\n")
                    f.write(f"Id (signal/M): {median_params[2] * 1e6:.2e}\n")
                    f.write(f"Ihd (signal/M): {median_params[3] * 1e6:.2e}\n")
                    f.write(f"RMSE: {rmse:.3f}\n")
                    f.write(f"R²: {r_squared:.3f}\n")
                
                    # Acceptable Fit Parameters
                    f.write("\nAcceptable Fit Parameters:\n")
                    f.write("Kg (M^-1)\tI0\tId (signal/M)\tIhd (signal/M)\tRMSE\tR²\n")
                    for params, fit_rmse, fit_r2 in filtered_results:
                        f.write(f"{params[1] * 1e6:.2e}\t{params[0]:.2e}\t{params[2] * 1e6:.2e}\t{params[3] * 1e6:.2e}\t{fit_rmse:.3f}\t{fit_r2:.3f}\n")
                
                    # Calculate standard deviations for Kg, I0, Id, and Ihd if there are filtered results
                    if filtered_results:
                        Kg_values = [params[1] * 1e6 for params, _, _ in filtered_results]
                        I0_values = [params[0] for params, _, _ in filtered_results]
                        Id_values = [params[2] * 1e6 for params, _, _ in filtered_results]
                        Ihd_values = [params[3] * 1e6 for params, _, _ in filtered_results]
                
                        Kg_std = np.std(Kg_values)
                        I0_std = np.std(I0_values)
                        Id_std = np.std(Id_values)
                        Ihd_std = np.std(Ihd_values)
                    else:
                        Kg_std = I0_std = Id_std = Ihd_std = np.nan  # Assign NaN if no filtered results
                
                    # Standard Deviations section
                    f.write("\nStandard Deviations:\n")
                    f.write(f"K_g Std Dev (M-1): {Kg_std:.2e}\n")
                    f.write(f"I_0 Std Dev: {I0_std:.2e}\n")
                    f.write(f"I_d Std Dev (signal/M): {Id_std:.2e}\n")
                    f.write(f"I_hd Std Dev (signal/M): {Ihd_std:.2e}\n")
                    
                    # Original Data section
                    f.write("\nOriginal Data:\nConcentration do (M)\tSignal\n")
                    for d0, signal in zip(d0_values / 1e6, Signal_observed):  
                        f.write(f"{d0:.6e}\t{signal:.6e}\n")
                    
                    # Fitting Curve section
                    f.write("\nFitting Curve:\n")
                    f.write("Simulated Concentration (M)\tSimulated Signal\n")
                    for x_sim, y_sim in zip(np.array(fitting_curve_x) / 1e6, fitting_curve_y):
                        f.write(f"{x_sim:.6e}\t{y_sim:.6e}\n")
                    
                    # Date of Export
                    f.write(f"\nDate of Export: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                    print(f"Results for Replica {replica_index} saved to {replica_file}")

        # Display all figures at once if display_plots is True
        if display_plots:
            plt.show()

    except Exception as e:
        messagebox.showerror("Error", str(e))

# Function to load data from file
def load_data(file_path):
    try:
        with open(file_path, 'r') as f:
            return f.readlines()
    except Exception as e:
        print(f"Error loading file: {e}")
        return None

# Function to split data into replicas
def split_replicas(data):
    if data is None:
        return None
    replicas, current_replica = [], []
    use_var_signal_split = any("var signal" in line.lower() for line in data)

    for line in data:
        if "var" in line.lower():
            if current_replica:
                replicas.append(np.array(current_replica))
                current_replica = []
        else:
            try:
                x, y = map(float, line.split())
                if use_var_signal_split:
                    current_replica.append((x, y))
                else:
                    if x == 0.0 and current_replica:
                        replicas.append(np.array(current_replica))
                        current_replica = []
                    current_replica.append((x, y))
            except ValueError:
                continue
    if current_replica:
        replicas.append(np.array(current_replica))
    return replicas if replicas else None

# Function to compute signal
def compute_signal(params, d0_values, Kd, h0, g0):
    I0, Kg, Id, Ihd = params
    Signal_values = []
    for d0 in d0_values:
        try:
            def equation_h(h):
                denom_Kd = 1 + Kd * h
                denom_Kg = 1 + Kg * h
                h_d = (Kd * h * d0) / denom_Kd
                h_g = (Kg * h * g0) / denom_Kg
                return h + h_d + h_g - h0

            h_sol = brentq(equation_h, 1e-20, h0, xtol=1e-14, maxiter=1000)
            denom_Kd = 1 + Kd * h_sol
            d_free = d0 / denom_Kd
            h_d = Kd * h_sol * d_free
            Signal = I0 + Id * d_free + Ihd * h_d
            Signal_values.append(Signal)
        except Exception:
            Signal_values.append(np.nan)
    return np.array(Signal_values)

# Function to compute residuals
def residuals(params, d0_values, Signal_observed, Kd, h0, g0):
    Signal_computed = compute_signal(params, d0_values, Kd, h0, g0)
    return np.nan_to_num(Signal_observed - Signal_computed, nan=1e6)

# Function to calculate fit metrics
def calculate_fit_metrics(Signal_observed, Signal_computed):
    rmse = np.sqrt(np.nanmean((Signal_observed - Signal_computed) ** 2))
    ss_res = np.nansum((Signal_observed - Signal_computed) ** 2)
    ss_tot = np.nansum((Signal_observed - np.nanmean(Signal_observed)) ** 2)
    r_squared = 1 - (ss_res / ss_tot) if ss_tot > 0 else np.nan
    return rmse, r_squared

# Function to browse file
class GDAFittingApp:
    def __init__(self, root):
        self.root = root
        self.root.title("GDA Fitting Interface")
        self.info_label = None
        
        # Variables
        self.file_path_var = tk.StringVar()
        self.use_results_file_var = tk.BooleanVar()
        self.results_file_path_var = tk.StringVar()
        self.Kd_var = tk.DoubleVar()
        self.h0_var = tk.DoubleVar()
        self.g0_var = tk.DoubleVar()
        self.fit_trials_var = tk.IntVar()
        self.rmse_threshold_var = tk.DoubleVar()
        self.r2_threshold_var = tk.DoubleVar()
        self.save_plots_var = tk.BooleanVar()
        self.results_dir_var = tk.StringVar()
        self.display_plots_var = tk.BooleanVar()
        self.save_results_var = tk.BooleanVar()  # New variable for saving results
        self.results_save_dir_var = tk.StringVar()  # New variable for results save directory
        
        # Set default values
        self.Kd_var.set(1.68e7)  # Binding constant for h_d binding in M^-1
        self.h0_var.set(4.3e-6)  # Initial host concentration (M)
        self.g0_var.set(6e-6)    # Initial guest concentration (M)
        self.fit_trials_var.set(200)  # Number of fit trials
        self.rmse_threshold_var.set(2)  # RMSE threshold factor
        self.r2_threshold_var.set(0.9)  # R² threshold
        self.display_plots_var.set(True)
        
        # Padding
        pad_x = 10
        pad_y = 5

        # Widgets
        tk.Label(self.root, text="Input File Path:").grid(row=0, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.file_path_entry = tk.Entry(self.root, textvariable=self.file_path_var, width=40, justify='left')
        self.file_path_entry.grid(row=0, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)
        tk.Button(self.root, text="Browse", command=self.browse_input_file).grid(row=0, column=2, padx=pad_x, pady=pad_y)

        tk.Checkbutton(self.root, text="Read Boundaries from File: ", variable=self.use_results_file_var, command=self.update_use_results_widgets).grid(row=1, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.results_file_path_entry = tk.Entry(self.root, textvariable=self.results_file_path_var, width=40, justify='left', state=tk.DISABLED)
        self.results_file_path_entry.grid(row=1, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)
        self.results_file_button = tk.Button(self.root, text="Browse", command=lambda: self.browse_file(self.results_file_path_entry), state=tk.DISABLED)
        self.results_file_button.grid(row=1, column=2, padx=pad_x, pady=pad_y)
        self.use_results_file_var.trace_add('write', lambda *args: self.update_use_results_widgets())

        tk.Label(self.root, text=r"Kₐ (M⁻¹):").grid(row=3, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.Kd_entry = tk.Entry(self.root, textvariable=self.Kd_var, justify='left')
        self.Kd_entry.grid(row=3, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)

        tk.Label(self.root, text="H₀ (M):").grid(row=4, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.h0_entry = tk.Entry(self.root, textvariable=self.h0_var, justify='left')
        self.h0_entry.grid(row=4, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)

        tk.Label(self.root, text="G₀ (M):").grid(row=5, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.g0_entry = tk.Entry(self.root, textvariable=self.g0_var, justify='left')
        self.g0_entry.grid(row=5, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)

        tk.Label(self.root, text="Number of Fit Trials:").grid(row=6, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.fit_trials_entry = tk.Entry(self.root, textvariable=self.fit_trials_var, justify='left')
        self.fit_trials_entry.grid(row=6, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)

        tk.Label(self.root, text="RMSE Threshold Factor:").grid(row=7, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.rmse_threshold_entry = tk.Entry(self.root, textvariable=self.rmse_threshold_var, justify='left')
        self.rmse_threshold_entry.grid(row=7, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)

        tk.Label(self.root, text="R² Threshold:").grid(row=8, column=0, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.r2_threshold_entry = tk.Entry(self.root, textvariable=self.r2_threshold_var, justify='left')
        self.r2_threshold_entry.grid(row=8, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)

        tk.Checkbutton(self.root, text="Save Plots To", variable=self.save_plots_var, command=self.update_save_plot_widgets).grid(row=9, column=0, columnspan=1, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.results_dir_entry = tk.Entry(self.root, textvariable=self.results_dir_var, width=40, state=tk.DISABLED, justify='left')
        self.results_dir_entry.grid(row=9, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)
        self.results_dir_button = tk.Button(self.root, text="Browse", command=lambda: self.browse_directory(self.results_dir_entry), state=tk.DISABLED)
        self.results_dir_button.grid(row=9, column=2, padx=pad_x, pady=pad_y)

        self.save_plots_var.trace_add('write', lambda *args: self.update_save_plot_widgets())

        tk.Checkbutton(self.root, text="Save Results To", variable=self.save_results_var, command=self.update_save_results_widgets).grid(row=10, column=0, columnspan=1, sticky=tk.W, padx=pad_x, pady=pad_y)
        self.results_save_dir_entry = tk.Entry(self.root, textvariable=self.results_save_dir_var, width=40, state=tk.DISABLED, justify='left')
        self.results_save_dir_entry.grid(row=10, column=1, padx=pad_x, pady=pad_y, sticky=tk.W)
        self.results_save_dir_button = tk.Button(self.root, text="Browse", command=lambda: self.browse_directory(self.results_save_dir_entry), state=tk.DISABLED)
        self.results_save_dir_button.grid(row=10, column=2, padx=pad_x, pady=pad_y)

        self.save_results_var.trace_add('write', lambda *args: self.update_save_results_widgets())

        tk.Checkbutton(self.root, text="Display Plots", variable=self.display_plots_var).grid(row=11, column=0, columnspan=3, sticky=tk.W, padx=pad_x, pady=pad_y)
        tk.Button(self.root, text="Run Fitting", command=self.run_fitting).grid(row=12, column=0, columnspan=3, pady=10, padx=pad_x)

    def browse_input_file(self):
        file_path = filedialog.askopenfilename()
        if file_path:
            self.file_path_var.set(file_path)
            root_dir = os.path.dirname(file_path)
            self.results_dir_var.set(root_dir)
            self.results_save_dir_var.set(root_dir)

    def browse_file(self, entry):
        initial_dir = os.path.dirname(self.file_path_var.get()) if self.file_path_var.get() else os.getcwd()
        file_path = filedialog.askopenfilename(initialdir=initial_dir)
        if file_path:
            entry.delete(0, tk.END)
            entry.insert(0, file_path)

    def browse_directory(self, entry):
        initial_dir = os.path.dirname(self.file_path_var.get()) if self.file_path_var.get() else os.getcwd()
        directory_path = filedialog.askdirectory(initialdir=initial_dir)
        if directory_path:
            entry.delete(0, tk.END)
            entry.insert(0, directory_path)

    def update_use_results_widgets(self):
        state = tk.NORMAL if self.use_results_file_var.get() else tk.DISABLED
        self.results_file_path_entry.config(state=state)
        self.results_file_button.config(state=state)

    def update_save_plot_widgets(self):
        state = tk.NORMAL if self.save_plots_var.get() else tk.DISABLED
        self.results_dir_entry.config(state=state)
        self.results_dir_button.config(state=state)

    def update_save_results_widgets(self):
        state = tk.NORMAL if self.save_results_var.get() else tk.DISABLED
        self.results_save_dir_entry.config(state=state)
        self.results_save_dir_button.config(state=state)
        
    def show_message(self, message, is_error=False):
        if self.info_label:
            self.info_label.destroy()
        fg_color = 'red' if is_error else 'green'
        self.info_label = tk.Label(self.root, text=message, fg=fg_color)
        self.info_label.grid(row=13, column=0, columnspan=3, pady=10)
        
    def run_fitting(self):
        # Adjust the run_fitting function to access self variables and implement the fitting logic
        try:
            # Get user inputs
            file_path = self.file_path_entry.get()
            results_file_path = self.results_file_path_entry.get() if self.use_results_file_var.get() else None
            Kd_in_M = self.Kd_var.get()
            h0_in_M = self.h0_var.get()
            g0_in_M = self.g0_var.get()
            number_of_fit_trials = self.fit_trials_var.get()
            rmse_threshold_factor = self.rmse_threshold_var.get()
            r2_threshold = self.r2_threshold_var.get()
            save_plots = self.save_plots_var.get()
            display_plots = self.display_plots_var.get()
            plots_dir = self.results_dir_entry.get()
            save_results = self.save_results_var.get()
            results_save_dir = self.results_save_dir_entry.get()
            
            # Run the fitting process
            run_fitting(file_path, results_file_path, Kd_in_M, h0_in_M, g0_in_M, number_of_fit_trials, rmse_threshold_factor, r2_threshold, save_plots, display_plots, plots_dir, save_results, results_save_dir)
            
            self.show_message(f"Fitting completed!", is_error=False)
            
        except Exception as e:
            self.show_message(f"Error: {str(e)}", is_error=True)

# Main function to run the GUI
if __name__ == "__main__":
    root = tk.Tk()
    root.title("Automation Project")

    GDAFittingApp(root)

    root.mainloop()