fix_dpg_ui_layout.py

#!/usr/bin/env python3
"""
Fix UI layout issues in the DearPyGui version of OneTrainer

This script corrects tab ordering, missing fields, and layout issues in the DPG UI.
"""

import os
import sys
import json

# Add the current directory to the path
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

def fix_app_py():
    """Fix tab ordering in app.py"""
    print("Fixing tab order in app.py...")
    
    app_py_path = os.path.join(current_dir, "dpg_ui", "app.py")
    
    if not os.path.exists(app_py_path):
        print(f"Error: {app_py_path} not found")
        return False
    
    # Create a backup
    backup_path = f"{app_py_path}.bak"
    if not os.path.exists(backup_path):
        with open(app_py_path, "r") as src:
            with open(backup_path, "w") as dst:
                dst.write(src.read())
        print(f"Created backup: {backup_path}")
    
    with open(app_py_path, "r") as f:
        content = f.read()
    
    # Fix tab order in the import_tabs method
    old_tab_modules = """        tab_modules = [
            ("training_tab", "TrainingTab"),
            ("model_tab", "ModelTab"),
            ("lora_tab", "LoraTab"),
            ("concept_tab", "ConceptTab"),
            ("sampling_tab", "SamplingTab"),
            ("cloud_tab", "CloudTab")
        ]"""
    
    new_tab_modules = """        tab_modules = [
            ("model_tab", "ModelTab"),
            ("training_tab", "TrainingTab"),
            ("lora_tab", "LoraTab"),
            ("concept_tab", "ConceptTab"),
            ("sampling_tab", "SamplingTab"),
            ("additional_embeddings_tab", "AdditionalEmbeddingsTab"),
            ("cloud_tab", "CloudTab")
        ]"""
    
    content = content.replace(old_tab_modules, new_tab_modules)
    
    # Fix tab order in the create_tabs method
    old_tab_config = """        tab_config = [
            ("training", "Training", "training_tab"),
            ("model", "Model", "model_tab"),
            ("lora", "LoRA", "lora_tab"),
            ("concept", "Concept", "concept_tab"),
            ("sampling", "Sampling", "sampling_tab"),
            ("cloud", "Cloud", "cloud_tab")
        ]"""
    
    new_tab_config = """        tab_config = [
            ("model", "Model", "model_tab"),
            ("training", "Training", "training_tab"),
            ("lora", "LoRA", "lora_tab"),
            ("concept", "Concept", "concept_tab"),
            ("sampling", "Sampling", "sampling_tab"),
            ("additional_embeddings", "Additional Embeddings", "additional_embeddings_tab"),
            ("cloud", "Cloud", "cloud_tab")
        ]"""
    
    content = content.replace(old_tab_config, new_tab_config)
    
    # Change the default selected tab
    content = content.replace('self.window.select_tab("training")', 'self.window.select_tab("model")')
    
    # Write the changes back
    with open(app_py_path, "w") as f:
        f.write(content)
    
    return True

