import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.widgets import LassoSelector
from matplotlib.path import Path
import json
import pydicom

class ROIManager:
    def __init__(self, ax, image):
        self.ax = ax
        self.image = image
        self.rois = []
        self.current_poly = None
        self.history = []
        self.colors = plt.cm.get_cmap('tab10')
        self.color_index = 0

        self.ax.imshow(self.image, cmap='gray')
        self.lasso = LassoSelector(ax, onselect=self.on_select)
        self.cid = self.ax.figure.canvas.mpl_connect('key_press_event', self.on_key_press)

    def on_select(self, verts):
        path = Path(verts)
        color = self.colors(self.color_index % 10)
        self.color_index += 1
        patch = Polygon(verts, closed=True, edgecolor=color, facecolor='none', label=f'ROI {len(self.rois) + 1}')
        self.rois.append({'path': path, 'patch': patch, 'verts': verts})
        self.ax.add_patch(patch)
        self.history.append(('add', patch))
        self.ax.legend()
        self.ax.figure.canvas.draw_idle()

    def on_key_press(self, event):
        if event.key == 'u':  # Undo
            if self.history:
                action, patch = self.history.pop()
                if action == 'add':
                    patch.remove()
                    self.rois = [roi for roi in self.rois if roi['patch'] != patch]
                elif action == 'remove':
                    self.ax.add_patch(patch)
                    self.rois.append({'path': patch.get_path(), 'patch': patch})
                self.ax.legend()
                self.ax.figure.canvas.draw_idle()
        elif event.key == 'd':  # Delete selected ROI
            if self.current_poly:
                confirm = input(f"Do you want to delete {self.current_poly.get_label()}? (y/n): ")
                if confirm.lower() == 'y':
                    self.current_poly.remove()
                    self.rois = [roi for roi in self.rois if roi['patch'] != self.current_poly]
                    self.history.append(('remove', self.current_poly))
                    self.current_poly = None
                    self.ax.legend()
                    self.ax.figure.canvas.draw_idle()
        elif event.key == 's':  # Save ROIs
            self.save_rois()

    def select_roi(self, event):
        for roi in self.rois:
            if roi['path'].contains_point((event.xdata, event.ydata)):
                self.current_poly = roi['patch']
                break

    def save_rois(self):
        roi_data = [{'label': roi['patch'].get_label(), 'vertices': roi['verts']} for roi in self.rois]
        with open('rois.json', 'w') as f:
            json.dump(roi_data, f)
        print("ROIs saved to rois.json")
    def get_rois_data(self):
        roi_data = [{'label': roi['patch'].get_label(), 'vertices': roi['verts']} for roi in self.rois]
        return roi_data


def main():
    image = np.zeros((512, 512))  # Placeholder for MR image
    # Load an image (replace this with actual MR image loading)
    path = r'C:\Users\user\Desktop\knee_seg\LITVYAK_D.I\LITVYAK_D.I\2025-01-22 181006\IMG-0001-00001.dcm'
    file_dcm = pydicom.dcmread(path)
    image = file_dcm.pixel_array

    fig, ax = plt.subplots()
    roi_manager = ROIManager(ax, image)
    fig.canvas.mpl_connect('button_press_event', roi_manager.select_roi)
    plt.show()
    print(roi_manager.get_rois_data())

if __name__ == "__main__":
    main()