# -*- encoding:utf-8 -*-
import sys
import cv2
import numpy as np
import os
import subprocess
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QSlider, QLabel
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap, QImage


class ImageComparisonApp(QWidget):
    def __init__(self, inpath, outpath):
        self.inpath = inpath
        self.outpath = outpath
        super().__init__()

        # 初始化变量
        self.original_img = cv2.imread(self.inpath)
        self.processed_img = cv2.imread(self.outpath)

        # 检查图像是否成功加载
        if self.original_img is None:
            print(f"Error: Unable to load original image: {self.inpath}")
            sys.exit(1)
        if self.processed_img is None:
            print(f"Error: Unable to load processed image: {self.outpath}")
            sys.exit(1)

        # 调整两张图像大小以匹配
        target_height = max(self.original_img.shape[0], self.processed_img.shape[0])
        self.original_resized = cv2.resize(
            self.original_img,
            (
                int(
                    self.original_img.shape[1]
                    * target_height
                    / self.original_img.shape[0]
                ),
                target_height,
            ),
        )
        self.processed_resized = cv2.resize(
            self.processed_img,
            (
                int(
                    self.processed_img.shape[1]
                    * target_height
                    / self.processed_img.shape[0]
                ),
                target_height,
            ),
        )

        self.initUI()

    def initUI(self):
        self.setWindowTitle("Image Comparison")
        self.setGeometry(100, 100, 800, 600)

        self.slider = QSlider(Qt.Horizontal)
        self.slider.setRange(0, self.original_resized.shape[1])
        self.slider.setValue(self.original_resized.shape[1] // 2)
        self.slider.setTickPosition(QSlider.TicksBelow)
        self.slider.setTickInterval(1)
        self.slider.valueChanged.connect(self.update_image)

        layout = QVBoxLayout()
        self.image_label = QLabel()
        layout.addWidget(self.image_label)
        layout.addWidget(self.slider)

        self.setLayout(layout)

    def update_image(self):
        split_pos = self.slider.value()
        left_img = self.original_resized[:, :split_pos]
        right_img = self.processed_resized[:, split_pos:]
        combined_img = np.hstack((left_img, right_img))
        height, width, channel = combined_img.shape
        bytes_per_line = 3 * width
        q_image = QImage(
            combined_img.data, width, height, bytes_per_line, QImage.Format_BGR888
        )
        pixmap = QPixmap(q_image)
        self.image_label.setPixmap(pixmap)


if __name__ == "__main__":
    inpath = input("Enter input image path: ")
    outpath = input("Enter output image path: ")
    output_image = os.path.join(outpath, "output.png")

    model_path = "models"

    command = [
        "realesrgan-ncnn-vulkan.exe",
        "-i",
        inpath,
        "-o",
        output_image,
        "-s",
        "4",
        "-m",
        model_path,
        "-n",
        "realesrgan-x4plus",
    ]

    try:
        subprocess.run(command, check=True)
        print(f"Image processed successfully. Output saved to: {output_image}")
        app = QApplication(sys.argv)
        ex = ImageComparisonApp(inpath, output_image)
        ex.show()
        sys.exit(app.exec_())
    except subprocess.CalledProcessError as e:
        print(f"Error occurred during processing: {e}")
        sys.exit(1)