def fix_top_bar():
    """Fix top bar issues"""
    print("Fixing top bar in top_bar.py...")
    
    top_bar_path = os.path.join(current_dir, "dpg_ui", "tabs", "top_bar.py")
    
    if not os.path.exists(top_bar_path):
        print(f"Error: {top_bar_path} not found")
        # Try adapter path
        top_bar_path = os.path.join(current_dir, "dpg_ui", "adapters", "top_bar.py")
        if not os.path.exists(top_bar_path):
            print(f"Error: {top_bar_path} not found")
            return False
    
    # Create a backup
    backup_path = f"{top_bar_path}.bak"
    if not os.path.exists(backup_path):
        with open(top_bar_path, "r") as src:
            with open(backup_path, "w") as dst:
                dst.write(src.read())
        print(f"Created backup: {backup_path}")
    
    with open(top_bar_path, "r") as f:
        content = f.read()
    
    # Enhance the top bar implementation
    if "create_preset_selector" not in content:
        improved_content = """#!/usr/bin/env python3
\"\"\"
Top bar implementation for the Dear PyGui version of OneTrainer

This module provides the top bar UI component for the OneTrainer DPG UI.
\"\"\"

import os
import sys
import dearpygui.dearpygui as dpg
import json
from typing import Any, Dict, Optional, List, Callable

# Add the parent directory to sys.path to allow importing from the modules package
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from modules.util.enum.ModelType import ModelType
from modules.util.enum.TrainingMethod import TrainingMethod
from modules.util import path_util

# Import components
from dpg_ui.components.preset_selector import PresetSelector

class TopBar:
    \"\"\"
    Top bar implementation for the OneTrainer UI
    
    This class provides the top bar UI component that contains the logo,
    version information, and main actions.
    \"\"\"
    
    def __init__(self, parent: int, state_manager: Any):
        \"\"\"
        Initialize the top bar
        
        Args:
            parent: Parent container ID
            state_manager: State manager instance
        \"\"\"
        self.parent = parent
        self.state_manager = state_manager
        self.presets_dir = "training_presets"
        self.presets = []
        self.preset_selector = None
        
        # Create the UI
        self.create_ui()
        
    def create_ui(self):
        \"\"\"Create the top bar UI\"\"\"
        with dpg.group(parent=self.parent, horizontal=True):
            # Logo and title
            self.create_logo_section()
            
            # Add spacer
            dpg.add_spacer(width=20)
            
            # Model type selector
            self.create_model_type_selector()
            
            # Add spacer
            dpg.add_spacer(width=20)
            
            # Training method selector
            self.create_training_method_selector()
            
            # Add spacer
            dpg.add_spacer(width=20)
            
            # Preset selector
            self.create_preset_selector()
            
            # Add spacer
            dpg.add_spacer(width=20)
            
            # Action buttons
            self.create_action_buttons()
            
            # Add right-aligned version info
            with dpg.group(horizontal=True):
                dpg.add_spacer(width=20)
                dpg.add_text("DPG Edition (LyCORIS Enabled)", tag="version_info")
    
    def create_logo_section(self):
        \"\"\"Create the logo and title\"\"\"
        # Logo would go here if available
        logo_path = os.path.join("resources", "icons", "icon.png")
        if os.path.exists(logo_path):
            try:
                width, height, channels, data = dpg.load_image(logo_path)
                texture_id = f"logo_texture_{dpg.generate_uuid()}"
                dpg.add_static_texture(width, height, data, tag=texture_id)
                dpg.add_image(texture_id, width=40, height=40)
            except Exception as e:
                print(f"Error loading logo: {e}")
                # Fallback to a colored rectangle
                with dpg.drawlist(width=40, height=40):
                    dpg.draw_rectangle((0, 0), (40, 40), color=(41, 83, 154, 255), 
                                    fill=(41, 83, 154, 255))
        else:
            # Use a colored rectangle instead
            with dpg.drawlist(width=40, height=40):
                dpg.draw_rectangle((0, 0), (40, 40), color=(41, 83, 154, 255), 
                                fill=(41, 83, 154, 255))
            
        # Add title text
        dpg.add_text("OneTrainer", color=(255, 255, 0), tag="app_title")
    
    def create_model_type_selector(self):
        \"\"\"Create the model type selector\"\"\"
        # Add model type label
        dpg.add_text("Model Type:", color=(220, 220, 220))
        
        # Get model types
        model_types = []
        for model_type in ModelType:
            model_types.append(model_type.name)
        
        # Add selector
        self.model_selector = dpg.add_combo(
            items=model_types,
            default_value="STABLE_DIFFUSION_15",
            callback=self.on_model_change,
            width=200
        )
    
    def create_training_method_selector(self):
        \"\"\"Create the training method selector\"\"\"
        # Add training method label
        dpg.add_text("Training Method:", color=(220, 220, 220))
        
        # Get training methods
        training_methods = []
        for method in TrainingMethod:
            training_methods.append(method.name)
        
        # Add selector
        self.training_method_selector = dpg.add_combo(
            items=training_methods,
            default_value="LORA",
            callback=self.on_training_method_change,
            width=150
        )
    
    def create_preset_selector(self):
        \"\"\"Create the preset selector\"\"\"
        # Add preset selector label
        dpg.add_text("Preset:", color=(220, 220, 220))
        
        # Load presets
        with dpg.group(horizontal=True):
            self.preset_combo = dpg.add_combo(
                items=self.get_preset_names(),
                default_value="-- Select Preset --",
                callback=self.on_preset_change,
                width=200
            )
            
            # Add load button
            load_button = dpg.add_button(
                label="Load",
                callback=self.on_load_preset,
                width=60
            )
            
            # Add save button
            save_button = dpg.add_button(
                label="Save",
                callback=self.on_save_preset,
                width=60
            )
            
    def create_action_buttons(self):
        \"\"\"Create action buttons\"\"\"
        # Load config button
        load_button = dpg.add_button(
            label="Load Config",
            callback=self.on_load_config,
            width=100
        )
        
        # Save config button
        save_button = dpg.add_button(
            label="Save Config",
            callback=self.on_save_config,
            width=100
        )
    
    def get_preset_names(self) -> List[str]:
        \"\"\"Get the list of available preset names\"\"\"
        self.presets = ["-- Select Preset --"]
        
        # Check if presets directory exists
        if not os.path.exists(self.presets_dir):
            try:
                os.makedirs(self.presets_dir)
            except Exception as e:
                print(f"Error creating presets directory: {e}")
                return self.presets
        
        # Load presets from directory
        try:
            for file in os.listdir(self.presets_dir):
                if file.endswith(".json"):
                    # Remove .json extension and add to list
                    preset_name = os.path.splitext(file)[0]
                    self.presets.append(preset_name)
        except Exception as e:
            print(f"Error loading presets: {e}")
        
        return self.presets
    
    def on_model_change(self, sender, app_data):
        \"\"\"
        Handle model type change
        
        Args:
            sender: Sender ID
            app_data: Selected model type name
        \"\"\"
        # Update status
        dpg.set_value("app_title", f"OneTrainer - {app_data}")
        
        # Update state manager if available
        if hasattr(self.state_manager, "set_state"):
            try:
                from modules.util.enum.ModelType import ModelType
                model_type = ModelType[app_data]
                self.state_manager.set_state("model_type", model_type)
            except Exception as e:
                print(f"Error setting model type: {e}")
    
    def on_training_method_change(self, sender, app_data):
        \"\"\"
        Handle training method change
        
        Args:
            sender: Sender ID
            app_data: Selected training method name
        \"\"\"
        # Update state manager if available
        if hasattr(self.state_manager, "set_state"):
            try:
                from modules.util.enum.TrainingMethod import TrainingMethod
                training_method = TrainingMethod[app_data]
                self.state_manager.set_state("training_method", training_method)
            except Exception as e:
                print(f"Error setting training method: {e}")
    
    def on_preset_change(self, sender, app_data):
        \"\"\"
        Handle preset selection change
        
        Args:
            sender: Sender ID
            app_data: Selected preset name
        \"\"\"
        pass
    
    def on_load_preset(self):
        \"\"\"Handle load preset button\"\"\"
        preset_name = dpg.get_value(self.preset_combo)
        
        if preset_name == "-- Select Preset --":
            return
        
        preset_path = os.path.join(self.presets_dir, f"{preset_name}.json")
        
        if os.path.exists(preset_path):
            try:
                # Load preset data
                with open(preset_path, "r") as f:
                    preset_data = json.load(f)
                
                # Update state manager
                if hasattr(self.state_manager, "set_config"):
                    self.state_manager.set_config(preset_data)
                    print(f"Loaded preset: {preset_name}")
            except Exception as e:
                print(f"Error loading preset {preset_name}: {e}")
    
    def on_save_preset(self):
        \"\"\"Handle save preset button\"\"\"
        # Create a save dialog
        with dpg.window(label="Save Preset", modal=True, width=400, height=150, 
                        pos=(dpg.get_viewport_width() // 2 - 200, dpg.get_viewport_height() // 2 - 75)):
            dpg.add_text("Enter preset name:")
            
            # Add input field for preset name
            input_id = dpg.add_input_text(width=380)
            
            with dpg.group(horizontal=True):
                dpg.add_button(label="Save", callback=lambda: self.save_preset(dpg.get_value(input_id)))
                dpg.add_button(label="Cancel", callback=lambda: dpg.delete_item(dpg.get_current_parent_id()))
    
    def save_preset(self, preset_name):
        \"\"\"
        Save the current config as a preset
        
        Args:
            preset_name: Name of the preset to save
        \"\"\"
        if not preset_name:
            print("Error: Preset name cannot be empty")
            return
        
        # Close the dialog
        dpg.delete_item(dpg.get_current_parent_id())
        
        # Get current config from state manager
        config_data = {}
        if hasattr(self.state_manager, "get_config"):
            config_data = self.state_manager.get_config()
        
        # Save config to preset file
        preset_path = os.path.join(self.presets_dir, f"{preset_name}.json")
        
        try:
            # Create presets directory if it doesn't exist
            if not os.path.exists(self.presets_dir):
                os.makedirs(self.presets_dir)
            
            # Write data to file
            with open(preset_path, "w") as f:
                json.dump(config_data, f, indent=2)
            
            print(f"Saved preset: {preset_name}")
            
            # Update preset list
            dpg.set_item_callback(self.preset_combo, None)
            dpg.configure_item(self.preset_combo, items=self.get_preset_names())
            dpg.set_value(self.preset_combo, preset_name)
            dpg.set_item_callback(self.preset_combo, self.on_preset_change)
        except Exception as e:
            print(f"Error saving preset {preset_name}: {e}")
    
    def on_load_config(self):
        \"\"\"Handle load config button\"\"\"
        # Update status
        dpg.set_value("app_title", "OneTrainer - Loading Config...")
        
        # Create a file dialog
        with dpg.file_dialog(label="Load Config", directory_selector=False, callback=self.load_config_callback, 
                            width=700, height=400):
            dpg.add_file_extension(".json", color=(0, 255, 0, 255))
    
    def load_config_callback(self, sender, app_data):
        \"\"\"
        Handle file dialog result for loading config
        
        Args:
            sender: Sender ID
            app_data: File dialog data
        \"\"\"
        if "file_path_name" in app_data:
            config_path = app_data["file_path_name"]
            
            try:
                # Load config data
                with open(config_path, "r") as f:
                    config_data = json.load(f)
                
                # Update state manager
                if hasattr(self.state_manager, "set_config"):
                    self.state_manager.set_config(config_data)
                    print(f"Loaded config: {config_path}")
            except Exception as e:
                print(f"Error loading config {config_path}: {e}")
        
        # Restore status
        dpg.set_value("app_title", "OneTrainer")
    
    def on_save_config(self):
        \"\"\"Handle save config button\"\"\"
        # Update status
        dpg.set_value("app_title", "OneTrainer - Saving Config...")
        
        # Create a file dialog
        with dpg.file_dialog(label="Save Config", directory_selector=False, callback=self.save_config_callback, 
                            width=700, height=400):
            dpg.add_file_extension(".json", color=(0, 255, 0, 255))
    
    def save_config_callback(self, sender, app_data):
        \"\"\"
        Handle file dialog result for saving config
        
        Args:
            sender: Sender ID
            app_data: File dialog data
        \"\"\"
        if "file_path_name" in app_data:
            config_path = app_data["file_path_name"]
            
            # Add .json extension if not present
            if not config_path.lower().endswith(".json"):
                config_path += ".json"
            
            # Get current config from state manager
            config_data = {}
            if hasattr(self.state_manager, "get_config"):
                config_data = self.state_manager.get_config()
            
            try:
                # Write data to file
                with open(config_path, "w") as f:
                    json.dump(config_data, f, indent=2)
                
                print(f"Saved config: {config_path}")
            except Exception as e:
                print(f"Error saving config {config_path}: {e}")
        
        # Restore status
        dpg.set_value("app_title", "OneTrainer")
"""
        
        # Write the improved content
        with open(top_bar_path, "w") as f:
            f.write(improved_content)
        
        # Create preset selector component if it doesn't exist
        preset_selector_path = os.path.join(current_dir, "dpg_ui", "components", "preset_selector.py")
        if not os.path.exists(preset_selector_path):
            os.makedirs(os.path.dirname(preset_selector_path), exist_ok=True)
            
            with open(preset_selector_path, "w") as f:
                f.write("""#!/usr/bin/env python3
\"\"\"
Preset selector component for DPG UI

This module provides a preset selector component for the OneTrainer DPG UI.
\"\"\"

import os
import json
import dearpygui.dearpygui as dpg
from typing import List, Callable, Dict, Any, Optional

class PresetSelector:
    \"\"\"Preset selector component for DPG UI\"\"\"
    
    def __init__(self, parent: int, presets_dir: str = "training_presets", on_select: Callable = None):
        \"\"\"
        Initialize the preset selector component
        
        Args:
            parent: Parent container ID
            presets_dir: Directory containing preset files
            on_select: Callback function to call when a preset is selected
        \"\"\"
        self.parent = parent
        self.presets_dir = presets_dir
        self.on_select_callback = on_select
        
        # State variables
        self.presets = []
        self.filtered_presets = []
        self.selected_index = -1
        self.search_text = ""
        self.current_filter = "All"
        self.available_types = ["All"]
        
        # Load presets
        self.load_presets()
        
        # Create UI
        self.create_ui()
    
    def load_presets(self):
        \"\"\"Load available presets\"\"\"
        self.presets = []
        
        # Add default "None" preset
        self.presets.append({
            "name": "-- None --",
            "path": "",
            "type": "None"
        })
        
        # Check if presets directory exists
        if not os.path.exists(self.presets_dir):
            try:
                os.makedirs(self.presets_dir)
            except Exception as e:
                print(f"Error creating presets directory: {e}")
                return
        
        # Add presets from directory
        for file in os.listdir(self.presets_dir):
            if file.endswith(".json"):
                preset_path = os.path.join(self.presets_dir, file)
                preset_name = os.path.splitext(file)[0]
                
                # Determine preset type from filename
                preset_type = "Other"
                if "#sd" in preset_name.lower():
                    preset_type = "SD"
                elif "#sdxl" in preset_name.lower():
                    preset_type = "SDXL"
                elif "#lora" in preset_name.lower():
                    preset_type = "LoRA"
                elif "#embedding" in preset_name.lower():
                    preset_type = "Embedding"
                
                # Add preset to list
                self.presets.append({
                    "name": preset_name,
                    "path": preset_path,
                    "type": preset_type
                })
                
                # Add type to available types if not already present
                if preset_type not in self.available_types:
                    self.available_types.append(preset_type)
        
        # Update filtered presets
        self.filter_presets()
    
    def filter_presets(self):
        \"\"\"Filter presets based on current filter and search text\"\"\"
        self.filtered_presets = []
        
        for preset in self.presets:
            # Apply type filter
            if self.current_filter != "All" and preset["type"] != self.current_filter:
                continue
            
            # Apply search filter
            if self.search_text and self.search_text.lower() not in preset["name"].lower():
                continue
            
            # Add to filtered list
            self.filtered_presets.append(preset)
    
    def create_ui(self):
        \"\"\"Create the preset selector UI\"\"\"
        with dpg.group(parent=self.parent, horizontal=True):
            # Filter by type
            dpg.add_text("Type:")
            
            combo_id = dpg.add_combo(
                items=self.available_types,
                default_value=self.current_filter,
                width=100,
                callback=self.on_filter_change
            )
            
            # Search box
            dpg.add_text("Search:")
            
            search_id = dpg.add_input_text(
                width=150,
                callback=self.on_search_change
            )
            
            # Preset selector
            dpg.add_text("Preset:")
            
            # Create preset combo with items from filtered_presets
            preset_names = [preset["name"] for preset in self.filtered_presets]
            
            self.combo_id = dpg.add_combo(
                items=preset_names,
                default_value=preset_names[0] if preset_names else "",
                width=200,
                callback=self.on_preset_selected
            )
            
            # Load button
            load_id = dpg.add_button(
                label="Load",
                callback=self.on_load_preset
            )
            
            # Save button
            save_id = dpg.add_button(
                label="Save",
                callback=self.on_save_preset
            )
    
    def on_filter_change(self, sender, app_data):
        \"\"\"
        Handle filter type change
        
        Args:
            sender: Sender ID
            app_data: Selected filter type
        \"\"\"
        self.current_filter = app_data
        self.filter_presets()
        
        # Update combo items
        preset_names = [preset["name"] for preset in self.filtered_presets]
        dpg.configure_item(self.combo_id, items=preset_names)
        
        if preset_names:
            dpg.set_value(self.combo_id, preset_names[0])
    
    def on_search_change(self, sender, app_data):
        \"\"\"
        Handle search text change
        
        Args:
            sender: Sender ID
            app_data: Search text
        \"\"\"
        self.search_text = app_data
        self.filter_presets()
        
        # Update combo items
        preset_names = [preset["name"] for preset in self.filtered_presets]
        dpg.configure_item(self.combo_id, items=preset_names)
        
        if preset_names:
            dpg.set_value(self.combo_id, preset_names[0])
    
    def on_preset_selected(self, sender, app_data):
        \"\"\"
        Handle preset selection
        
        Args:
            sender: Sender ID
            app_data: Selected preset name
        \"\"\"
        # Find the selected preset
        selected_preset = None
        for preset in self.filtered_presets:
            if preset["name"] == app_data:
                selected_preset = preset
                break
        
        # Call callback if set
        if selected_preset and self.on_select_callback:
            self.on_select_callback(selected_preset)
    
    def on_load_preset(self):
        \"\"\"Handle load preset button\"\"\"
        # Get selected preset name
        preset_name = dpg.get_value(self.combo_id)
        
        # Find the selected preset
        selected_preset = None
        for preset in self.filtered_presets:
            if preset["name"] == preset_name:
                selected_preset = preset
                break
        
        # Skip if no preset or empty path
        if not selected_preset or not selected_preset["path"]:
            print("No preset selected or preset has no path")
            return
        
        # Load the preset
        try:
            with open(selected_preset["path"], "r") as f:
                preset_data = json.load(f)
            
            # Call callback if set
            if self.on_select_callback:
                self.on_select_callback(preset_data)
            
            print(f"Loaded preset: {selected_preset['name']}")
        except Exception as e:
            print(f"Error loading preset: {e}")
    
    def on_save_preset(self):
        \"\"\"Handle save preset button\"\"\"
        # Create a save dialog
        with dpg.window(label="Save Preset", modal=True, width=400, height=150, 
                      pos=(dpg.get_viewport_width() // 2 - 200, dpg.get_viewport_height() // 2 - 75)):
            dpg.add_text("Enter preset name:")
            
            # Add input field for preset name
            input_id = dpg.add_input_text(width=380)
            
            with dpg.group(horizontal=True):
                dpg.add_button(label="Save", callback=lambda: self.save_preset(dpg.get_value(input_id)))
                dpg.add_button(label="Cancel", callback=lambda: dpg.delete_item(dpg.get_current_parent_id()))
    
    def save_preset(self, preset_name):
        \"\"\"
        Save the current config as a preset
        
        Args:
            preset_name: Name of the preset to save
        \"\"\"
        if not preset_name:
            print("Error: Preset name cannot be empty")
            return
        
        # Close the dialog
        dpg.delete_item(dpg.get_current_parent_id())
        
        # Get current config from host
        # This will be implemented by the host application
        # For now, just create an empty config
        config_data = {}
        
        # Save config to preset file
        preset_path = os.path.join(self.presets_dir, f"{preset_name}.json")
        
        try:
            # Create presets directory if it doesn't exist
            if not os.path.exists(self.presets_dir):
                os.makedirs(self.presets_dir)
            
            # Write data to file
            with open(preset_path, "w") as f:
                json.dump(config_data, f, indent=2)
            
            print(f"Saved preset: {preset_name}")
            
            # Reload presets
            self.load_presets()
            
            # Update combo items
            preset_names = [preset["name"] for preset in self.filtered_presets]
            dpg.configure_item(self.combo_id, items=preset_names)
            
            # Select the new preset
            for i, preset in enumerate(self.filtered_presets):
                if preset["name"] == preset_name:
                    dpg.set_value(self.combo_id, preset_name)
                    break
        except Exception as e:
            print(f"Error saving preset: {e}")
""")
    else:
        print("Top bar already has preset selector implementation")
    
    return True

