from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsView,QMessageBox
import subprocess,shutil,os, urllib.request
from PyQt5.QtGui import QPixmap
import pandas as pd
import matplotlib.pyplot as plt
import os
import glob
import subprocess
import imaplib
from PIL import Image
import copy,cv2
import numpy as np
destination_folder = "cell_application/uploaded_images/"
if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)
output_folder="cell_application/output"
if not os.path.exists(output_folder):
        os.makedirs(output_folder)
output_red="cell_application/outred"
if not os.path.exists(output_red):
        os.makedirs(output_red)
output_pl="cell_application/outpl"
if not os.path.exists(output_pl):
        os.makedirs(output_pl)
model_path = "./models/yolov7"  # specify your model path

if not os.path.exists(model_path):
    
    # Create the directory if it doesn't exist
    os.makedirs(model_path, exist_ok=True)

    # Clone the repository
    subprocess.check_call(['git', 'clone', 'https://github.com/WongKinYiu/yolov7', model_path])

    # Install the requirements
    subprocess.check_call(['pip', 'install', '-r', os.path.join(model_path, 'requirements.txt')])
class GraphicsView(QGraphicsView):
    def __init__(self, parent=None):
        super(GraphicsView, self).__init__(parent)
        self._isPanning = False
        self._panStartX = 0
        self._panStartY = 0

    def wheelEvent(self, event):
        factor = 1.15 if event.angleDelta().y() > 0 else 1 / 1.15
        self.scale(factor, factor)

    def mousePressEvent(self, event):
        if event.button() == QtCore.Qt.MiddleButton:
            self._isPanning = True
            self._panStartX = event.x()
            self._panStartY = event.y()
            self.setCursor(QtCore.Qt.ClosedHandCursor)
            event.accept()
        else:
            super(GraphicsView, self).mousePressEvent(event)

    def mouseMoveEvent(self, event):
        if self._isPanning:
            self.horizontalScrollBar().setValue(self.horizontalScrollBar().value() - (event.x() - self._panStartX))
            self.verticalScrollBar().setValue(self.verticalScrollBar().value() - (event.y() - self._panStartY))
            self._panStartX = event.x()
            self._panStartY = event.y()
            event.accept()
        else:
            super(GraphicsView, self).mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        if event.button() == QtCore.Qt.MiddleButton:
            self._isPanning = False
            self.setCursor(QtCore.Qt.ArrowCursor)
            event.accept()
        else:
            super(GraphicsView, self).mouseReleaseEvent(event)
