main.py

import os
import sys
import customtkinter as ctk
from tkinter import messagebox, filedialog
import threading
import subprocess

# Import app components
from components.sidebar import SidebarFrame
from components.main_panel import MainPanel
from components.results_panel import ResultsPanel
from components.database_manager import DatabaseManager
from utils.config import AppConfig

class RMFSamplingApp(ctk.CTk):
    def __init__(self):
        super().__init__()
        
        # Configure the app
        self.title("CompactObject CEDF EOS Database")
        self.geometry("1200x800")
        
        # Load configuration
        self.config = AppConfig()
        
        # Set CustomTkinter appearance mode and color theme
        ctk.set_appearance_mode(self.config.get("appearance_mode", "System"))
        ctk.set_default_color_theme(self.config.get("color_theme", "blue"))
        ctk.set_default_color_theme(os.path.join("Classic.json"))
        
        # Create database manager
        self.database_manager = DatabaseManager(self.config)
        
        # Configure grid layout (3x1)
        self.grid_columnconfigure(1, weight=1)
        self.grid_rowconfigure(0, weight=1)
        
        # Initialize components
        self.initialize_components()
        
        # Check if database exists
        self.check_database()
    
    def initialize_components(self):
        # Create sidebar
        self.sidebar = SidebarFrame(self, self.config, command_callback=self.run_command)
        self.sidebar.grid(row=0, column=0, sticky="nsew")
        
        # Create main panel
        self.main_panel = MainPanel(self, self.config)
        self.main_panel.grid(row=0, column=1, sticky="nsew", padx=20, pady=20)
        
        # Create results panel
        self.results_panel = ResultsPanel(self, self.config)
        self.results_panel.grid(row=0, column=2, sticky="nsew")
        
        # Add manual initialization button in case automatic fails
        initialize_btn = ctk.CTkButton(
            self,
            text="Manual Initialize",
            command=self.manual_initialize
        )
        initialize_btn.grid(row=1, column=1, padx=20, pady=10)
    
    def manual_initialize(self):
        """Force initialization of models and quantities"""
        # Show output message
        self.results_panel.clear_output()
        self.results_panel.add_output("Manually initializing...\n")
        
        # Get available models
        models = self.database_manager.get_available_models()
        if models:
            self.main_panel.set_models(models)
            self.results_panel.add_output(f"Found {len(models)} models\n")
            
            # If we have models, get quantities for first one
            if models:
                quantities = self.database_manager.get_available_quantities(models[0])
                if quantities:
                    self.main_panel.set_quantities(quantities)
                    self.results_panel.add_output(f"Loaded quantities for {models[0]}\n")
                else:
                    self.results_panel.add_output(f"No quantities found for {models[0]}\n")
        else:
            self.results_panel.add_output("No models found. Please check the database file.\n")
            
            # Fallback to dummy data if needed
            fallback = messagebox.askyesno("No Models Found", "Would you like to initialize with dummy data for testing?")
            if fallback:
                # Create dummy models and quantities
                dummy_models = ["GDFMX", "DDB", "DDH", "RMFNL"]
                self.main_panel.set_models(dummy_models)
                
                dummy_quantities = {
                    'eos': ['rho', 'p', 'e', 'cs2'],
                    'mr': ['M', 'R', 'T'],
                    'nmp': ['rho0', 'k0', 'e0', 'jsym0', 'lsym0']
                }
                self.main_panel.set_quantities(dummy_quantities)
                self.results_panel.add_output("Initialized with dummy data for testing\n")
    
    def check_database(self):
        """Check if database exists and offer to download if not"""
        if not self.database_manager.check_database_exists():
            download = messagebox.askyesno(
                "Database Not Found", 
                "The CompactObject CEDF EOS Database was not found. Would you like to download it now?\n\n"
                "Note: The database is approximately 1.71 GB in size.",
                parent=self
            )
            if download:
                self.download_database()
            else:
                messagebox.showinfo(
                    "Database Required", 
                    "You can download the database later from the Database menu or manually place it in the application directory.",
                    parent=self
                )
        else:
            # Database exists, initialize models
            self.update_models()
    
    def download_database(self):
        """Download the database in a separate thread and show progress"""
        # Disable the sidebar while downloading
        self.sidebar.set_state("disabled")
        
        # Show download status in results panel
        self.results_panel.clear_output()
        self.results_panel.add_output("Starting database download...\n")
        
        # Start download in separate thread
        thread = threading.Thread(target=self._download_database_thread)
        thread.daemon = True
        thread.start()
    
    def _download_database_thread(self):
        """Background thread for database download"""
        try:
            def progress_callback(percent, downloaded, total):
                self.results_panel.clear_output()
                self.results_panel.add_output(
                    f"\rDownloading: {percent:.1f}% ({downloaded:.1f} MB / {total:.1f} MB)",
                    replace_last=True
                )
            
            success = self.database_manager.download_database(progress_callback)
            
            if success:
                self.results_panel.add_output("\nDatabase download completed successfully!")
                # Update models dropdown after successful download
                self.after(100, self.update_models)
            else:
                self.results_panel.add_output("\nError downloading database. Please try again later.")
        except Exception as e:
            self.results_panel.add_output(f"\nError during download: {str(e)}")
        finally:
            # Re-enable sidebar
            self.after(100, lambda: self.sidebar.set_state("normal"))
    
    def update_models(self):
        """Update the models dropdown with available models"""
        models = self.database_manager.get_available_models()
        if models:
            self.main_panel.set_models(models)
            # Also update quantities for the first model
            if models:
                self.update_quantities(models[0])
    
    def update_quantities(self, model):
        """Update quantities based on selected model"""
        # If "Mix All Models" is selected, collect quantities from all models
        if model == "Mix All Models":
            all_models = self.database_manager.get_available_models()
            combined_quantities = {}
            
            # Get quantities from first model to start with
            if all_models:
                combined_quantities = self.database_manager.get_available_quantities(all_models[0])
            
            # If we have more than one model, merge quantities
            for model_name in all_models[1:]:
                model_quantities = self.database_manager.get_available_quantities(model_name)
                # Merge quantities from each category
                for category, items in model_quantities.items():
                    if category in combined_quantities:
                        # Add unique items from this model
                        combined_quantities[category] = list(set(combined_quantities[category] + items))
                    else:
                        combined_quantities[category] = items
            
            # Update the UI with combined quantities
            if combined_quantities:
                self.main_panel.set_quantities(combined_quantities)
        else:
            # Standard single model behavior
            quantities = self.database_manager.get_available_quantities(model)
            if quantities:
                self.main_panel.set_quantities(quantities)
    
    def run_command(self):
        """Run the CompactObject RMF sampling command"""
        # Get parameters from UI
        params = self.main_panel.get_parameters()
        if not params:
            return
        
        # Disable sidebar while command is running
        self.sidebar.set_state("disabled")
        
        # Clear previous results
        self.results_panel.clear_output()
        self.results_panel.add_output("Running command...\n\n")
        
        # Build command
        cmd = ["python", "CompactObject_RMF_sample.py"]
        
        # Add parameters
        if params.get("model") and params["model"] != "Mix All Models":
            cmd.extend(["--model", params["model"]])
        
        if params.get("n_samples"):
            cmd.extend(["--n", params["n_samples"]])
        
        if params.get("quantities"):
            for q in params["quantities"]:
                cmd.extend(["--quantity", q])
        
        if params.get("output_prefix"):
            cmd.extend(["--output", params["output_prefix"]])
        
        if params.get("output_format"):
            cmd.extend(["--format", params["output_format"]])
        
        # Handle confidence interval or highest weights
        if params.get("highest_weights", False):
            cmd.append("--highest-weights")
        elif params.get("ci_value") and not params.get("ci_optional", False):
            cmd.extend(["--ci", params["ci_value"]])
            
            # Add CI bound type if specified
            if params.get("ci_type"):
                cmd.extend(["--ci-bound", params["ci_type"]])
        
        if params.get("filters"):
            for f in params["filters"]:
                cmd.extend(["--filter", f])
        
        if params.get("all_matching", False):
            cmd.append("--filter-all-matching")
        
        if params.get("save_separate", False):
            cmd.append("--save-separate")
        
        # Show command
        self.results_panel.add_output("Command: " + " ".join(cmd) + "\n\n")
        
        # Run command in thread
        threading.Thread(target=lambda: self._run_command_thread(cmd)).start()
    
    def _run_command_thread(self, cmd):
        """Run command in background thread"""
        try:
            # Run the command
            process = subprocess.Popen(
                cmd, 
                stdout=subprocess.PIPE, 
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,
                universal_newlines=True
            )
            
            # Read output line by line
            for line in iter(process.stdout.readline, ""):
                self.results_panel.add_output(line)
            
            # Read errors
            for line in iter(process.stderr.readline, ""):
                self.results_panel.add_output("ERROR: " + line, error=True)
            
            # Wait for process to complete
            return_code = process.wait()
            
            # Check return code
            if return_code == 0:
                self.results_panel.add_output("\nCommand completed successfully!")
                
                # Add button to open results folder
                self.results_panel.add_button("Open Results Folder", self.open_results_folder)
            else:
                self.results_panel.add_output(f"\nCommand failed with return code {return_code}", error=True)
        
        except Exception as e:
            self.results_panel.add_output(f"Error executing command: {str(e)}", error=True)
        
        finally:
            # Re-enable sidebar
            self.after(100, lambda: self.sidebar.set_state("normal"))
    
    def open_results_folder(self):
        """Open the folder containing results"""
        try:
            # Get current directory
            output_dir = os.path.abspath(os.getcwd())
            
            # Open folder
            if sys.platform == 'win32':
                os.startfile(output_dir)
            elif sys.platform == 'darwin':  # macOS
                subprocess.run(['open', output_dir])
            else:  # Linux
                subprocess.run(['xdg-open', output_dir])
        except Exception as e:
            messagebox.showerror("Error", f"Could not open folder: {str(e)}", parent=self)

if __name__ == "__main__":
    app = RMFSamplingApp()
    app.mainloop()