def create_additional_embeddings_tab():
    """Create additional embeddings tab if it doesn't exist"""
    print("Creating additional embeddings tab...")
    
    tab_dir = os.path.join(current_dir, "dpg_ui", "tabs")
    target_path = os.path.join(tab_dir, "additional_embeddings_tab.py")
    
    if os.path.exists(target_path):
        print(f"Additional embeddings tab already exists at {target_path}")
        return True
    
    # Create the file with a basic implementation
    with open(target_path, "w") as f:
        f.write("""#!/usr/bin/env python3
\"\"\"
Additional Embeddings tab implementation for the Dear PyGui version of OneTrainer

This module provides the additional embeddings tab UI component.
\"\"\"

import os
import sys
import dearpygui.dearpygui as dpg
from typing import Any, Dict, List, Optional

# Add parent directory to system path
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# Import components
from dpg_ui.components.basic import Components

class AdditionalEmbeddingsTab:
    \"\"\"
    Additional Embeddings tab implementation
    
    This class provides the additional embeddings tab UI component.
    \"\"\"
    
    def __init__(self, parent: int, state_manager: Any):
        \"\"\"
        Initialize the additional embeddings tab
        
        Args:
            parent: Parent container ID
            state_manager: State manager instance
        \"\"\"
        self.parent = parent
        self.state_manager = state_manager
        
        # Create the UI
        self.create_ui()
    
    def create_ui(self):
        \"\"\"Create the UI elements\"\"\"
        with dpg.group(parent=self.parent):
            # Header
            Components.label(dpg.last_item(), "Additional Text Embeddings", color=(255, 255, 0))
            
            dpg.add_text("Control which text embeddings are used for text conditioning.")
            
            # Add separator
            dpg.add_separator()
            
            # Settings section
            with dpg.collapsing_header(label="Text Embedding Models", default_open=True):
                # CLIP Embeddings section
                Components.label(dpg.last_item(), "CLIP Text Embeddings", color=(220, 220, 150))
                
                # Use Open CLIP
                use_open_clip = Components.switch(
                    dpg.last_item(),
                    label="Use OpenCLIP",
                    default_value=False,
                    callback=self.on_use_open_clip_change,
                    tooltip="Use OpenCLIP model for text embeddings"
                )
                
                # CLIP skip
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "CLIP Skip:", tooltip="Number of layers to skip in CLIP model")
                    clip_skip = dpg.add_slider_int(
                        default_value=1,
                        min_value=1,
                        max_value=12,
                        callback=self.on_clip_skip_change,
                        width=200
                    )
                
                dpg.add_separator()
                
                # T5 Embeddings section
                Components.label(dpg.last_item(), "T5 Text Embeddings", color=(220, 220, 150))
                
                # Use T5
                use_t5 = Components.switch(
                    dpg.last_item(),
                    label="Use T5",
                    default_value=False,
                    callback=self.on_use_t5_change,
                    tooltip="Use T5 model for additional text embeddings"
                )
                
                # T5 model
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "T5 Model:", tooltip="T5 model to use")
                    t5_model = Components.options(
                        dpg.last_item(),
                        values=["t5-small", "t5-base", "t5-large", "flan-t5-small", "flan-t5-base", "flan-t5-large"],
                        default_value="t5-base",
                        callback=self.on_t5_model_change,
                        width=200,
                        tooltip="Select T5 model variant"
                    )
                
                dpg.add_separator()
                
                # BERT Embeddings section
                Components.label(dpg.last_item(), "BERT Text Embeddings", color=(220, 220, 150))
                
                # Use BERT
                use_bert = Components.switch(
                    dpg.last_item(),
                    label="Use BERT",
                    default_value=False,
                    callback=self.on_use_bert_change,
                    tooltip="Use BERT model for additional text embeddings"
                )
                
                # BERT model
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "BERT Model:", tooltip="BERT model to use")
                    bert_model = Components.options(
                        dpg.last_item(),
                        values=["bert-base-uncased", "bert-large-uncased", "bert-base-cased", "bert-large-cased"],
                        default_value="bert-base-uncased",
                        callback=self.on_bert_model_change,
                        width=200,
                        tooltip="Select BERT model variant"
                    )
                
                dpg.add_separator()
                
                # LLaMA Embeddings section
                Components.label(dpg.last_item(), "LLaMA Text Embeddings", color=(220, 220, 150))
                
                # Use LLaMA
                use_llama = Components.switch(
                    dpg.last_item(),
                    label="Use LLaMA",
                    default_value=False,
                    callback=self.on_use_llama_change,
                    tooltip="Use LLaMA model for additional text embeddings"
                )
                
                # LLaMA model path
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "LLaMA Model Path:", tooltip="Path to LLaMA model")
                    llama_path = Components.file_entry(
                        dpg.last_item(),
                        default_path="",
                        width=300,
                        tooltip="Select LLaMA model path",
                        callback=self.on_llama_path_change
                    )
            
            # Weight settings section
            with dpg.collapsing_header(label="Embedding Weights", default_open=True):
                Components.label(dpg.last_item(), "Embedding Weight Factors", color=(220, 220, 150))
                
                # CLIP weight
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "CLIP Weight:", tooltip="Weight of CLIP embeddings")
                    clip_weight = dpg.add_slider_float(
                        default_value=1.0,
                        min_value=0.0,
                        max_value=2.0,
                        callback=self.on_clip_weight_change,
                        width=200
                    )
                
                # T5 weight
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "T5 Weight:", tooltip="Weight of T5 embeddings")
                    t5_weight = dpg.add_slider_float(
                        default_value=0.5,
                        min_value=0.0,
                        max_value=2.0,
                        callback=self.on_t5_weight_change,
                        width=200
                    )
                
                # BERT weight
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "BERT Weight:", tooltip="Weight of BERT embeddings")
                    bert_weight = dpg.add_slider_float(
                        default_value=0.5,
                        min_value=0.0,
                        max_value=2.0,
                        callback=self.on_bert_weight_change,
                        width=200
                    )
                
                # LLaMA weight
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "LLaMA Weight:", tooltip="Weight of LLaMA embeddings")
                    llama_weight = dpg.add_slider_float(
                        default_value=0.5,
                        min_value=0.0,
                        max_value=2.0,
                        callback=self.on_llama_weight_change,
                        width=200
                    )
    
    def on_use_open_clip_change(self, sender, app_data):
        \"\"\"
        Handle use OpenCLIP change
        
        Args:
            sender: Sender ID
            app_data: Switch value
        \"\"\"
        pass
    
    def on_clip_skip_change(self, sender, app_data):
        \"\"\"
        Handle CLIP skip change
        
        Args:
            sender: Sender ID
            app_data: Slider value
        \"\"\"
        pass
    
    def on_use_t5_change(self, sender, app_data):
        \"\"\"
        Handle use T5 change
        
        Args:
            sender: Sender ID
            app_data: Switch value
        \"\"\"
        pass
    
    def on_t5_model_change(self, sender, app_data):
        \"\"\"
        Handle T5 model change
        
        Args:
            sender: Sender ID
            app_data: Selected model
        \"\"\"
        pass
    
    def on_use_bert_change(self, sender, app_data):
        \"\"\"
        Handle use BERT change
        
        Args:
            sender: Sender ID
            app_data: Switch value
        \"\"\"
        pass
    
    def on_bert_model_change(self, sender, app_data):
        \"\"\"
        Handle BERT model change
        
        Args:
            sender: Sender ID
            app_data: Selected model
        \"\"\"
        pass
    
    def on_use_llama_change(self, sender, app_data):
        \"\"\"
        Handle use LLaMA change
        
        Args:
            sender: Sender ID
            app_data: Switch value
        \"\"\"
        pass
    
    def on_llama_path_change(self, path):
        \"\"\"
        Handle LLaMA path change
        
        Args:
            path: Selected path
        \"\"\"
        pass
    
    def on_clip_weight_change(self, sender, app_data):
        \"\"\"
        Handle CLIP weight change
        
        Args:
            sender: Sender ID
            app_data: Slider value
        \"\"\"
        pass
    
    def on_t5_weight_change(self, sender, app_data):
        \"\"\"
        Handle T5 weight change
        
        Args:
            sender: Sender ID
            app_data: Slider value
        \"\"\"
        pass
    
    def on_bert_weight_change(self, sender, app_data):
        \"\"\"
        Handle BERT weight change
        
        Args:
            sender: Sender ID
            app_data: Slider value
        \"\"\"
        pass
    
    def on_llama_weight_change(self, sender, app_data):
        \"\"\"
        Handle LLaMA weight change
        
        Args:
            sender: Sender ID
            app_data: Slider value
        \"\"\"
        pass
""")

    # Create the adapter for the additional embeddings tab
    adapter_dir = os.path.join(current_dir, "dpg_ui", "adapters")
    adapter_path = os.path.join(adapter_dir, "additional_embeddings_tab.py")
    
    if not os.path.exists(adapter_path):
        with open(adapter_path, "w") as f:
            f.write("""#!/usr/bin/env python3
\"\"\"
Additional Embeddings tab adapter for Dear PyGui version of OneTrainer

This module provides an adapter for the original Additional Embeddings tab.
\"\"\"

import os
import sys
import dearpygui.dearpygui as dpg
from typing import Any, Dict, List, Optional

# Add parent directory to system path
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from dpg_ui.adapters.base_adapter import BaseAdapter
from dpg_ui.components.basic import Components

class AdditionalEmbeddingsTabAdapter(BaseAdapter):
    \"\"\"
    Adapter for the Additional Embeddings tab
    
    This class adapts the original Additional Embeddings tab to the DPG UI.
    \"\"\"
    
    def __init__(self, parent: int, state_manager: Any):
        \"\"\"
        Initialize the adapter
        
        Args:
            parent: Parent container ID
            state_manager: State manager instance
        \"\"\"
        super().__init__(parent, state_manager)
        
        # Create the UI
        self.create_ui()
    
    def create_ui(self):
        \"\"\"Create the UI elements\"\"\"
        # Use the same layout as in the tabs implementation
        with dpg.group(parent=self.parent):
            # Header
            Components.label(dpg.last_item(), "Additional Text Embeddings", color=(255, 255, 0))
            
            dpg.add_text("Control which text embeddings are used for text conditioning.")
            
            # Add separator
            dpg.add_separator()
            
            # Settings section
            with dpg.collapsing_header(label="Text Embedding Models", default_open=True):
                # CLIP Embeddings section
                Components.label(dpg.last_item(), "CLIP Text Embeddings", color=(220, 220, 150))
                
                # Use Open CLIP
                use_open_clip = Components.switch(
                    dpg.last_item(),
                    label="Use OpenCLIP",
                    default_value=False,
                    callback=lambda s, v: self.update_config("use_open_clip", v),
                    tooltip="Use OpenCLIP model for text embeddings"
                )
                
                # CLIP skip
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "CLIP Skip:", tooltip="Number of layers to skip in CLIP model")
                    clip_skip = dpg.add_slider_int(
                        default_value=1,
                        min_value=1,
                        max_value=12,
                        callback=lambda s, v: self.update_config("clip_skip", v),
                        width=200
                    )
                
                dpg.add_separator()
                
                # T5 Embeddings section
                Components.label(dpg.last_item(), "T5 Text Embeddings", color=(220, 220, 150))
                
                # Use T5
                use_t5 = Components.switch(
                    dpg.last_item(),
                    label="Use T5",
                    default_value=False,
                    callback=lambda s, v: self.update_config("use_t5", v),
                    tooltip="Use T5 model for additional text embeddings"
                )
                
                # T5 model
                with dpg.group(horizontal=True):
                    Components.label(dpg.last_item(), "T5 Model:", tooltip="T5 model to use")
                    t5_model = Components.options(
                        dpg.last_item(),
                        values=["t5-small", "t5-base", "t5-large", "flan-t5-small", "flan-t5-base", "flan-t5-large"],
                        default_value="t5-base",
                        callback=lambda s, v: self.update_config("t5_model", v),
                        width=200,
                        tooltip="Select T5 model variant"
                    )
            
            # Add information text
            dpg.add_text(
                "Note: This is an adapter implementation for the Additional Embeddings tab. "
                "It provides a simplified version of the full functionality."
            )
    
    def update_config(self, key, value):
        \"\"\"
        Update configuration value
        
        Args:
            key: Configuration key
            value: New value
        \"\"\"
        if hasattr(self.state_manager, "set_state"):
            self.state_manager.set_state(key, value)
""")
    
    # Update __init__.py in adapters directory to include the new adapter
    init_path = os.path.join(adapter_dir, "__init__.py")
    if os.path.exists(init_path):
        with open(init_path, "r") as f:
            content = f.read()
        
        if "additional_embeddings_tab" not in content:
            with open(init_path, "a") as f:
                f.write("\nfrom dpg_ui.adapters.additional_embeddings_tab import AdditionalEmbeddingsTabAdapter\n")
    
    return True