class Ui_Dialog(object):
    def setupUi(self, Dialog):
        Dialog.setObjectName("Dialog")
        Dialog.setMinimumSize(800, 600)
        Dialog.resize(1147, 840)

        # Create a vertical layout for the main window
        main_layout = QtWidgets.QVBoxLayout(Dialog)

        # Create a horizontal layout for the buttons
        button_layout = QtWidgets.QHBoxLayout()
        main_layout.addLayout(button_layout)

        # Create the upload button and add it to the button layout
        self.uploadB = QtWidgets.QPushButton("upload")
        self.uploadB.setObjectName("uploadB")
        button_layout.addWidget(self.uploadB)

        # Create a stretchable spacer and add it to the button layout
        button_layout.addStretch()
        
        # Create an input field and add it to the button layout
        self.inputField = QtWidgets.QLineEdit()
        self.inputField.setObjectName("inputField")
        button_layout.addWidget(self.inputField)

       

        # Create labels for the results and add them to the results layout
        # Create a horizontal layout for the results
        self.results_layout = QtWidgets.QHBoxLayout()
        main_layout.addLayout(self.results_layout)

        # Create labels for the results and add them to the results layout
        self.results = {'Basophil':0, 'Eosinophil':0, 'Erythroblast':0, 'Ig':0, 'Lymphocyte':0, 'Monocyte':0, 'Neutrophil':0,'White Cells':0, 'Red Cells':0, 'platelets':0}
        self.labels = {}  # Dictionary to store the QLabel objects

        i=0
        for d in self.results:
            result_label = QtWidgets.QLabel()
            result_label.setObjectName(f"result{i+1}")
            result_label.setText(f"{d} {self.results[d]}")  # Set some initial text
            self.results_layout.addWidget(result_label)
            self.labels[d] = result_label  # Store the QLabel object in self.labels
            i+=1
           
       # Create the extract_lymphocyte button and add it to the button layout
        self.extract_lymphocyte = QtWidgets.QPushButton("extract_lymphocyte")
        self.extract_lymphocyte.setObjectName("extract_lymphocyte")
        button_layout.addWidget(self.extract_lymphocyte)

        self.applypl = QtWidgets.QPushButton("pl")
        self.applypl.setObjectName("applypl")
        button_layout.addWidget(self.applypl)

        # Create the apply button and add it to the button layout
        self.applyred = QtWidgets.QPushButton("red")
        self.applyred.setObjectName("applyb")
        button_layout.addWidget(self.applyred)

        # Create the apply button and add it to the button layout
        self.applyb = QtWidgets.QPushButton("white")
        self.applyb.setObjectName("applyb")
        button_layout.addWidget(self.applyb)
        
      

       # Create the QGraphicsView and add it to the main layout
        self.graphicsView = GraphicsView()
        main_layout.addWidget(self.graphicsView, stretch=3)

        # Create a QGraphicsScene for the QGraphicsView
        self.scene = QGraphicsScene()
        self.graphicsView.setScene(self.scene)

        self.graphicsViewly = GraphicsView()
        main_layout.addWidget(self.graphicsViewly, stretch=1)

        # Create a QGraphicsScene for the second QGraphicsView
        self.scenely = QGraphicsScene()
        self.graphicsViewly.setScene(self.scenely)
        
    

        self.uploadB.clicked.connect(self.handle_upload_button_clicked)
        self.extract_lymphocyte.clicked.connect(self.handle_extract_lymphocyte_button_clicked)
        self.applyb.clicked.connect(self.handle_detect_button_clicked)
        self.applyred.clicked.connect(self.handle_red_button_clicked)
        self.applypl.clicked.connect(self.handle_pl_button_clicked)
    def change_image(self, image_path):
        # Create a QPixmap from the new image path
        pixmap = QtGui.QPixmap(image_path)

        # Clear the QGraphicsScene
        self.scene.clear()

        # Add the QPixmap to the QGraphicsScene
        self.scene.addPixmap(pixmap)

        # Fit the QGraphicsView to the scene's content
        self.graphicsView.fitInView(self.scene.itemsBoundingRect(), QtCore.Qt.KeepAspectRatio)
        
    def handle_upload_button_clicked(self):
        global take_lymphocyte, yolo_img_resize
        
        # Open a file dialog and get the selected image file path
        file_path, _ = QtWidgets.QFileDialog.getOpenFileName(None, 'Select Image')
        global upload_name

        if file_path:
            upload_name  = os.path.basename(file_path)
            self.change_image(file_path)

            # Load the image into a QPixmap object
            pixmap = QPixmap(file_path)

            # Get the width and height
            width = pixmap.width()
            height = pixmap.height()
            ma = max(width, height)
            yolo_img_resize = max((ma // 640) *640 ,640)
            #yolo_img_resize = max((ma // 384) *384 ,384)
            #yolo_img_resize = (yolo_img_resize // 32) * 32 + (yolo_img_resize % 32 != 0) * 32
            self.inputField.setText(str(yolo_img_resize))

            # Define the destination folders
            destination_folder = "cell_application/uploaded_images/"

            source = file_path
            destination = os.path.join(destination_folder, upload_name)

            if os.path.normpath(source) != os.path.normpath(destination):
                # Copy the file
                shutil.copy(source, destination)
            else:
                # If source and destination are the same, delete and copy
                if os.path.exists(destination):
                    os.remove(destination)
                    shutil.copy(source, destination)
            
            # If the image is not a png or jpg make it jpg then save it
            if not (file_path.lower().endswith(('.png', '.jpg' ))):
               
                im = Image.open(file_path)
                rgb_im = im.convert('RGB')
                rgb_im.save(file_path[:-4] + 'jpg')
            lb = ['Basophil', 'Eosinophil', 'Erythroblast', 'Ig', 'Lymphocyte', 'Monocyte', 'Neutrophil','White Cells', 'Red Cells', 'platelets']
            for l in lb:
                self.labels[l].setText(l+': 0')

            

    
    def handle_detect_button_clicked(self):
        global yolo_img_resize
        yolo_img_resize = int(self.inputField.text())
        global take_lymphocyte
        image_path = "cell_application/uploaded_images/"+upload_name
        print(upload_name)
       
        model_path="models_weights/lastv7_aug_0.978_4.pt"
        command = f"python ./models/yolov7/detect.py --source {image_path} --weights {model_path} --img {yolo_img_resize} --conf 0.5 --project {output_folder} --save-txt"
        process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE)
        process.wait()
        
        if os.path.exists("cell_application/output/exp"):
            if os.path.exists("cell_application/output/"+upload_name[0:upload_name.find('.')]):
                 shutil.rmtree("cell_application/output/"+upload_name[0:upload_name.find('.')])
            os.rename("cell_application/output/exp","cell_application/output/"+upload_name[0:upload_name.find('.')])

        image_out_path="cell_application/output/"+upload_name[0:upload_name.find('.')]+"/"+upload_name
        take_lymphocyte=True
       
        txt_file = glob.glob(os.path.join("cell_application/output/"+upload_name[0:upload_name.find('.')]+"/labels/", '*.txt'))[0]
        
  
       
        df = pd.read_csv(txt_file,delimiter=' ', header=None)
        df.columns = ['label', 'x_center', 'y_center', 'width', 'height']
      
        # Count the appearance of each label in the DataFrame
        label_counts = df['label'].value_counts()
       
        # Map the label numbers to their names
        label_names = {0: 'Basophil', 1: 'Eosinophil', 2: 'Erythroblast', 3: 'Ig', 4: 'Lymphocyte', 5: 'Monocyte', 6: 'Neutrophil'}

        # Update the results dictionary with the counts
        w=0
        for label_number, label_name in label_names.items():
            if label_number in label_counts:
                self.labels[label_name].setText(label_name+": "+str(label_counts[label_number]))
                w+=label_counts[label_number]
                
    
        self.labels['White Cells'].setText('White Cells: '+str(w))
       
      
        
   
        
        print(image_out_path)
        self.change_image(image_out_path)
    def handle_red_button_clicked(self):
        img_size = int(self.inputField.text())
        
        image_path = "cell_application/uploaded_images/"+upload_name
        
        model_path="models_weights/bestv7_red.pt"
        command = f"python ./models/yolov7/detect.py --source {image_path} --weights {model_path} --img {img_size} --conf 0.5 --project {output_red} --save-txt"
        process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE)
        process.wait()
        if os.path.exists("cell_application/outred/exp"):
            if os.path.exists("cell_application/outred/"+upload_name[0:upload_name.find('.')]):
                 shutil.rmtree("cell_application/outred/"+upload_name[0:upload_name.find('.')])
            os.rename("cell_application/outred/exp","cell_application/outred/"+upload_name[0:upload_name.find('.')])
        image_out_path="cell_application/outred/"+upload_name[0:upload_name.find('.')]+"/"+upload_name
        txt_file = glob.glob(os.path.join("cell_application/outred/"+upload_name[0:upload_name.find('.')]+"/labels/", '*.txt'))[0]
        df = pd.read_csv(txt_file,delimiter=' ', header=None)
        df.columns = ['label', 'x_center', 'y_center', 'width', 'height']
       
        
    
        self.labels['Red Cells'].setText('Red Cells: '+str(len(df)))
       
      
     
        
        self.change_image(image_out_path)
    def handle_pl_button_clicked(self):
        img_size = int(self.inputField.text())
        
        image_path = "cell_application/uploaded_images/"+upload_name
        print(upload_name)
        model_path="models_weights/bestv7_pl_640_3.pt"
        command = f"python ./models/yolov7/detect.py --source {image_path} --weights {model_path} --img {img_size} --conf 0.5 --project {output_pl} --save-txt"
        process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE)
        process.wait()
        if os.path.exists("cell_application/outpl/exp"):
            if os.path.exists("cell_application/outpl/"+upload_name[0:upload_name.find('.')]):
                 shutil.rmtree("cell_application/outpl/"+upload_name[0:upload_name.find('.')])
            os.rename("cell_application/outpl/exp","cell_application/outpl/"+upload_name[0:upload_name.find('.')])
        image_out_path="cell_application/outpl/"+upload_name[0:upload_name.find('.')]+"/"+upload_name
        txt_file = glob.glob(os.path.join("cell_application/outpl/"+upload_name[0:upload_name.find('.')]+"/labels/", '*.txt'))[0]
        df = pd.read_csv(txt_file,delimiter=' ', header=None)
        df.columns = ['label', 'x_center', 'y_center', 'width', 'height']
       
        
    
        self.labels['platelets'].setText('platelets: '+str(len(df)))
       
      
        
        
        self.change_image(image_out_path)
    

    def handle_extract_lymphocyte_button_clicked(self):
            global take_lymphocyte
                
                

            def get_most_recently_modified_folder(rootDir):
                # Create a list to store tuples of (folder_path, last_modified_time)
                folders = [(os.path.join(rootDir, dir), os.path.getmtime(os.path.join(rootDir, dir)))
                        for dir in os.listdir(rootDir) if os.path.isdir(os.path.join(rootDir, dir))]

                # Sort the list by last_modified_time
                folders.sort(key=lambda x: x[1], reverse=True)

                # The first element in the list is the most recently modified folder
                most_recently_modified_folder = folders[0][0] if folders else None
               
                return most_recently_modified_folder

            if not take_lymphocyte:
                msg = QMessageBox()
                msg.setIcon(QMessageBox.Information)
                msg.setText("You should make a detection for white blood cell first press White button.")
                msg.setWindowTitle("Information")
                msg.exec_()
                
               
                
            
            else :
                main_folder_path = "cell_application/output"
                lastfolder_path=get_most_recently_modified_folder(main_folder_path)
            
                
                txt_file = glob.glob(os.path.join(lastfolder_path+"/labels/", '*.txt'))[0]
                
                df = pd.read_csv(txt_file,delimiter=' ', header=None)
                df.columns = ['label', 'x_center', 'y_center', 'width', 'height']
                df=df[df['label']==4]
                imagename=txt_file[txt_file.rfind('\\')+1:txt_file.find('.')]
                imagepath=os.path.join("cell_application/uploaded_images",imagename)
                
                image = Image.open(imagepath+".jpg")
                npimage=np.array(image)
                # Get the original size of the image
                ma=max(npimage.shape[0],npimage.shape[1])
                yolo_img_resize= (ma*2144) //2592
                yolo_img_resize=(yolo_img_resize//32)*32 +(yolo_img_resize%32!=0)*32
            
                height, width = yolo_img_resize,yolo_img_resize
                height_i, width_i=npimage.shape[:2]
            
                height_scale = height_i/height
                width_scale=width_i/width
                df['x_center'] *= width_i
                df['y_center'] *= height_i
                df['width'] *= width_i
                df['height'] *= height_i
                    
                # Calculate the top-left and bottom-right corners of the bounding boxes
                df1 = pd.DataFrame()
                df1['xmin'] = df['x_center'] - df['width'] / 2
                df1['ymin'] = df['y_center'] - df['height'] / 2
                df1['xmax'] = df['x_center'] + df['width'] / 2
                df1['ymax'] = df['y_center'] + df['height'] / 2
                
                new_box=[]
                crop_dir = lastfolder_path+"/crop_results"
                sam_dir = lastfolder_path+"/sam_results"

                # Check if the directory already exists
                if not os.path.exists(crop_dir):
                    # If not, create the directory
                    os.mkdir(crop_dir)
                if not os.path.exists(sam_dir):
                    # If not, create the directory
                    os.mkdir(sam_dir)
                boxes = [np.array(row) for row in df1.values]
                for i in range(len(boxes)):
                    box =copy.copy(boxes[i])
                    padding = 70
                    box[0] = max(0, box[0] - padding)
                    box[1] = max(0, box[1] - padding)
                    box[2] = min(image.width, box[2] + padding)
                    box[3] = min(image.height, box[3] + padding)

                    # Crop the image using the bounding box coordinates
                    cropped_image1 = np.array(image.crop(boxes[i]))
                    cropped_image = image.crop(box)
                    cropped_image.save(f'{lastfolder_path}/crop_results/{imagename}{i}.jpg')
                    # Save the cropped image
                    # Calculate the new bounding box coordinates
                    new_box.append(np.array([ padding,  padding,  padding+cropped_image1.shape[1],   padding+cropped_image1.shape[0]]))
                box = new_box
                import torch
                CHECKPOINT_PATH="models_weights/sam_vit_h_4b8939.pth"
                DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
                MODEL_TYPE = "vit_h"
                import cv2
                    
                from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
                sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH)
                mask_predictor = SamPredictor(sam)
                mask_generator = SamAutomaticMaskGenerator(sam)
                from keras.models import load_model
                import cv2
                import supervision as sv
                # Load the model from the .h5 file
                model = load_model('models_weights/EfficientNetB3-leukemia-0.96.h5')
                i=-1
                image_list=[]
                cropimagespath=lastfolder_path+"/crop_results"
                for file in os.listdir(cropimagespath):
                
                    i+=1
                    image_bgr = cv2.imread(os.path.join(cropimagespath,file))
                    

                    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

                    mask_predictor.set_image(image_rgb)

                    masks, scores, logits = mask_predictor.predict(
                        box=new_box[i],
                        multimask_output=True
                    )
                        

                    img=np.zeros(image_bgr.shape)

                    img[:,:,:]=[0,0,0]
                    img[masks[2,:,:]]=image_bgr[masks[2,:,:]]
                    #imgtostore=img.copy()
                    #cv2.resize(imgtostore,(300,300))
                    cv2.imwrite(f'{sam_dir}/{imagename}{i}.jpg',img)
                
                    annotated_image = img.astype(np.uint8)
                
                
                    
                    """sam_result = mask_generator.generate(img)
                    mask_annotator = sv.MaskAnnotator()
                    
                    detections = sv.Detections.from_sam(sam_result=sam_result)

                    annotated_image = mask_annotator.annotate(scene=img.copy(), detections=detections)
                    print("ok2")"""
                    
                

                        
                    img_size_ly = (300, 300)
                    
                    def evaluate(model ,img):
                        

                        img = cv2.resize(img, img_size_ly)
                        

                        img = np.expand_dims(img, axis=0)
                        
                        # now predict the image
                        pred = model.predict(img,verbose=0)
                        print('the shape of prediction is ', pred.shape)
                        print(pred)
                
                        index = np.argmax(pred[0])

                        return index
                    image_list.append([(evaluate(model,annotated_image)),f'{sam_dir}/{imagename}{i}.jpg'])
            
                print(image_list)
                x_offset = 0 
                label_dict = {0: "all", 1: "hem"}
                for image_info in image_list:
                    label_number, image_path = image_info
                    label = label_dict[label_number]
                    pixmap = QtGui.QPixmap(image_path)

                    # Scale the pixmap without preserving aspect ratio
                    scaled_pixmap = pixmap.scaled(300, 300)

                    item = QtWidgets.QGraphicsPixmapItem(scaled_pixmap)
                    item.setPos(x_offset, 0)
                    self.scenely.addItem(item)

                    # Create a QGraphicsTextItem for the label
                    text_item = QtWidgets.QGraphicsTextItem(label)

                    # Set the color of the text to white
                    text_item.setDefaultTextColor(QtGui.QColor('white'))

                    # Calculate the center position of the image for x-coordinate
                    center_x = x_offset + scaled_pixmap.width() / 2 - text_item.boundingRect().width() / 2

                    # Set the position of the text item on the image and centered horizontally
                    text_item.setPos(center_x, scaled_pixmap.height() / 2)
                    
                    self.scenely.addItem(text_item)

                    # Update x offset for next image
                    x_offset += scaled_pixmap.width()

                

                




                 
                 
                 
    
            
      


        



if __name__ == "__main__":
    import sys
    upload_name = ""
    yolo_img_resize=1216
    lastfolder_path=""
    take_lymphocyte=False
    app = QtWidgets.QApplication(sys.argv)
    Dialog = QtWidgets.QDialog()
    ui = Ui_Dialog()
    ui.setupUi(Dialog)
    
    # Add a maximize button to the window
    Dialog.setWindowFlag(QtCore.Qt.WindowMaximizeButtonHint, True)

    Dialog.show()
    sys.exit(app.exec_())