import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)

import datetime
import gc
from pathlib import Path

import torch
from torch.cuda import empty_cache

from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PyQt5.QtWidgets import (QFileDialog, QMessageBox, QVBoxLayout,
                             QLineEdit, QDialog, QHBoxLayout, QPushButton,
                             QSpinBox, QFormLayout, QComboBox,
                             QGroupBox, QLabel, QDoubleSpinBox)

from coralnet_toolbox.MachineLearning.SMP import SemanticModel

from coralnet_toolbox.Icons import get_icon


# ----------------------------------------------------------------------------------------------------------------------
# Classes
# ----------------------------------------------------------------------------------------------------------------------


class EvaluateModelWorker(QThread):
    """
    Worker thread for evaluating a model.

    Signals:
        evaluation_started: Emitted when the evaluation starts.
        evaluation_completed: Emitted when the evaluation completes.
        evaluation_error: Emitted when an error occurs during evaluation.
    """
    evaluation_started = pyqtSignal()
    evaluation_completed = pyqtSignal()
    evaluation_error = pyqtSignal(str)

    def __init__(self, model, params):
        """
        Initialize the EvaluateModelWorker.

        Args:
            model: The model to be evaluated.
            params: A dictionary of parameters for evaluation.
        """
        super().__init__()
        self.model = model
        self.params = params

    def run(self):
        """
        Run the evaluation process in a separate thread.
        """
        try:
            # Emit signal to indicate evaluation has started
            self.evaluation_started.emit()

            # Evaluate the model
            self.model.eval(
                data_yaml=self.params['data_yaml'],
                split=self.params['split'],
                num_vis_samples=self.params['num_vis_samples'],
                output_dir=self.params['save_dir'],
                device=self.params['device'],
                imgsz=self.params.get('imgsz'),
                batch=self.params.get('batch', 1),
                confidence=self.params.get('conf', 0.5)
            )

            # Emit signal to indicate evaluation has completed
            self.evaluation_completed.emit()

        except Exception as e:
            self.evaluation_error.emit(str(e))
            
            