def enhance_concept_tab():
    """Enhance the concept tab with more functionality"""
    print("Enhancing concept tab...")
    
    concept_tab_path = os.path.join(current_dir, "dpg_ui", "tabs", "concept_tab.py")
    
    if not os.path.exists(concept_tab_path):
        print(f"Error: {concept_tab_path} not found")
        return False
    
    # Create concept_tab_helper.py if it doesn't exist
    helper_path = os.path.join(current_dir, "dpg_ui", "tabs", "concept_tab_helper.py")
    
    if not os.path.exists(helper_path):
        with open(helper_path, "w") as f:
            f.write("""#!/usr/bin/env python3
\"\"\"
Helper functions for concept tab implementation

This module provides helper functions for the concept tab implementation.
\"\"\"

import os
import json
from typing import Dict, List, Any, Optional

def load_concepts(file_path: str, default_func=None) -> List[Dict[str, Any]]:
    \"\"\"
    Load concepts from a JSON file
    
    Args:
        file_path: Path to the concepts JSON file
        default_func: Function to call to get default values if file doesn't exist
        
    Returns:
        List of concept dictionaries
    \"\"\"
    if os.path.exists(file_path):
        try:
            with open(file_path, "r") as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading concepts from {file_path}: {e}")
            
    # If file doesn't exist or error occurred, return empty list or default
    if default_func:
        return [default_func()]
    return []

def save_concepts(concepts: List[Dict[str, Any]], file_path: str) -> bool:
    \"\"\"
    Save concepts to a JSON file
    
    Args:
        concepts: List of concept dictionaries
        file_path: Path to save the concepts JSON file
        
    Returns:
        True if successfully saved, False otherwise
    \"\"\"
    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    
    try:
        with open(file_path, "w") as f:
            json.dump(concepts, f, indent=2)
        return True
    except Exception as e:
        print(f"Error saving concepts to {file_path}: {e}")
        return False

def get_concept_image_path(concept: Dict[str, Any]) -> Optional[str]:
    \"\"\"
    Get the path to a representative image for the concept
    
    Args:
        concept: Concept dictionary
        
    Returns:
        Path to an image file or None if not found
    \"\"\"
    # Check if concept has image directory
    if "image" in concept and "directory" in concept["image"]:
        dir_path = concept["image"]["directory"]
        
        if os.path.exists(dir_path):
            # Look for image files
            for file in os.listdir(dir_path):
                if file.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
                    return os.path.join(dir_path, file)
    
    return None

def get_concept_name(concept: Dict[str, Any], index: int) -> str:
    \"\"\"
    Get the display name for a concept
    
    Args:
        concept: Concept dictionary
        index: Index of the concept in the list
        
    Returns:
        Display name for the concept
    \"\"\"
    # Check if concept has name field
    if "name" in concept:
        name = concept["name"]
    else:
        name = f"Concept {index + 1}"
    
    # For very long names, truncate and add ellipsis
    if len(name) > 20:
        name = name[:17] + "..."
        
    return name

def create_default_concept() -> Dict[str, Any]:
    \"\"\"
    Create a default concept with empty values
    
    Returns:
        Default concept dictionary
    \"\"\"
    return {
        "name": "New Concept",
        "enabled": True,
        "image": {
            "directory": "",
            "flip_p": 0.0,
            "random_crop": False,
            "caption_extension": ".txt",
            "mask_extension": ".mask.png",
            "resize": {
                "size": 512,
                "method": "BILINEAR"
            }
        },
        "mask": {
            "mask_prompt_p": 0.0,
            "path": "",
            "use_internal_mask": False
        },
        "trigger": {
            "words": [""],
            "class_words": [""],
            "prepend": False,
            "weight": 1.0
        }
    }
""")
    
    return True

