gui.py

"""
Protein Location Predictor GUI
This module provides a Tkinter-based GUI for loading FASTA files
     and running protein location prediction tools.
"""
import tkinter as tk
from tkinter import Menu, filedialog, messagebox
import torch
from src.my_utils import predict_with_prost, predict_with_esm


FASTA_FILE_PATH = None  # Global or instance variable

def load_fasta_file():
    """
    Opens a file dialog for the user to select a FASTA file and
      stores the selected file path in a global variable.

    If a file is selected, displays an information message with the file path.
    If no file is selected, displays a warning message.

    Uses:
        - filedialog.askopenfilename for file selection.
        - messagebox.showinfo and messagebox.showwarning for user feedback.

    Global Variables:
        FASTA_FILE_PATH (str): Path to the selected FASTA file.
    """
    global FASTA_FILE_PATH # pylint: disable=global-statement
    FASTA_FILE_PATH = filedialog.askopenfilename(
        filetypes=[("FASTA files", "*.fasta *.fa")],
        title="Select a FASTA file"
    )
    if FASTA_FILE_PATH:
        messagebox.showinfo("File Loaded", f"Loaded file:\n{FASTA_FILE_PATH}")
    else:
        messagebox.showwarning("No file", "No file was selected.")

def run_prost():
    """
    Runs the protein location prediction process.

    Checks if a FASTA file path is provided. If not, displays an error message to the user.
    If a FASTA file is loaded, proceeds to run the prediction using the PROST model.

    Raises:
        Shows a message box error if no FASTA file is loaded.
    """
    if not FASTA_FILE_PATH:
        messagebox.showerror("Error", "Please load a FASTA file first.")
        return
    predict_with_prost(FASTA_FILE_PATH)
    torch.cuda.empty_cache()

def run_esm300():
    """
    Runs the protein location prediction process.

    Checks if a FASTA file path is provided. If not, displays an error message to the user.
    If a FASTA file is loaded, proceeds to run the prediction using the PROST model.

    Raises:
        Shows a message box error if no FASTA file is loaded.
    """
    if not FASTA_FILE_PATH:
        messagebox.showerror("Error", "Please load a FASTA file first.")
        return
    predict_with_esm(fasta_path = FASTA_FILE_PATH,
                     model = 'esmc_300m')
    torch.cuda.empty_cache()

def run_esm600():
    """
    Runs the protein location prediction process.

    Checks if a FASTA file path is provided. If not, displays an error message to the user.
    If a FASTA file is loaded, proceeds to run the prediction using the PROST model.

    Raises:
        Shows a message box error if no FASTA file is loaded.
    """
    if not FASTA_FILE_PATH:
        messagebox.showerror("Error", "Please load a FASTA file first.")
        return
    predict_with_esm(fasta_path = FASTA_FILE_PATH,
                     model = 'esmc_600m')
    torch.cuda.empty_cache()
def menu():
    """
    Displays the main GUI menu for the Protein Tools application.

    This function creates a Tkinter window with a menu bar containing 'File' and 'Help' menus,
    and buttons for running protein prediction tools and exiting the application.

    Menus:
        - File: Options to load a FASTA file or close the application.
        - Help: Options for welcome information and about dialog.

    Buttons:
        - Predict with Prost: Runs the Prost prediction tool.
        - Predict with ESM C: Placeholder for ESM prediction functionality (not yet implemented).
        - Exit: Closes the application.
    """
    # root window
    root = tk.Tk()
    root.geometry('320x200')
    root.title('Protein Location Predictor')

    # create a menubar
    menubar = Menu(root)
    root.config(menu=menubar)

    # create the file_menu
    file_menu = Menu(menubar, tearoff=0)
    file_menu.add_command(label='Load FASTA', command=load_fasta_file)
    file_menu.add_command(label='Close', command=root.quit)
    file_menu.add_separator()
    menubar.add_cascade(label="File", menu=file_menu, underline=0)

    # help menu
    help_menu = Menu(menubar, tearoff=0)
    help_menu.add_command(label='Welcome')
    help_menu.add_command(label='About...')
    menubar.add_cascade(label="Help", menu=help_menu, underline=0)

    btn_prost = tk.Button(root, text="Predict with Prost T5",
                           command=run_prost) #Predict with Prost
    btn_prost.pack(pady=5)

    btn_esm300 = tk.Button(root, text="Predict with ESM C 300m",
                            command = run_esm300) #Predict with ESM C 300m
    btn_esm300.pack(pady=5)

    btn_esm300 = tk.Button(root, text="Predict with ESM C 600m",
                            command = run_esm600) #Predict with ESM C 600m
    btn_esm300.pack(pady=5)

    btn_exit = tk.Button(root, text="Exit", command=root.quit)
    btn_exit.pack(pady=5)

    root.mainloop()

if __name__ == "__main__":
    menu()

# End-of-file (EOF)