"""
Key to STL Generator - Educational Tool
Analyzes key images and creates 3D STL models of the key teeth profile
Enhanced version with manual selection tools
"""

from PIL import Image, ImageDraw, ImageFilter, ImageTk, ImageEnhance
import numpy as np
from scipy import ndimage, signal
from scipy.ndimage import median_filter, gaussian_filter
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import struct

class KeyToSTL:
    def __init__(self):
        self.original_image = None
        self.image = None
        self.processed_image = None
        self.key_profile = []
        self.key_bounds = None
        self.manual_bounds = None
        self.groove_regions = []
        
    def load_image(self, filepath):
        """Load and preprocess the key image"""
        self.original_image = Image.open(filepath)
        self.original_image = self.original_image.convert('RGB')
        self.image = self.original_image.convert('L')
        print(f"Image loaded: {self.image.size}")
        return self.image
    
    def set_manual_bounds(self, x1, y1, x2, y2):
        """Set manual key bounds from user selection"""
        self.manual_bounds = {
            'left': min(x1, x2),
            'right': max(x1, x2),
            'top': min(y1, y2),
            'bottom': max(y1, y2)
        }
        print(f"Manual bounds set: {self.manual_bounds}")
    
    def add_groove_region(self, x1, y1, x2, y2):
        """Add a groove region that should be emphasized"""
        region = {
            'left': min(x1, x2),
            'right': max(x1, x2),
            'top': min(y1, y2),
            'bottom': max(y1, y2)
        }
        self.groove_regions.append(region)
        print(f"Groove region added: {region}")
    
    def clear_groove_regions(self):
        """Clear all groove regions"""
        self.groove_regions = []
        print("Groove regions cleared")
    
    def denoise_image(self, median_size=3, gaussian_sigma=1.5):
        """Apply strong denoising to handle static/noise in images"""
        if self.image is None:
            raise ValueError("No image loaded")
        
        img_array = np.array(self.image)
        denoised = median_filter(img_array, size=median_size)
        denoised = gaussian_filter(denoised, sigma=gaussian_sigma)
        
        return denoised
    
    def auto_detect_key(self, noise_handling=True):
        """Automatically detect the key region with noise handling"""
        if self.image is None:
            raise ValueError("No image loaded")
        
        if noise_handling:
            img_array = self.denoise_image(median_size=5, gaussian_sigma=2.0)
        else:
            img_array = np.array(self.image)
        
        img_pil = Image.fromarray(img_array.astype(np.uint8))
        enhancer = ImageEnhance.Contrast(img_pil)
        enhanced = enhancer.enhance(2.5)
        
        img_array = np.array(enhanced)
        img_array = gaussian_filter(img_array, sigma=1.0)
        
        sobel_x = ndimage.sobel(img_array, axis=1)
        sobel_y = ndimage.sobel(img_array, axis=0)
        edge_magnitude = np.hypot(sobel_x, sobel_y)
        
        edge_magnitude = (edge_magnitude / edge_magnitude.max() * 255).astype(np.uint8)
        
        threshold = np.percentile(edge_magnitude[edge_magnitude > 0], 90)
        binary_edges = edge_magnitude > threshold
        
        from scipy.ndimage import binary_closing, binary_opening, binary_dilation, binary_erosion
        
        binary_edges = binary_opening(binary_edges, structure=np.ones((2, 2)))
        binary_edges = binary_closing(binary_edges, structure=np.ones((3, 3)))
        binary_edges = binary_erosion(binary_edges, structure=np.ones((2, 2)))
        binary_edges = binary_dilation(binary_edges, structure=np.ones((2, 2)))
        
        self.processed_image = Image.fromarray((binary_edges * 255).astype(np.uint8))
        return self.processed_image
    
    def find_key_bounds(self, margin_percent=0.1):
        """Find the bounding box of the key with noise-resistant algorithm"""
        if self.processed_image is None:
            raise ValueError("Process image first")
        
        # Use manual bounds if available
        if self.manual_bounds:
            self.key_bounds = self.manual_bounds.copy()
            print(f"Using manual bounds: {self.key_bounds}")
            return self.key_bounds
        
        edges_array = np.array(self.processed_image)
        height, width = edges_array.shape
        
        row_content = np.sum(edges_array, axis=1)
        col_content = np.sum(edges_array, axis=0)
        
        row_smooth = gaussian_filter(row_content.astype(float), sigma=height * 0.02)
        col_smooth = gaussian_filter(col_content.astype(float), sigma=width * 0.01)
        
        row_threshold = np.percentile(row_smooth[row_smooth > 0], 25)
        col_threshold = np.percentile(col_smooth[col_smooth > 0], 15)
        
        rows_with_content = np.where(row_smooth > row_threshold)[0]
        cols_with_content = np.where(col_smooth > col_threshold)[0]
        
        if len(rows_with_content) > 0 and len(cols_with_content) > 0:
            row_margin = int(height * margin_percent)
            col_margin = int(width * margin_percent * 0.5)
            
            self.key_bounds = {
                'top': max(0, rows_with_content[0] - row_margin),
                'bottom': min(height - 1, rows_with_content[-1] + row_margin),
                'left': max(0, cols_with_content[0] - col_margin),
                'right': min(width - 1, cols_with_content[-1] + col_margin)
            }
            print(f"Key bounds detected: {self.key_bounds}")
        else:
            self.key_bounds = {
                'top': height // 4,
                'bottom': 3 * height // 4,
                'left': width // 10,
                'right': 9 * width // 10
            }
            print("Using fallback bounds")
        
        return self.key_bounds
    
    def extract_key_profile(self, use_top=True, consensus_window=5):
        """Extract key profile with noise-resistant consensus approach"""
        if self.processed_image is None or self.key_bounds is None:
            raise ValueError("Process image and find bounds first")
        
        edges_array = np.array(self.processed_image)
        bounds = self.key_bounds
        
        roi = edges_array[bounds['top']:bounds['bottom'], 
                         bounds['left']:bounds['right']]
        
        height, width = roi.shape
        profile = []
        
        for x in range(width):
            start_col = max(0, x - consensus_window // 2)
            end_col = min(width, x + consensus_window // 2 + 1)
            
            window = roi[:, start_col:end_col]
            
            edge_positions = []
            for col in range(window.shape[1]):
                column = window[:, col]
                edge_pixels = np.where(column > 128)[0]
                
                if len(edge_pixels) > 0:
                    if use_top:
                        edge_positions.append(edge_pixels[0])
                    else:
                        edge_positions.append(edge_pixels[-1])
            
            if len(edge_positions) > 0:
                y_pos = int(np.median(edge_positions))
                actual_y = y_pos + bounds['top']
                actual_x = x + bounds['left']
                profile.append((actual_x, actual_y))
            else:
                if len(profile) > 0:
                    profile.append((x + bounds['left'], profile[-1][1]))
                else:
                    profile.append((x + bounds['left'], bounds['top'] + height // 2))
        
        profile = self._fill_profile_gaps(profile)
        
        # Apply groove emphasis
        if self.groove_regions:
            profile = self._emphasize_grooves(profile)
        
        self.key_profile = profile
        return profile
    
    def _emphasize_grooves(self, profile, emphasis_factor=1.5):
        """Emphasize groove regions in the profile"""
        emphasized = []
        
        for x, y in profile:
            # Check if point is in any groove region
            in_groove = False
            for groove in self.groove_regions:
                if groove['left'] <= x <= groove['right']:
                    in_groove = True
                    break
            
            if in_groove:
                # Find average Y in this region
                region_ys = [py for px, py in profile if groove['left'] <= px <= groove['right']]
                if region_ys:
                    avg_y = np.mean(region_ys)
                    # Emphasize deviation from average
                    deviation = y - avg_y
                    new_y = avg_y + deviation * emphasis_factor
                    emphasized.append((x, int(new_y)))
                else:
                    emphasized.append((x, y))
            else:
                emphasized.append((x, y))
        
        return emphasized
    
    def _fill_profile_gaps(self, profile, max_gap=10):
        """Fill gaps in profile with interpolation"""
        if len(profile) < 2:
            return profile
        
        filled = [profile[0]]
        
        for i in range(1, len(profile)):
            prev_x, prev_y = filled[-1]
            curr_x, curr_y = profile[i]
            
            if abs(curr_y - prev_y) > max_gap:
                steps = int(abs(curr_x - prev_x))
                if steps > 1:
                    for j in range(1, steps):
                        interp_x = prev_x + j
                        interp_y = int(prev_y + (curr_y - prev_y) * j / steps)
                        filled.append((interp_x, interp_y))
            
            filled.append((curr_x, curr_y))
        
        return filled
    
    def smooth_profile_advanced(self, window_size=11, poly_order=3):
        """Apply Savitzky-Golay filter with noise handling"""
        if not self.key_profile:
            raise ValueError("Extract profile first")
        
        profile_array = np.array(self.key_profile)
        x_coords = profile_array[:, 0]
        y_coords = profile_array[:, 1]
        
        y_median = median_filter(y_coords, size=5)
        
        window_size = min(window_size, len(y_median))
        if window_size % 2 == 0:
            window_size -= 1
        window_size = max(window_size, poly_order + 2)
        if window_size % 2 == 0:
            window_size += 1
        
        try:
            smoothed_y = signal.savgol_filter(y_median, window_size, poly_order)
        except:
            smoothed_y = gaussian_filter(y_median, sigma=window_size / 3.0)
        
        self.key_profile = [(int(x), int(y)) for x, y in zip(x_coords, smoothed_y)]
        return self.key_profile
    
    def remove_outliers(self, threshold=3.0, window_size=10):
        """Remove outlier points using local statistics"""
        if not self.key_profile:
            return
        
        profile_array = np.array(self.key_profile)
        y_coords = profile_array[:, 1].astype(float)
        
        cleaned_y = y_coords.copy()
        
        for i in range(len(y_coords)):
            start = max(0, i - window_size // 2)
            end = min(len(y_coords), i + window_size // 2 + 1)
            window = y_coords[start:end]
            
            local_median = np.median(window)
            mad = np.median(np.abs(window - local_median))
            
            if mad > 0:
                z_score = abs(y_coords[i] - local_median) / (1.4826 * mad)
                if z_score > threshold:
                    cleaned_y[i] = local_median
        
        self.key_profile = [(int(x), int(y)) for x, y in zip(profile_array[:, 0], cleaned_y)]
        return self.key_profile
    
    def normalize_profile(self):
        """Normalize profile to 0-1 range for depth"""
        if not self.key_profile:
            raise ValueError("Extract profile first")
        
        profile_array = np.array(self.key_profile)
        y_coords = profile_array[:, 1]
        
        y_min, y_max = y_coords.min(), y_coords.max()
        
        if y_max - y_min < 1:
            return [(x, 0.5) for x, _ in self.key_profile]
        
        normalized = []
        for x, y in self.key_profile:
            normalized_depth = (y - y_min) / (y_max - y_min)
            normalized.append((x, normalized_depth))
        
        return normalized
    
    def generate_stl(self, output_path, depth_scale=2.0, base_thickness=1.0):
        """Generate STL file from the key profile"""
        if not self.key_profile:
            raise ValueError("Extract profile first")
        
        normalized_profile = self.normalize_profile()
        
        profile_array = np.array(self.key_profile)
        x_coords = profile_array[:, 0]
        x_min, x_max = x_coords.min(), x_coords.max()
        
        key_length = float(x_max - x_min)
        key_width = 10.0
        
        scaled_profile = []
        for (x, _), (_, normalized_depth) in zip(self.key_profile, normalized_profile):
            scaled_x = float(x - x_min)
            actual_depth = normalized_depth * depth_scale
            scaled_profile.append((scaled_x, actual_depth))
        
        vertices, faces = self._create_mesh(scaled_profile, key_width, base_thickness)
        self._write_stl(output_path, vertices, faces)
        
        print(f"STL file saved: {output_path}")
        print(f"Dimensions: {key_length:.1f} x {key_width:.1f} x {base_thickness + depth_scale:.1f} units")
    
    def _create_mesh(self, profile, width, base_thickness):
        """Create 3D mesh from 2D profile"""
        vertices = []
        faces = []
        
        n = len(profile)
        
        for x, depth in profile:
            z_top = base_thickness + depth
            vertices.append([x, 0.0, z_top])
            vertices.append([x, width, z_top])
            vertices.append([x, 0.0, 0.0])
            vertices.append([x, width, 0.0])
        
        for i in range(n - 1):
            base = i * 4
            next_base = (i + 1) * 4
            
            v1, v2 = base, base + 1
            v3, v4 = next_base, next_base + 1
            faces.append([v1, v3, v2])
            faces.append([v2, v3, v4])
            
            v1, v2 = base + 2, base + 3
            v3, v4 = next_base + 2, next_base + 3
            faces.append([v1, v2, v3])
            faces.append([v2, v4, v3])
            
            v1, v2 = base, base + 2
            v3, v4 = next_base, next_base + 2
            faces.append([v1, v3, v2])
            faces.append([v2, v3, v4])
            
            v1, v2 = base + 1, base + 3
            v3, v4 = next_base + 1, next_base + 3
            faces.append([v1, v2, v3])
            faces.append([v2, v4, v3])
        
        faces.append([0, 2, 1])
        faces.append([1, 2, 3])
        
        last = (n - 1) * 4
        faces.append([last, last + 1, last + 2])
        faces.append([last + 1, last + 3, last + 2])
        
        return vertices, faces
    
    def _write_stl(self, filepath, vertices, faces):
        """Write vertices and faces to binary STL file"""
        with open(filepath, 'wb') as f:
            header = b'Key STL - Python Generated' + b' ' * (80 - 26)
            f.write(header)
            f.write(struct.pack('<I', len(faces)))
            
            for face in faces:
                v1 = vertices[face[0]]
                v2 = vertices[face[1]]
                v3 = vertices[face[2]]
                
                normal = self._calculate_normal(v1, v2, v3)
                
                f.write(struct.pack('<fff', *normal))
                f.write(struct.pack('<fff', *v1))
                f.write(struct.pack('<fff', *v2))
                f.write(struct.pack('<fff', *v3))
                f.write(struct.pack('<H', 0))
    
    def _calculate_normal(self, v1, v2, v3):
        """Calculate face normal vector"""
        v1, v2, v3 = np.array(v1), np.array(v2), np.array(v3)
        edge1 = v2 - v1
        edge2 = v3 - v1
        normal = np.cross(edge1, edge2)
        
        length = np.linalg.norm(normal)
        if length > 0:
            normal = normal / length
        else:
            normal = np.array([0.0, 0.0, 1.0])
        
        return normal.tolist()


class KeySTLGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Key to STL Generator - Manual Selection Tools")
        self.root.geometry("1200x800")
        
        self.key_processor = KeyToSTL()
        self.canvas_image = None
        self.photo_image = None
        
        # Selection state
        self.selection_mode = None  # 'key_bounds' or 'groove'
        self.selection_start = None
        self.current_rect = None
        self.drawn_rectangles = []
        
        self._create_widgets()
        self._bind_canvas_events()
    
    def _create_widgets(self):
        main_container = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        main_container.pack(fill=tk.BOTH, expand=True)
        
        # Control panel
        control_frame = ttk.Frame(main_container, padding="10")
        main_container.add(control_frame, weight=0)
        
        ttk.Label(control_frame, text="Key to STL Generator", 
                 font=('Arial', 14, 'bold')).pack(pady=10)
        
        # Step 1
        ttk.Label(control_frame, text="Step 1: Load Image", 
                 font=('Arial', 10, 'bold')).pack(pady=(10, 5))
        ttk.Button(control_frame, text="📁 Load Image", 
                  command=self.load_image, width=22).pack(pady=5)
        
        ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
        
        # Step 2 - Manual Selection Tools
        ttk.Label(control_frame, text="Step 2: Manual Selection (Optional)", 
                 font=('Arial', 10, 'bold')).pack(pady=(5, 5))
        
        tool_frame = ttk.LabelFrame(control_frame, text="Selection Tools", padding="5")
        tool_frame.pack(fill='x', pady=5)
        
        ttk.Button(tool_frame, text="🔲 Select Key Area", 
                  command=self.start_key_selection, width=20).pack(pady=2)
        ttk.Label(tool_frame, text="Drag rectangle around key", 
                 font=('Arial', 8), foreground='gray').pack()
        
        ttk.Separator(tool_frame, orient='horizontal').pack(fill='x', pady=5)
        
        ttk.Button(tool_frame, text="🎯 Mark Groove Region", 
                  command=self.start_groove_selection, width=20).pack(pady=2)
        ttk.Label(tool_frame, text="Mark areas with teeth cuts", 
                 font=('Arial', 8), foreground='gray').pack()
        
        ttk.Button(tool_frame, text="🗑️ Clear Grooves", 
                  command=self.clear_grooves, width=20).pack(pady=2)
        
        ttk.Button(tool_frame, text="↩️ Cancel Selection", 
                  command=self.cancel_selection, width=20).pack(pady=2)
        
        self.selection_status = ttk.Label(tool_frame, text="No active tool", 
                                         font=('Arial', 8, 'italic'), foreground='blue')
        self.selection_status.pack(pady=5)
        
        ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
        
        # Step 3 - Auto Process
        ttk.Label(control_frame, text="Step 3: Auto Process", 
                 font=('Arial', 10, 'bold')).pack(pady=(5, 5))
        
        self.noise_handling_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(control_frame, text="Enable noise reduction", 
                       variable=self.noise_handling_var).pack(pady=2)
        
        ttk.Button(control_frame, text="🔍 Auto Detect Key", 
                  command=self.auto_process, width=22).pack(pady=5)
        
        self.edge_top_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(control_frame, text="Use top edge", 
                       variable=self.edge_top_var).pack(pady=2)
        
        ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
        
        # Step 4 - Refine
        ttk.Label(control_frame, text="Step 4: Refine Profile", 
                 font=('Arial', 10, 'bold')).pack(pady=(5, 5))
        
        ttk.Label(control_frame, text="Smoothing:").pack()
        self.smooth_var = tk.IntVar(value=11)
        ttk.Scale(control_frame, from_=5, to=21, variable=self.smooth_var, 
                 orient=tk.HORIZONTAL).pack()
        
        ttk.Button(control_frame, text="✨ Smooth", 
                  command=self.smooth_profile, width=22).pack(pady=2)
        ttk.Button(control_frame, text="🧹 Remove Outliers", 
                  command=self.remove_outliers, width=22).pack(pady=2)
        
        ttk.Separator(control_frame, orient='horizontal').pack(fill='x', pady=10)
        
        # Step 5 - Export
        ttk.Label(control_frame, text="Step 5: Export STL", 
                 font=('Arial', 10, 'bold')).pack(pady=(5, 5))
        
        ttk.Label(control_frame, text="Teeth Depth (mm):").pack()
        self.depth_var = tk.DoubleVar(value=2.0)
        ttk.Entry(control_frame, textvariable=self.depth_var, width=15).pack()
        
        ttk.Label(control_frame, text="Base (mm):").pack()
        self.base_var = tk.DoubleVar(value=1.0)
        ttk.Entry(control_frame, textvariable=self.base_var, width=15).pack()
        
        ttk.Button(control_frame, text="💾 Generate STL", 
                  command=self.generate_stl, width=22).pack(pady=10)
        
        # Canvas
        canvas_frame = ttk.Frame(main_container)
        main_container.add(canvas_frame, weight=1)
        
        self.canvas = tk.Canvas(canvas_frame, bg='#f0f0f0', highlightthickness=0, cursor='crosshair')
        
        h_scrollbar = ttk.Scrollbar(canvas_frame, orient=tk.HORIZONTAL, command=self.canvas.xview)
        v_scrollbar = ttk.Scrollbar(canvas_frame, orient=tk.VERTICAL, command=self.canvas.yview)
        
        self.canvas.configure(xscrollcommand=h_scrollbar.set, yscrollcommand=v_scrollbar.set)
        
        h_scrollbar.pack(side=tk.BOTTOM, fill=tk.X)
        v_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        # Status
        self.status_label = ttk.Label(self.root, text="Ready - Load a key image to start", 
                                     relief=tk.SUNKEN, padding=(5, 2))
        self.status_label.pack(side=tk.BOTTOM, fill=tk.X)
    
    def _bind_canvas_events(self):
        """Bind mouse events for manual selection"""
        self.canvas.bind('<Button-1>', self.on_canvas_click)
        self.canvas.bind('<B1-Motion>', self.on_canvas_drag)
        self.canvas.bind('<ButtonRelease-1>', self.on_canvas_release)
    
    def start_key_selection(self):
        """Start key bounds selection mode"""
        if self.photo_image is None:
            messagebox.showwarning("Warning", "Please load an image first")
            return
        
        self.selection_mode = 'key_bounds'
        self.selection_status.config(text="✓ Drag rectangle around key area")
        self.canvas.config(cursor='crosshair')
        self.status_label.config(text="SELECTION MODE: Draw rectangle around the key")
    
    def start_groove_selection(self):
        """Start groove region selection mode"""
        if self.photo_image is None:
            messagebox.showwarning("Warning", "Please load an image first")
            return
        
        self.selection_mode = 'groove'
        self.selection_status.config(text="✓ Drag rectangle over groove area")
        self.canvas.config(cursor='crosshair')
        self.status_label.config(text="GROOVE MODE: Draw rectangles over areas with teeth cuts")
    
    def cancel_selection(self):
        """Cancel current selection mode"""
        self.selection_mode = None
        self.selection_start = None
        if self.current_rect:
            self.canvas.delete(self.current_rect)
            self.current_rect = None
        self.selection_status.config(text="No active tool")
        self.canvas.config(cursor='')
        self.status_label.config(text="Selection cancelled")
    
    def clear_grooves(self):
        """Clear all groove regions"""
        self.key_processor.clear_groove_regions()
        # Redraw without groove rectangles
        self._redraw_with_annotations()
        self.status_label.config(text="✓ Groove regions cleared")
    
    def on_canvas_click(self, event):
        """Handle canvas click for selection start"""
        if self.selection_mode:
            self.selection_start = (self.canvas.canvasx(event.x), 
                                   self.canvas.canvasy(event.y))
            if self.current_rect:
                self.canvas.delete(self.current_rect)
    
    def on_canvas_drag(self, event):
        """Handle canvas drag for drawing selection rectangle"""
        if self.selection_mode and self.selection_start:
            x1, y1 = self.selection_start
            x2 = self.canvas.canvasx(event.x)
            y2 = self.canvas.canvasy(event.y)
            
            if self.current_rect:
                self.canvas.delete(self.current_rect)
            
            # Draw rectangle with different colors for different modes
            color = 'green' if self.selection_mode == 'key_bounds' else 'orange'
            self.current_rect = self.canvas.create_rectangle(
                x1, y1, x2, y2,
                outline=color, width=3, dash=(5, 5)
            )
    
    def on_canvas_release(self, event):
        """Handle canvas release to finalize selection"""
        if self.selection_mode and self.selection_start:
            x1, y1 = self.selection_start
            x2 = self.canvas.canvasx(event.x)
            y2 = self.canvas.canvasy(event.y)
            
            # Ensure we have a valid rectangle
            if abs(x2 - x1) > 5 and abs(y2 - y1) > 5:
                if self.selection_mode == 'key_bounds':
                    self.key_processor.set_manual_bounds(int(x1), int(y1), int(x2), int(y2))
                    self.status_label.config(text="✓ Key area selected")
                    # Redraw with permanent rectangle
                    if self.current_rect:
                        self.canvas.delete(self.current_rect)
                    self.current_rect = self.canvas.create_rectangle(
                        x1, y1, x2, y2,
                        outline='green', width=2
                    )
                    self.selection_mode = None
                    self.selection_status.config(text="Key area set")
                    
                elif self.selection_mode == 'groove':
                    self.key_processor.add_groove_region(int(x1), int(y1), int(x2), int(y2))
                    self.status_label.config(text=f"✓ Groove region added (Total: {len(self.key_processor.groove_regions)})")
                    # Draw permanent groove rectangle
                    if self.current_rect:
                        self.canvas.delete(self.current_rect)
                    groove_rect = self.canvas.create_rectangle(
                        x1, y1, x2, y2,
                        outline='orange', width=2, dash=(3, 3)
                    )
                    self.drawn_rectangles.append(groove_rect)
                    # Stay in groove mode for multiple selections
            
            self.selection_start = None
            self.current_rect = None
    
    def _redraw_with_annotations(self):
        """Redraw image with all annotations"""
        if self.key_processor.original_image:
            self._display_image(self.key_processor.original_image.convert('L'))
            
            # Redraw key bounds if set
            if self.key_processor.manual_bounds:
                bounds = self.key_processor.manual_bounds
                self.canvas.create_rectangle(
                    bounds['left'], bounds['top'],
                    bounds['right'], bounds['bottom'],
                    outline='green', width=2
                )
            
            # Redraw groove regions
            self.drawn_rectangles = []
            for groove in self.key_processor.groove_regions:
                rect = self.canvas.create_rectangle(
                    groove['left'], groove['top'],
                    groove['right'], groove['bottom'],
                    outline='orange', width=2, dash=(3, 3)
                )
                self.drawn_rectangles.append(rect)
    
    def load_image(self):
        filepath = filedialog.askopenfilename(
            title="Select Key Image",
            filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp *.gif"), ("All files", "*.*")]
        )
        
        if filepath:
            try:
                self.key_processor.load_image(filepath)
                self._display_image(self.key_processor.image)
                self.status_label.config(text=f"✓ Loaded: {filepath.split('/')[-1]}")
            except Exception as e:
                messagebox.showerror("Error", f"Failed to load image: {str(e)}")
    
    def auto_process(self):
        try:
            noise_handling = self.noise_handling_var.get()
            
            self.status_label.config(text="Processing... applying noise reduction" if noise_handling else "Processing... detecting edges")
            self.root.update()
            
            self.key_processor.auto_detect_key(noise_handling=noise_handling)
            
            self.status_label.config(text="Processing... finding key bounds")
            self.root.update()
            
            self.key_processor.find_key_bounds()
            
            self.status_label.config(text="Processing... extracting profile")
            self.root.update()
            
            use_top = self.edge_top_var.get()
            profile = self.key_processor.extract_key_profile(use_top=use_top, consensus_window=7)
            
            self._display_image_with_profile()
            
            groove_info = f" (with {len(self.key_processor.groove_regions)} groove regions)" if self.key_processor.groove_regions else ""
            self.status_label.config(text=f"✓ Profile extracted: {len(profile)} points{groove_info}")
            
        except Exception as e:
            messagebox.showerror("Error", f"Auto processing failed: {str(e)}")
            import traceback
            traceback.print_exc()
    
    def smooth_profile(self):
        try:
            window = self.smooth_var.get()
            if window % 2 == 0:
                window += 1
            self.key_processor.smooth_profile_advanced(window_size=window)
            self._display_image_with_profile()
            self.status_label.config(text="✓ Profile smoothed")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to smooth: {str(e)}")
    
    def remove_outliers(self):
        try:
            self.key_processor.remove_outliers()
            self._display_image_with_profile()
            self.status_label.config(text="✓ Outliers removed")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to remove outliers: {str(e)}")
    
    def generate_stl(self):
        if not self.key_processor.key_profile:
            messagebox.showwarning("Warning", "Please extract key profile first (Step 3)")
            return
        
        filepath = filedialog.asksaveasfilename(
            title="Save STL File",
            defaultextension=".stl",
            filetypes=[("STL files", "*.stl"), ("All files", "*.*")]
        )
        
        if filepath:
            try:
                self.key_processor.generate_stl(
                    filepath,
                    depth_scale=self.depth_var.get(),
                    base_thickness=self.base_var.get()
                )
                messagebox.showinfo("Success", 
                    f"✓ STL file generated!\n\n{filepath}\n\nYou can now scale it in your 3D slicer.")
                self.status_label.config(text=f"✓ STL saved: {filepath.split('/')[-1]}")
            except Exception as e:
                messagebox.showerror("Error", f"Failed to generate STL: {str(e)}")
                import traceback
                traceback.print_exc()
    
    def _display_image(self, image):
        if image is None:
            return
        
        self.photo_image = ImageTk.PhotoImage(image)
        self.canvas.delete("all")
        self.canvas.create_image(0, 0, image=self.photo_image, anchor=tk.NW)
        self.canvas.configure(scrollregion=self.canvas.bbox("all"))
    
    def _display_image_with_profile(self):
        if self.key_processor.processed_image and self.key_processor.key_profile:
            display_img = self.key_processor.processed_image.convert('RGB')
            draw = ImageDraw.Draw(display_img)
            
            # Draw key bounds
            if self.key_processor.key_bounds:
                bounds = self.key_processor.key_bounds
                draw.rectangle(
                    [bounds['left'], bounds['top'], bounds['right'], bounds['bottom']],
                    outline='green', width=3
                )
            
            # Draw groove regions
            for groove in self.key_processor.groove_regions:
                draw.rectangle(
                    [groove['left'], groove['top'], groove['right'], groove['bottom']],
                    outline='orange', width=2
                )
            
            # Draw profile
            if len(self.key_processor.key_profile) > 1:
                draw.line(self.key_processor.key_profile, fill='red', width=3)
                
                for x, y in self.key_processor.key_profile[::10]:
                    draw.ellipse([x-2, y-2, x+2, y+2], fill='yellow', outline='red')
            
            self._display_image(display_img)


if __name__ == "__main__":
    print("Key to STL Generator - Manual Selection Tools")
    print("Requirements: pip install scipy pillow numpy")
    print("\nHow to use manual selection:")
    print("1. Load your key image")
    print("2. Click 'Select Key Area' and drag rectangle around key")
    print("3. Click 'Mark Groove Region' and drag rectangles over teeth cuts")
    print("4. Click 'Auto Detect Key' to process with your selections")
    print("5. Refine and export STL")
    
    root = tk.Tk()
    app = KeySTLGUI(root)
    root.mainloop()