def check_required_directories():
    """Check and create required directories"""
    print("Checking required directories...")
    
    # Create components directory if it doesn't exist
    components_dir = os.path.join(current_dir, "dpg_ui", "components")
    if not os.path.exists(components_dir):
        os.makedirs(components_dir, exist_ok=True)
        print(f"Created components directory: {components_dir}")
    
    # Create basic components __init__.py if it doesn't exist
    init_path = os.path.join(components_dir, "__init__.py")
    if not os.path.exists(init_path):
        with open(init_path, "w") as f:
            f.write("""#!/usr/bin/env python3
\"\"\"
UI components for Dear PyGui implementation of OneTrainer

This module provides basic UI components for the OneTrainer application.
\"\"\"

from dpg_ui.components.basic import Components
""")
        print(f"Created components/__init__.py")
    
    # Create tabs directory if it doesn't exist
    tabs_dir = os.path.join(current_dir, "dpg_ui", "tabs")
    if not os.path.exists(tabs_dir):
        os.makedirs(tabs_dir, exist_ok=True)
        print(f"Created tabs directory: {tabs_dir}")
    
    # Create tabs __init__.py if it doesn't exist
    init_path = os.path.join(tabs_dir, "__init__.py")
    if not os.path.exists(init_path):
        with open(init_path, "w") as f:
            f.write("""#!/usr/bin/env python3
\"\"\"
UI tabs for Dear PyGui implementation of OneTrainer

This module provides tab UI components for the OneTrainer application.
\"\"\"
""")
        print(f"Created tabs/__init__.py")
    
    return True

