import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QLabel, QVBoxLayout, QHBoxLayout, QPushButton,
    QSlider, QFileDialog, QWidget, QComboBox, QSpinBox
)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap, QImage
from skimage.segmentation import active_contour
from scipy import ndimage as ndi
from skimage.filters import threshold_otsu
from skimage.segmentation import watershed
from skimage.feature import peak_local_max

class ImageSegmentationApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Image Segmentation Explorer")
        self.image = None
        self.segmented_image = None

        self.initUI()

    def initUI(self):
        # Main layout
        main_layout = QVBoxLayout()

        # Image display
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(self.image_label)

        # Controls layout
        controls_layout = QHBoxLayout()

        # Load image button
        load_button = QPushButton("Load Image")
        load_button.clicked.connect(self.load_image)
        controls_layout.addWidget(load_button)

        # Segmentation method selection
        self.method_combo = QComboBox()
        self.method_combo.addItems([
            "Thresholding", "Active Contours", "Region-Based", "Watershed"
        ])
        self.method_combo.currentIndexChanged.connect(self.update_parameters)
        controls_layout.addWidget(self.method_combo)

        # Parameter controls
        self.param_label = QLabel("Parameter:")
        controls_layout.addWidget(self.param_label)

        self.param_slider = QSlider(Qt.Horizontal)
        self.param_slider.setMinimum(0)
        self.param_slider.setMaximum(255)
        self.param_slider.setValue(128)
        self.param_slider.valueChanged.connect(self.apply_segmentation)
        controls_layout.addWidget(self.param_slider)

        self.param_spinbox = QSpinBox()
        self.param_spinbox.setMinimum(0)
        self.param_spinbox.setMaximum(255)
        self.param_spinbox.setValue(128)
        self.param_spinbox.valueChanged.connect(self.param_slider.setValue)
        self.param_slider.valueChanged.connect(self.param_spinbox.setValue)
        controls_layout.addWidget(self.param_spinbox)

        main_layout.addLayout(controls_layout)

        # Apply segmentation button
        apply_button = QPushButton("Apply Segmentation")
        apply_button.clicked.connect(self.apply_segmentation)
        main_layout.addWidget(apply_button)

        # Set central widget
        container = QWidget()
        container.setLayout(main_layout)
        self.setCentralWidget(container)

    def load_image(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getOpenFileName(
            self, "Open Image File", "", "Images (*.png *.jpg *.bmp)", options=options
        )
        if file_name:
            self.image = cv2.imread(file_name, cv2.IMREAD_GRAYSCALE)
            self.display_image(self.image)

    def display_image(self, image):
        height, width = image.shape
        bytes_per_line = width
        q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_Grayscale8)
        pixmap = QPixmap.fromImage(q_image)
        self.image_label.setPixmap(pixmap.scaled(
            self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation
        ))

    def update_parameters(self):
        method = self.method_combo.currentText()
        if method == "Thresholding":
            self.param_label.setText("Threshold:")
            self.param_slider.setMaximum(255)
            self.param_slider.setValue(128)
        elif method == "Active Contours":
            self.param_label.setText("Iterations:")
            self.param_slider.setMaximum(500)
            self.param_slider.setValue(100)
        elif method == "Region-Based":
            self.param_label.setText("Sigma:")
            self.param_slider.setMaximum(10)
            self.param_slider.setValue(2)
        elif method == "Watershed":
            self.param_label.setText("Markers:")
            self.param_slider.setMaximum(10)
            self.param_slider.setValue(2)

    def apply_segmentation(self):
        if self.image is None:
            return

        method = self.method_combo.currentText()
        param_value = self.param_slider.value()

        if method == "Thresholding":
            _, self.segmented_image = cv2.threshold(
                self.image, param_value, 255, cv2.THRESH_BINARY
            )
        elif method == "Active Contours":
            s = np.linspace(0, 2 * np.pi, 400)
            x = 100 + 50 * np.cos(s)
            y = 100 + 50 * np.sin(s)
            init = np.array([x, y]).T
            self.segmented_image = active_contour(
                self.image, init, max_iterations=param_value
            )
        elif method == "Region-Based":
            denoised = ndi.gaussian_filter(self.image, sigma=param_value)
            self.segmented_image = denoised > threshold_otsu(denoised)
        elif method == "Watershed":
            distance = ndi.distance_transform_edt(self.image)
            local_maxi = peak_local_max(
                distance, indices=False, footprint=np.ones((3, 3)), labels=self.image
            )
            markers, _ = ndi.label(local_maxi)
            self.segmented_image = watershed(-distance, markers, mask=self.image)

        self.display_image(self.segmented_image.astype(np.uint8) * 255)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ImageSegmentationApp()
    window.show()
    sys.exit(app.exec_())