"""
Simple tool for annotating tumors in medical images
"""
import os
import cv2
import numpy as np
import argparse
import json
from glob import glob

class TumorAnnotator:
    def __init__(self, images_dir, annotations_dir):
        self.images_dir = images_dir
        self.annotations_dir = annotations_dir
        self.image_files = sorted(glob(os.path.join(images_dir, "*.png")))
        self.current_idx = 0
        self.drawing = False
        self.roi_pt1 = None
        self.roi_pt2 = None
        self.annotations = {}
        
        # Create annotations directory if it doesn't exist
        os.makedirs(annotations_dir, exist_ok=True)
        
        # Load existing annotations if available
        self.load_annotations()
    
    def load_annotations(self):
        """Load existing annotations"""
        for image_file in self.image_files:
            image_name = os.path.basename(image_file)
            ann_file = os.path.join(self.annotations_dir, f"{os.path.splitext(image_name)[0]}.json")
            if os.path.exists(ann_file):
                with open(ann_file, 'r') as f:
                    self.annotations[image_name] = json.load(f)
    
    def save_annotation(self, image_name):
        """Save annotation for current image"""
        if image_name not in self.annotations or not self.annotations[image_name]:
            return
        
        ann_file = os.path.join(self.annotations_dir, f"{os.path.splitext(image_name)[0]}.json")
        with open(ann_file, 'w') as f:
            json.dump(self.annotations[image_name], f)
    
    def mouse_callback(self, event, x, y, flags, param):
        """Mouse callback for drawing bounding boxes"""
        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.roi_pt1 = (x, y)
            self.roi_pt2 = (x, y)
        
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing:
                self.roi_pt2 = (x, y)
        
        elif event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
            self.roi_pt2 = (x, y)
            
            # Add bounding box to annotations
            image_name = os.path.basename(self.image_files[self.current_idx])
            if image_name not in self.annotations:
                self.annotations[image_name] = []
            
            x1, y1 = min(self.roi_pt1[0], self.roi_pt2[0]), min(self.roi_pt1[1], self.roi_pt2[1])
            x2, y2 = max(self.roi_pt1[0], self.roi_pt2[0]), max(self.roi_pt1[1], self.roi_pt2[1])
            
            # Store as [x, y, width, height]
            bbox = [x1, y1, x2 - x1, y2 - y1]
            self.annotations[image_name].append({
                "bbox": bbox,
                "label": "tumor"
            })
    
    def run(self):
        """Run the annotation tool"""
        cv2.namedWindow("Tumor Annotator")
        cv2.setMouseCallback("Tumor Annotator", self.mouse_callback)
        
        while True:
            if self.current_idx >= len(self.image_files) or self.current_idx < 0:
                break
            
            image_path = self.image_files[self.current_idx]
            image_name = os.path.basename(image_path)
            image = cv2.imread(image_path)
            display_image = image.copy()
            
            # Draw existing annotations
            if image_name in self.annotations:
                for ann in self.annotations[image_name]:
                    bbox = ann["bbox"]
                    cv2.rectangle(
                        display_image,
                        (bbox[0], bbox[1]),
                        (bbox[0] + bbox[2], bbox[1] + bbox[3]),
                        (0, 255, 0),
                        2
                    )
            
            # Draw current selection
            if self.drawing:
                cv2.rectangle(
                    display_image,
                    self.roi_pt1,
                    self.roi_pt2,
                    (0, 0, 255),
                    2
                )
            
            # Display image with annotations
            cv2.imshow("Tumor Annotator", display_image)
            
            # Display instructions
            print("\nImage:", image_name)
            print("Controls:")
            print("  Left-click and drag: Draw bounding box")
            print("  n: Next image")
            print("  p: Previous image")
            print("  d: Delete last annotation")
            print("  s: Save annotations")
            print("  q: Quit")
            
            key = cv2.waitKey(1) & 0xFF
            
            if key == ord('n'):  # Next image
                self.save_annotation(image_name)
                self.current_idx = min(self.current_idx + 1, len(self.image_files) - 1)
            
            elif key == ord('p'):  # Previous image
                self.save_annotation(image_name)
                self.current_idx = max(self.current_idx - 1, 0)
            
            elif key == ord('d'):  # Delete last annotation
                if image_name in self.annotations and self.annotations[image_name]:
                    self.annotations[image_name].pop()
            
            elif key == ord('s'):  # Save annotations
                self.save_annotation(image_name)
                print("Annotations saved!")
            
            elif key == ord('q'):  # Quit
                self.save_annotation(image_name)
                break
        
        cv2.destroyAllWindows()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Tumor annotation tool')
    parser.add_argument('--images', type=str, required=True, help='Directory containing images')
    parser.add_argument('--annotations', type=str, default='./annotations', help='Directory to save annotations')
    
    args = parser.parse_args()
    
    annotator = TumorAnnotator(args.images, args.annotations)
    annotator.run()