def fix_dark_theme():
    """Fix the dark theme implementation"""
    print("Fixing dark theme...")
    
    theme_path = os.path.join(current_dir, "dpg_ui", "themes", "dark_theme.py")
    
    if not os.path.exists(theme_path):
        print(f"Error: {theme_path} not found")
        return False
    
    # Create a backup
    backup_path = f"{theme_path}.bak"
    if not os.path.exists(backup_path):
        with open(theme_path, "r") as src:
            with open(backup_path, "w") as dst:
                dst.write(src.read())
        print(f"Created backup: {backup_path}")
    
    with open(theme_path, "r") as f:
        content = f.read()
    
    # Check if the theme needs enhancing
    if "create_onetrainer_button_theme" in content and "primary_color" not in content:
        # The theme implementation lacks full customization
        improved_theme = """#!/usr/bin/env python3
\"\"\"
Dark theme implementation for Dear PyGui version of OneTrainer

This module provides a dark theme for the OneTrainer application using Dear PyGui.
\"\"\"

import dearpygui.dearpygui as dpg

# Define OneTrainer theme colors
primary_color = (41, 83, 154, 255)  # Blue
secondary_color = (63, 119, 183, 255)  # Lighter blue
accent_color = (255, 255, 0, 255)  # Yellow
background_color = (30, 30, 30, 255)  # Dark gray
text_color = (255, 255, 255, 255)  # White
disabled_color = (120, 120, 120, 255)  # Gray

def apply_theme():
    \"\"\"
    Apply a dark theme to the application
    
    Returns:
        The theme ID
    \"\"\"
    with dpg.theme() as theme:
        with dpg.theme_component(dpg.mvAll):
            # Standard theme settings
            dpg.add_theme_color(dpg.mvThemeCol_WindowBg, (30, 30, 30, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ChildBg, (25, 25, 25, 255))
            dpg.add_theme_color(dpg.mvThemeCol_PopupBg, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_Border, (60, 60, 60, 255))
            dpg.add_theme_color(dpg.mvThemeCol_BorderShadow, (0, 0, 0, 0))
            dpg.add_theme_color(dpg.mvThemeCol_FrameBg, (45, 45, 45, 255))
            dpg.add_theme_color(dpg.mvThemeCol_FrameBgHovered, (50, 50, 50, 255))
            dpg.add_theme_color(dpg.mvThemeCol_FrameBgActive, (55, 55, 55, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TitleBg, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TitleBgActive, primary_color)
            dpg.add_theme_color(dpg.mvThemeCol_TitleBgCollapsed, (35, 35, 35, 100))
            dpg.add_theme_color(dpg.mvThemeCol_MenuBarBg, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarBg, (30, 30, 30, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrab, (80, 80, 80, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrabHovered, (90, 90, 90, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrabActive, (100, 100, 100, 255))
            dpg.add_theme_color(dpg.mvThemeCol_CheckMark, accent_color)
            dpg.add_theme_color(dpg.mvThemeCol_SliderGrab, primary_color)
            dpg.add_theme_color(dpg.mvThemeCol_SliderGrabActive, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_Button, primary_color)
            dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_Header, primary_color)
            dpg.add_theme_color(dpg.mvThemeCol_HeaderHovered, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_HeaderActive, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_Tab, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TabHovered, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_TabActive, primary_color)
            dpg.add_theme_color(dpg.mvThemeCol_TabUnfocused, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TabUnfocusedActive, (45, 45, 45, 255))
            dpg.add_theme_color(dpg.mvThemeCol_Text, text_color)
            dpg.add_theme_color(dpg.mvThemeCol_TextDisabled, disabled_color)
            
            # Style settings
            dpg.add_theme_style(dpg.mvStyleVar_WindowRounding, 6)
            dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_ChildRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_PopupRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_TabRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_GrabRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_ScrollbarRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_WindowTitleAlign, 0.5, 0.5)
            dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 8, 4)
            dpg.add_theme_style(dpg.mvStyleVar_ItemSpacing, 8, 6)
            dpg.add_theme_style(dpg.mvStyleVar_ItemInnerSpacing, 4, 4)
            
    # Apply the theme globally
    dpg.bind_theme(theme)
    
    return theme

def create_onetrainer_button_theme(is_primary=True):
    \"\"\"
    Create a theme for OneTrainer buttons
    
    Args:
        is_primary: Whether this is a primary button
        
    Returns:
        The theme ID
    \"\"\"
    with dpg.theme() as theme:
        with dpg.theme_component(dpg.mvButton):
            if is_primary:
                # Primary button colors
                dpg.add_theme_color(dpg.mvThemeCol_Button, primary_color)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, secondary_color)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, secondary_color)
            else:
                # Secondary button colors
                dpg.add_theme_color(dpg.mvThemeCol_Button, (60, 60, 60, 255))
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (80, 80, 80, 255))
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (90, 90, 90, 255))
            
            # Button style
            dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 8, 4)
    
    return theme

def create_onetrainer_tab_theme():
    \"\"\"
    Create a theme for OneTrainer tabs
    
    Returns:
        The theme ID
    \"\"\"
    with dpg.theme() as theme:
        with dpg.theme_component(dpg.mvTab):
            dpg.add_theme_color(dpg.mvThemeCol_Tab, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TabHovered, secondary_color)
            dpg.add_theme_color(dpg.mvThemeCol_TabActive, primary_color)
            dpg.add_theme_color(dpg.mvThemeCol_TabUnfocused, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TabUnfocusedActive, (45, 45, 45, 255))
            
            # Tab style
            dpg.add_theme_style(dpg.mvStyleVar_TabRounding, 4)
    
    return theme

def create_onetrainer_frame_theme():
    \"\"\"
    Create a theme for OneTrainer frames
    
    Returns:
        The theme ID
    \"\"\"
    with dpg.theme() as theme:
        with dpg.theme_component(dpg.mvAll):
            dpg.add_theme_color(dpg.mvThemeCol_ChildBg, (25, 25, 25, 255))
            dpg.add_theme_color(dpg.mvThemeCol_Border, (60, 60, 60, 255))
            
            # Frame style
            dpg.add_theme_style(dpg.mvStyleVar_ChildRounding, 4)
            dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 8, 8)
    
    return theme
"""
        
        # Write the improved theme
        with open(theme_path, "w") as f:
            f.write(improved_theme)
    
    return True

def main():
    print("Fixing DearPyGui UI layout issues...")
    
    # Check and create required directories
    check_required_directories()
    
    # Fix the tab order in app.py
    fix_app_py()
    
    # Fix top bar implementation
    fix_top_bar()
    
    # Create additional embeddings tab
    create_additional_embeddings_tab()
    
    # Enhance concept tab
    enhance_concept_tab()
    
    # Fix dark theme
    fix_dark_theme()
    
    print("\nAll UI layout fixes have been applied. You can now run the UI with:")
    print("./fix_and_launch_dpg_ui.sh\n")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())