#!/usr/bin/python3

import os
import sys
import tkinter as tk
from tkinter import filedialog, ttk
import SimpleITK as sitk
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from platipy.imaging.projects.cardiac.run import run_hybrid_segmentation
from platipy.imaging import ImageVisualiser
from platipy.imaging.label.utils import get_com

class HeartSegmentationGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Heart Segmentation Tool")
        
        # Add protocol handler for window closing
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        
        # Set minimum window size
        self.root.minsize(800, 600)
        
        # Configure grid weights
        self.root.grid_columnconfigure(0, weight=1)
        self.root.grid_rowconfigure(0, weight=1)
        
        # Create main frame with padding
        self.main_frame = ttk.Frame(root, padding="20")
        self.main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        self.main_frame.grid_columnconfigure(1, weight=1)
        
        # Create control panel frame
        control_frame = ttk.LabelFrame(self.main_frame, text="Controls", padding="10")
        control_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N), padx=5, pady=5)
        
        # Input file selection
        ttk.Button(control_frame, text="Select Input File", width=20, command=self.select_file).grid(
            row=0, column=0, padx=(0,10), pady=5)
        self.file_label = ttk.Label(control_frame, text="No file selected")
        self.file_label.grid(row=0, column=1, sticky=tk.W, pady=5)
        
        # Output directory selection
        ttk.Button(control_frame, text="Select Output Directory", width=20, command=self.select_output_dir).grid(
            row=1, column=0, padx=(0,10), pady=5)
        self.output_label = ttk.Label(control_frame, text="No directory selected")
        self.output_label.grid(row=1, column=1, sticky=tk.W, pady=5)
        
        # Process button
        self.process_button = ttk.Button(control_frame, text="Process", width=20, command=self.process_image)
        self.process_button.grid(row=2, column=0, columnspan=2, pady=10)
        self.process_button.state(['disabled'])
        
        # Create status frame
        status_frame = ttk.LabelFrame(self.main_frame, text="Status", padding="10")
        status_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), padx=5, pady=5)
        
        # Status label
        self.status_label = ttk.Label(status_frame, text="")
        self.status_label.grid(row=0, column=0, sticky=(tk.W, tk.E), padx=5, pady=5)
        
        # Create results frame
        results_frame = ttk.LabelFrame(self.main_frame, text="Results", padding="10")
        results_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), padx=5, pady=5)
        results_frame.grid_columnconfigure(0, weight=1)
        results_frame.grid_rowconfigure(0, weight=1)
        
        # Result display
        self.figure_frame = ttk.Frame(results_frame)
        self.figure_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

    def select_output_dir(self):
        directory = filedialog.askdirectory()
        if directory:
            self.output_dir = directory
            self.output_label.config(text=directory)
            self.check_process_ready()
            
    def select_file(self):
        filepath = filedialog.askopenfilename(
            filetypes=[("NIFTI files", "*.nii.gz"), ("All files", "*.*")]
        )
        if filepath:
            self.input_file = filepath
            self.file_label.config(text=os.path.basename(filepath))
            self.check_process_ready()
            
    def check_process_ready(self):
        if hasattr(self, 'input_file') and hasattr(self, 'output_dir'):
            self.process_button.state(['!disabled'])
        else:
            self.process_button.state(['disabled'])
            
    def process_image(self):
        if not hasattr(self, 'input_file') or not hasattr(self, 'output_dir'):
            return
            
        self.status_label.config(text="Processing...")
        self.process_button.state(['disabled'])
        
        self.root.after(100, self.run_segmentation)
        
    def run_segmentation(self):
        try:
            # Create output directory
            filename = os.path.splitext(os.path.splitext(os.path.basename(self.input_file))[0])[0]
            output_dir = os.path.join(self.output_dir, filename)
            os.makedirs(output_dir, exist_ok=True)
            
            # Read image
            self.status_label.config(text="Loading image...")
            self.root.update()
            test_image = sitk.ReadImage(self.input_file)
            
            # Process image
            self.status_label.config(text="Running heart segmentation...")
            self.root.update()
            auto_structures, _ = run_hybrid_segmentation(test_image)
            
            # Save segmentations
            self.status_label.config(text="Saving segmentation results...")
            self.root.update()
            for struct_name in auto_structures.keys():
                output_path = os.path.join(output_dir, f"{filename}_{struct_name}.nii.gz")
                sitk.WriteImage(auto_structures[struct_name], output_path)
            
            # Create visualization
            self.status_label.config(text="Generating visualization...")
            self.root.update()
            vis = ImageVisualiser(test_image, cut=get_com(auto_structures["Heart"]))
            vis.add_contour({struct: auto_structures[struct] for struct in auto_structures.keys()})
            fig = vis.show()
            
            # Adjust figure size
            fig.set_size_inches(10, 10)
            fig.tight_layout()
            
            # Save visualization
            plt.savefig(os.path.join(output_dir, f"{filename}_heart_visualization.png"))
            
            # Display in GUI
            for widget in self.figure_frame.winfo_children():
                widget.destroy()
            
            canvas = FigureCanvasTkAgg(fig, master=self.figure_frame)
            canvas.draw()
            canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
            
            self.status_label.config(text=f"Processing complete. Results saved to {output_dir}")
            
        except Exception as e:
            self.status_label.config(text=f"Error: {str(e)}")
        
        finally:
            self.process_button.state(['!disabled'])

    def on_closing(self):
        self.root.quit()
        self.root.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    app = HeartSegmentationGUI(root)
    root.mainloop()
    sys.exit(0)