class Semantic(QDialog):
    def __init__(self, main_window, parent=None):
        super().__init__(parent)
        self.main_window = main_window
        
        self.setWindowIcon(get_icon("coral.png"))
        self.setWindowTitle("Evaluate SemanticModel")
        self.resize(400, 600)  # Increased height for additional parameters
        
        self.setWindowFlags(Qt.Window |
                            Qt.WindowCloseButtonHint |
                            Qt.WindowMinimizeButtonHint |
                            Qt.WindowMaximizeButtonHint |
                            Qt.WindowTitleHint)
        
        self.imgsz = 640
        self.task = None
        self.params = {}
        self.class_mapping = {}
        
        self.layout = QVBoxLayout(self)
        
        # Setup the info layout
        self.setup_info_layout()
        # Setup the dataset layout
        self.setup_dataset_layout()
        # Setup the output layout
        self.setup_output_layout()
        # Setup the parameters layout
        self.setup_parameters_layout()
        # Setup the buttons layout
        self.setup_buttons_layout()
        
    def setup_info_layout(self):
        """Set up the layout and widgets for the info layout."""
        group_box = QGroupBox("Information")
        layout = QVBoxLayout()
        
        # Create a QLabel with explanatory text and hyperlink
        info_label = QLabel("Evaluate a Semantic model on a dataset, and the split you want to evaluate on.")
        
        info_label.setOpenExternalLinks(True)
        info_label.setWordWrap(True)
        layout.addWidget(info_label)
        
        group_box.setLayout(layout)
        self.layout.addWidget(group_box)
        
    def setup_dataset_layout(self):
        """Setup the dataset layout."""
        group_box = QGroupBox("Dataset")
        layout = QFormLayout()
        
        self.model_edit = QLineEdit()
        self.model_button = QPushButton("Browse...")
        self.model_button.clicked.connect(self.browse_model_file)
        model_layout = QHBoxLayout()
        model_layout.addWidget(self.model_edit)
        model_layout.addWidget(self.model_button)
        layout.addRow("Existing Model:", model_layout)
            
        self.dataset_edit = QLineEdit()
        self.dataset_button = QPushButton("Browse...")
        self.dataset_button.clicked.connect(self.browse_dataset_yaml)
        dataset_layout = QHBoxLayout()
        dataset_layout.addWidget(self.dataset_edit)
        dataset_layout.addWidget(self.dataset_button)
        layout.addRow("Dataset YAML:", dataset_layout)
        
        group_box.setLayout(layout)
        self.layout.addWidget(group_box)
        
    def setup_output_layout(self):
        """Setup the output layout."""
        group_box = QGroupBox("Output Parameters")
        layout = QFormLayout()
        
        self.save_dir_edit = QLineEdit()
        self.save_dir_button = QPushButton("Browse...")
        self.save_dir_button.clicked.connect(self.browse_save_dir)
        save_dir_layout = QHBoxLayout()
        save_dir_layout.addWidget(self.save_dir_edit)
        save_dir_layout.addWidget(self.save_dir_button)
        layout.addRow("Save Directory:", save_dir_layout)
        
        self.name_edit = QLineEdit()
        layout.addRow("Name:", self.name_edit)
        
        self.split_combo = QComboBox()
        self.split_combo.addItems(["train", "val", "test"])
        self.split_combo.setCurrentText("test")
        layout.addRow("Split:", self.split_combo)
        
        group_box.setLayout(layout)
        self.layout.addWidget(group_box)
        
    def setup_parameters_layout(self):
        """Setup the parameters layout."""
        group_box = QGroupBox("Parameters")
        layout = QFormLayout()
        
        # Image size
        self.imgsz_spinbox = QSpinBox()
        self.imgsz_spinbox.setMinimum(16)
        self.imgsz_spinbox.setMaximum(4096)
        self.imgsz_spinbox.setValue(self.imgsz)
        layout.addRow("Image Size:", self.imgsz_spinbox)
        
        # Batch size
        self.batch_spinbox = QSpinBox()
        self.batch_spinbox.setMinimum(1)
        self.batch_spinbox.setMaximum(1024)
        self.batch_spinbox.setValue(16)
        layout.addRow("Batch:", self.batch_spinbox)
        
        # Confidence threshold
        self.conf_spinbox = QDoubleSpinBox()
        self.conf_spinbox.setMinimum(0.0)
        self.conf_spinbox.setMaximum(1.0)
        self.conf_spinbox.setSingleStep(0.001)
        self.conf_spinbox.setDecimals(3)
        self.conf_spinbox.setValue(0.001)
        layout.addRow("Confidence:", self.conf_spinbox)
        
        # Number of visualization samples
        self.num_vis_samples_spinbox = QSpinBox()
        self.num_vis_samples_spinbox.setMinimum(1)
        self.num_vis_samples_spinbox.setMaximum(100)
        self.num_vis_samples_spinbox.setValue(5)
        layout.addRow("Num Vis Samples:", self.num_vis_samples_spinbox)
        
        group_box.setLayout(layout)
        self.layout.addWidget(group_box)
        
    def setup_buttons_layout(self):
        """Setup the buttons layout."""        
        self.buttons = QPushButton("OK")
        self.buttons.clicked.connect(self.accept)
        self.layout.addWidget(self.buttons)
        self.cancel_button = QPushButton("Cancel")
        self.cancel_button.clicked.connect(self.reject)
        self.layout.addWidget(self.cancel_button)

    def browse_model_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select Model File")
        if file_path:
            self.model_edit.setText(file_path)

    def browse_dataset_yaml(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select YAML File", "", "YAML Files (*.yaml)")
        if file_path:
            self.dataset_edit.setText(file_path)

    def browse_save_dir(self):
        dir_path = QFileDialog.getExistingDirectory(self, "Select Save Directory")
        if dir_path:
            self.save_dir_edit.setText(dir_path)

    def accept(self):
        self.evaluate_model()
        super().accept()

    def get_evaluation_parameters(self):
        params = {
            'exist_ok': True,
        }
        params['model'] = self.model_edit.text()
        params['data'] = self.dataset_edit.text()
        params['save_dir'] = self.save_dir_edit.text()
        params['name'] = self.name_edit.text()
        params['split'] = self.split_combo.currentText()
        params['imgsz'] = int(self.imgsz_spinbox.value())
        params['batch'] = int(self.batch_spinbox.value())
        params['conf'] = float(self.conf_spinbox.value())
        params['num_vis_samples'] = int(self.num_vis_samples_spinbox.value())
        params['device'] = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        now = datetime.datetime.now()
        now = now.strftime("%Y-%m-%d_%H-%M-%S")
        params['name'] = params['name'] if params['name'] else now
        save_dir = params['save_dir']
        save_dir = Path(save_dir) / params['name']
        params['save_dir'] = save_dir
        return params

    def evaluate_model(self):
        self.params = self.get_evaluation_parameters()
        try:
            self.model = SemanticModel(self.params['model'])
            self.worker = EvaluateModelWorker(self.model, self.params)
            self.worker.evaluation_started.connect(self.on_evaluation_started)
            self.worker.evaluation_completed.connect(self.on_evaluation_completed)
            self.worker.evaluation_error.connect(self.on_evaluation_error)
            self.worker.start()
            
            del self.model
            gc.collect()
            empty_cache()
            
        except Exception as e:
            error_message = f"An error occurred when evaluating model: {e}"
            QMessageBox.critical(self, "Error", error_message)
            print(error_message)

    def on_evaluation_started(self):
        message = "Model evaluation has commenced.\nMonitor the console for real-time progress."
        QMessageBox.information(self, "Model Evaluation Status", message)

    def on_evaluation_completed(self):
        message = "Model evaluation has successfully been completed."
        QMessageBox.information(self, "Model Evaluation Status", message)

    def on_evaluation_error(self, error_message):
        QMessageBox.critical(self, "Error", error_message)
        print(error_message)