import sys
import numpy as np
import cv2
from PyQt5.QtWidgets import (
    QApplication, QWidget, QLabel, QSlider, QVBoxLayout,
    QHBoxLayout, QGridLayout, QPushButton, QFileDialog
)
from PyQt5.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import scipy.signal

import os
os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)


def generate_kernel(radius, sigma=None):
    """
    Generate a 2D Gaussian kernel with a given radius.

    Parameters:
    - radius: int, the radius of the kernel (size will be 2*radius + 1)
    - sigma: float (optional), standard deviation of the Gaussian. If None, sigma = radius / 3

    Returns:
    - kernel: 2D numpy array of shape (2*radius+1, 2*radius+1)
    """
    size = 2 * radius + 1
    if sigma is None:
        sigma = radius / 3.0  # Common default choice

    print(f"radius: {radius}, sigma: {sigma}")

    # Create a grid of (x,y) coordinates
    ax = np.arange(-radius, radius + 1)
    xx, yy = np.meshgrid(ax, ax)

    # Apply the 2D Gaussian formula
    kernel = np.exp(-(xx**2 + yy**2) / (2 * sigma**2))
    kernel /= 2 * np.pi * sigma**2  # Normalize based on Gaussian PDF
    kernel /= kernel.sum()          # Normalize to sum to 1

    return kernel


class KernelVisualizer(QWidget):
    def __init__(self, image_path=None):
        super().__init__()
        self.setWindowTitle("Gaussian Kernel Visualizer")
        self.image = None

        self.load_button = QPushButton("Load Image")
        self.load_button.clicked.connect(self.load_image)

        self.radius_slider = QSlider(Qt.Horizontal)
        self.radius_slider.setRange(1, 30)
        self.radius_slider.setValue(5)

        self.sigma_slider = QSlider(Qt.Horizontal)
        self.sigma_slider.setRange(1, 100)
        self.sigma_slider.setValue(15)

        self.radius_slider.valueChanged.connect(self.update_visualization)
        self.sigma_slider.valueChanged.connect(self.update_visualization)

        self.kernel_fig = Figure(figsize=(3, 3))
        self.kernel_canvas = FigureCanvas(self.kernel_fig)

        self.image_fig = Figure(figsize=(6, 3))
        self.image_canvas = FigureCanvas(self.image_fig)

        layout = QVBoxLayout()
        layout.addWidget(self.load_button)

        sliders_layout = QGridLayout()
        sliders_layout.addWidget(QLabel("Radius:"), 0, 0)
        sliders_layout.addWidget(self.radius_slider, 0, 1)
        sliders_layout.addWidget(QLabel("Sigma:"), 1, 0)
        sliders_layout.addWidget(self.sigma_slider, 1, 1)

        layout.addLayout(sliders_layout)
        layout.addWidget(QLabel("Kernel Visualization:"))
        layout.addWidget(self.kernel_canvas)
        layout.addWidget(QLabel("Original and Deconvolved Image:"))
        layout.addWidget(self.image_canvas)

        self.setLayout(layout)
        
        if image_path:
            self.load_image(image_path)
        else:
            self.update_visualization()

    def load_image(self, image_path=None):
        if not image_path:
            fname, _ = QFileDialog.getOpenFileName(self, "Open Image", "", "Images (*.png *.jpg *.bmp *.jpeg)")
            image_path = fname
        
        if image_path:
            img = cv2.imread(image_path)
            if img is not None:
                self.image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                self.update_visualization()

    def update_visualization(self):
        radius = self.radius_slider.value()
        sigma = self.sigma_slider.value() / 10.0
        kernel = generate_kernel(radius, sigma)

        # Kernel Visualization
        self.kernel_fig.clear()
        ax = self.kernel_fig.add_subplot(111)
        cax = ax.imshow(kernel, cmap='hot')
        self.kernel_fig.colorbar(cax, ax=ax)
        ax.set_title(f"Kernel (r={radius}, σ={sigma:.2f})")
        self.kernel_canvas.draw()

        if self.image is not None:
            kernel_size = 2 * radius + 1

            # Apply row-by-row deconvolution
            deconvolved_image = np.zeros_like(self.image)

            # Perform row-wise deconvolution with padded kernel
            for i in range(self.image.shape[0]):
                padded_kernel = np.pad(kernel, ((0, self.image.shape[1] - kernel_size), (0, 0)), mode='constant')
                deconvolved_image[i, :], _ = scipy.signal.deconvolve(self.image[i, :], padded_kernel[i, :])

            # Perform column-wise deconvolution with padded kernel
            for j in range(self.image.shape[1]):
                padded_kernel = np.pad(kernel, ((0, self.image.shape[0] - kernel_size), (0, 0)), mode='constant')
                deconvolved_image[:, j], _ = scipy.signal.deconvolve(self.image[:, j], padded_kernel[:, j])

            deconvolved_image = np.clip(deconvolved_image, 0, 255).astype(np.uint8)  # Ensure valid range

            self.image_fig.clear()
            ax1 = self.image_fig.add_subplot(121)
            ax1.imshow(self.image, cmap='gray')
            ax1.set_title("Original")
            ax1.axis('off')

            ax2 = self.image_fig.add_subplot(122)
            ax2.imshow(deconvolved_image, cmap='gray')
            ax2.set_title("Deconvolved")
            ax2.axis('off')

            self.image_canvas.draw()
        else:
            self.image_fig.clear()
            ax = self.image_fig.add_subplot(111)
            ax.text(0.5, 0.5, "No image loaded", fontsize=14, ha='center', va='center')
            ax.axis('off')
            self.image_canvas.draw()


if __name__ == "__main__":
    image_path = None
    if len(sys.argv) > 1:
        image_path = sys.argv[1]  # Get image path from command-line argument
        print(image_path)

    app = QApplication(sys.argv)
    viewer = KernelVisualizer(image_path=image_path)
    viewer.show()
    sys.exit(app.exec_())