# -*- coding: utf-8 -*-
"""
/***************************************************************************
 SELCAAnalysisDialog
                                 A QGIS plugin
 This plugin analyzes the correlation between land cover changes and dependent variables such as population and GDRP, and provides probabilistic estimating of future land cover patterns.
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                             -------------------
        begin                : 2025-02-27
        git sha              : $Format:%H$
        copyright            : (C) 2025 by Author
        email                : email@gmail.com
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

import os
from qgis.PyQt import uic
from qgis.PyQt import QtWidgets
from PyQt5.QtCore import QUrl
from PyQt5.QtGui import QDesktopServices
from qgis.gui import QgsMapLayerComboBox, QgsFieldComboBox
from qgis.PyQt.QtGui import QPixmap
from qgis.PyQt.QtWidgets import QLabel, QLineEdit, QFileDialog, QComboBox, QTabWidget, QWidget, QVBoxLayout, QGraphicsScene, QGraphicsPixmapItem, QGridLayout, QGraphicsView, QTableView, QHeaderView
from qgis.core import QgsRasterLayer, QgsVectorLayer, QgsMapLayerProxyModel, QgsApplication, QgsTask
from .Processing.crosstab import (update_graphics_luas, update_table_transition, update_graphics_transition, 
                                    compute_area_statistics, compute_transition_matrix, start_change_map_task, TableModel)
from .Processing.assignClass import get_class_aliases, get_unique_classes
from .Processing.processDependent import update_dependent_tables, export_dependent_data
from .Processing.Regression import run_regression, export_regression_data
from .Processing.crosstab import generate_descriptive_text
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import pandas as pd
from qgis.gui import QgsMapCanvas, QgsMapToolPan, QgsMapToolZoom
from qgis.core import QgsProject, QgsVectorLayer, QgsFeature, QgsGeometry, QgsPointXY, QgsField, QgsFields
from PyQt5.QtCore import QVariant
from qgis.PyQt.QtCore import Qt
from qgis.core import QgsProject, QgsLayerTreeLayer
from .Processing.forecast import fetch_forecast_data, perform_forecasting

# Load the .ui file
FORM_CLASS, _ = uic.loadUiType(os.path.join(
    os.path.dirname(__file__), 'selca_base.ui'))

class LandCoverAnalysisTask(QgsTask):
    """Background task for computing land cover statistics and transition matrix."""
    def __init__(self, initial_path, final_path, additional_paths, dialog):
        super().__init__('Land Cover Analysis', QgsTask.CanCancel)
        self.initial_path = initial_path
        self.final_path = final_path
        self.additional_paths = additional_paths
        self.dialog = dialog
        self.initial_stats = None
        self.final_stats = None
        self.additional_stats = []
        self.transition_matrix = None
        self.exception = None
        

    def run(self):
        """Perform computations in the background."""
        try:
            self.setProgress(0)
            # Compute area statistics for initial raster
            self.initial_stats = compute_area_statistics(self.initial_path)
            if self.isCanceled():
                return False
            self.setProgress(20)

            # Compute area statistics for additional rasters
            for path in self.additional_paths:
                stats = compute_area_statistics(path)
                self.additional_stats.append(stats)
                if self.isCanceled():
                    return False
            self.setProgress(40)

            # Compute area statistics for final raster
            self.final_stats = compute_area_statistics(self.final_path)
            if self.isCanceled():
                return False
            self.setProgress(50)

            # Compute transition matrix
            self.transition_matrix = compute_transition_matrix(self.initial_path, self.final_path)
            if self.isCanceled():
                return False
            self.setProgress(80)

            return True
        except Exception as e:
            self.exception = e
            return False

    def finished(self, result):
        """Handle task completion in the main thread."""
        if result:
            #update_table_luas(self.dialog, self.initial_stats, self.final_stats, self.additional_stats)
            update_graphics_luas(self.dialog, self.initial_stats, self.final_stats, self.additional_stats)
            update_table_transition(self.dialog, self.transition_matrix)
            update_graphics_transition(self.dialog, self.transition_matrix)
            self.dialog.progressBar.setValue(100)
            self.dialog.logWidget.append("Analysis completed.")
        else:
            self.dialog.logWidget.append(f"Error: {str(self.exception) if self.exception else 'Task was canceled.'}")
            self.dialog.progressBar.setValue(0)
        # Ensure temporary files are deleted
        if hasattr(self, 'temp_files'):
            for temp_file in self.temp_files:
                try:
                    os.remove(temp_file)
                except OSError:
                    pass
    
    

class SelcaAnalysisDialog(QtWidgets.QDialog, FORM_CLASS):
    def __init__(self, parent=None):
        """Constructor."""
        super(SelcaAnalysisDialog, self).__init__(parent)
        self.setupUi(self)
        self.temp_files = []
        self.landCoverCount = 2
        self.dependentVariableCount = 2
        self.landCoverComboBoxes = []
        self.aliasLineEdits = []
        self.dependentVariableComboBoxes = []
        self.observedField = []
        self.classAliases = {}
        self.master_data = pd.DataFrame()  
        self.graphicsSum = None
        self.graphicsCenterSum = None
        self.tableSum = None # Initialize master_data as an empty DataFrame
        self.change_task = None

        # Set stretch factors for landCoverLayout and aliasLayout
        rowLayout = self.findChild(QtWidgets.QHBoxLayout, "rowLayout")
        if rowLayout:
            rowLayout.setStretch(0, 5)  # Set stretch for landCoverLayout
            rowLayout.setStretch(1, 1)  # Set stretch for aliasLayout

        self.addLandCoverButton.clicked.connect(self.add_land_cover)
        self.pushUpdateAtas.clicked.connect(self.start_analysis)
        self.pushhelp.clicked.connect(self.open_help)
        self.pushFetchUniqueValues.clicked.connect(self.fetch_unique_values)
        self.pushAssignClass.clicked.connect(self.assign_class)
        self.layerDependent.layerChanged.connect(self.update_dependent_variable_fields)
        self.pushButtonDependent.clicked.connect(self.create_dependent_tabs)
        self.pushButtonUpdate.clicked.connect(lambda: update_dependent_tables(self))
        self.pushRegresi.clicked.connect(lambda: run_regression(self))
        self.ExportLandCoverInsight.clicked.connect(self.Export_LandCover_Insight)
        self.ExportDataRegresi.clicked.connect(lambda: export_regression_data(self))
        self.ExportDedendent.clicked.connect(lambda: export_dependent_data(self))
        self.CreateChangeMap.clicked.connect(lambda: start_change_map_task(self))
        self.pushFetch.clicked.connect(lambda: fetch_forecast_data(self))
        self.pushForecast.clicked.connect(lambda: perform_forecasting(self))

        self.plotWidget = self.findChild(QtWidgets.QWidget, "plotWidget")
        if self.plotWidget is not None:
            self.canvas = FigureCanvas(Figure())
            self.plotWidget.setLayout(QtWidgets.QVBoxLayout())
            self.plotWidget.layout().addWidget(self.canvas)
            self.ax = self.canvas.figure.add_subplot(111)

        # Set filters for QgsMapLayerComboBox
        self.initialLC.setFilters(QgsMapLayerProxyModel.RasterLayer)
        self.finalLC.setFilters(QgsMapLayerProxyModel.RasterLayer)
        self.layerDependent.setFilters(QgsMapLayerProxyModel.PolygonLayer)
        self.task = None  # Add this line to initialize the task attribute
    
    def open_help(self):
        help_path = os.path.join(os.path.dirname(__file__), 'selcahelpEN.html')
        if os.path.exists(help_path):
            QDesktopServices.openUrl(QUrl.fromLocalFile(help_path))
        else:
            self.logWidget.append("Help file not found.")
        

    def add_land_cover(self):
        """Add a new land cover layer input."""
        if self.landCoverCount < 10:
            comboBox = QgsMapLayerComboBox(self)
            alias = QLineEdit(self)

            # Set size policy to allow horizontal stretching
            comboBox.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed)
            alias.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed)

            comboBox.setFilters(QgsMapLayerProxyModel.RasterLayer)
            comboBox.setObjectName(f"LC_{self.landCoverCount + 2}")
            alias.setObjectName(f"alias_{self.landCoverCount + 2}")
            self.landCoverLayout.addWidget(comboBox)
            self.aliasLayout.addWidget(alias)
            self.landCoverComboBoxes.append(comboBox)
            self.aliasLineEdits.append(alias)
            self.landCoverCount += 1
            # Create the remove button
            removeButton = QtWidgets.QPushButton("✖", self)
            removeButton.setFixedSize(20, 20)
            removeButton.setObjectName(f"remove_{self.landCoverCount + 2}")
            # Create a horizontal layout to group the widgets
            rowLayout = QtWidgets.QHBoxLayout()
            rowLayout.addWidget(comboBox)
            rowLayout.addWidget(alias)
            rowLayout.addWidget(removeButton)
            self.landCoverLayout.addLayout(rowLayout)

            removeButton.clicked.connect(lambda: self.remove_land_cover(rowLayout, comboBox, alias, comboBox2, removeButton))
            rowLayout.addWidget(removeButton)

            comboBox2 = QgsFieldComboBox(self)
            comboBox2.setLayer(self.layerDependent.currentLayer())
            comboBox2.setObjectName(f"dependent_{self.dependentVariableCount + 2}")
            self.dependentVariableLayout.addWidget(comboBox2)
            self.dependentVariableComboBoxes.append(comboBox2)
            self.dependentVariableCount += 1

    def remove_land_cover(self, rowLayout, comboBox, alias, ComboBox2, removeButton):
        """Remove a land cover layer input and its associated widgets."""
        # Remove the widgets from the layout and delete them
        for widget in [comboBox, alias, ComboBox2, removeButton]:
            rowLayout.removeWidget(widget)
            widget.deleteLater()

        # Remove the layout itself
        self.landCoverLayout.removeItem(rowLayout)
        rowLayout.deleteLater()

        # Update the internal lists
        self.landCoverComboBoxes.remove(comboBox)
        self.aliasLineEdits.remove(alias)
        self.dependentVariableComboBoxes.remove(ComboBox2)
        self.landCoverCount -= 1

    def fetch_unique_values(self):
        """Fetch unique classes from the initial land cover layer."""
        self.progressBar.setValue(0)
        self.logWidget.clear()
        self.logWidget.append("Fetching unique values...")
        self.classAliases = {}
        initialLC = self.initialLC.currentLayer()
        finalLC = self.finalLC.currentLayer()  # Add this line to get final layer

        # Clear existing content in classLayout
        while self.classLayout.count():
            child = self.classLayout.takeAt(0)
            if child.widget():
                child.widget().deleteLater()
        
        if isinstance(initialLC, QgsRasterLayer) and isinstance(finalLC, QgsRasterLayer):
            # Get unique classes from both initial and final layers
            initial_classes = set(get_unique_classes(initialLC.source()))
            final_classes = set(get_unique_classes(finalLC.source()))
            # Combine unique classes from both layers
            unique_classes = sorted(initial_classes.union(final_classes))
            
            for cls in unique_classes:
                label = QLabel(f"Class {cls}:")
                line_edit = QLineEdit()
                line_edit.setObjectName(f"class_{cls}")
                self.classLayout.addWidget(label)
                self.classLayout.addWidget(line_edit)
                
        self.progressBar.setValue(100)
        self.logWidget.append("Unique values fetched.")

    def assign_class(self):
        """Assign aliases to land cover classes."""
        self.progressBar.setValue(0)
        self.logWidget.clear()
        self.logWidget.append("Assigning class aliases...")
        for i in range(self.classLayout.count()):
            widget = self.classLayout.itemAt(i).widget()
            if isinstance(widget, QLineEdit):
                class_value = int(widget.objectName().split('_')[1])
                self.classAliases[class_value] = widget.text()
        self.progressBar.setValue(100)
        self.logWidget.append("Class aliases assigned.")

    def start_analysis(self):
        """Initiate the background task for land cover analysis."""
        initialLC = self.initialLC.currentLayer()
        finalLC = self.finalLC.currentLayer()
        additionalLCs = [combo.currentLayer() for combo in self.landCoverComboBoxes]

        if not isinstance(initialLC, QgsRasterLayer) or not isinstance(finalLC, QgsRasterLayer):
            self.logWidget.append("Error: Initial and final layers must be valid raster layers.")
            return

        initial_path = initialLC.source()
        final_path = finalLC.source()
        additional_paths = [combo.source() for combo in additionalLCs if isinstance(combo, QgsRasterLayer)]

        self.task = LandCoverAnalysisTask(initial_path, final_path, additional_paths, self)  # Assign the task to self.task
        self.task.progressChanged.connect(self.update_progress)
        self.task.taskCompleted.connect(lambda: self.on_task_completed(self.task))
        self.task.taskTerminated.connect(self.on_task_failed)

        QgsApplication.taskManager().addTask(self.task)
        self.logWidget.append("Starting analysis...")

    def update_progress(self):
        """Update the progress bar during task execution."""
        task = self.sender()
        self.progressBar.setValue(int(task.progress()))

    def on_task_completed(self, task):
        """Handle task completion."""
        if task.exception is None:
            self.logWidget.append("Analysis completed successfully.")
            # Generate descriptive text and populate textDeskriptif
            descriptive_text = generate_descriptive_text(
                task.initial_stats, 
                task.final_stats, 
                task.transition_matrix, 
                self.classAliases  # Pass class aliases
            )
            self.textDeskriptif.setPlainText(descriptive_text)
        else:
            self.logWidget.append(f"Error: {str(task.exception)}")
            self.progressBar.setValue(0)

    def on_task_failed(self):
        """Handle task failure or cancellation."""
        self.logWidget.append("Task failed or was canceled.")
        self.progressBar.setValue(0)

    def update_dependent_variable_fields(self):
        """Update dependent variable fields when the layer changes."""
        for comboBox in self.dependentVariableComboBoxes:
            comboBox.setLayer(self.layerDependent.currentLayer())


    def create_dependent_tabs(self):
        """Create tabs for dependent variable analysis."""
        self.progressBar.setValue(0)
        self.logWidget.clear()
        self.logWidget.append("Creating dependent tabs...")

        # Remove all existing tabs
        while self.tabWidgetAbangku.count() > 0:
            self.tabWidgetAbangku.removeTab(0)
        
        dependent_layer = self.layerDependent.currentLayer()
        if isinstance(dependent_layer, QgsVectorLayer):
            unique_values = dependent_layer.uniqueValues(dependent_layer.fields().indexOf(self.identifierField.currentText()))
            tab = QWidget()
            tab.setObjectName(f"tabSum")
            layout = QGridLayout(tab)
                
            label_14 = QLabel("Land Cover vs Dependent Table")
            layout.addWidget(label_14, 0, 0)
                
            graphicSum = QGraphicsView()
            graphicSum.setObjectName("graphicsSum")  # Ensure the object name is set
            layout.addWidget(graphicSum, 2, 0, 1, 2)
            self.graphicsSum = graphicSum  # Store reference  
                
            label_15 = QLabel("Land Cover vs Dependent Graph")
            layout.addWidget(label_15, 3, 0)
                
            graphicsCenter = QgsMapCanvas() 
            layout.addWidget(graphicsCenter, 1, 1)
            graphicsCenter.setObjectName(f"graphicsCenterSum")
            self.graphicsCenterSum = graphicsCenter   

            tableSum = QTableView()
            layout.addWidget(tableSum, 1, 0)
            tableSum.setObjectName(f"tableSum")  
            self.tableSum = tableSum             
            
            label_17 = QLabel("Center Shift")
            layout.addWidget(label_17, 0, 1)

            # Set symmetrical row and column stretch factors
            layout.setRowStretch(0, 0)
            layout.setRowStretch(1, 3)
            layout.setRowStretch(2, 3)
            layout.setRowStretch(3, 0)
            layout.setColumnStretch(0, 1)
            layout.setColumnStretch(1, 1)
            
            # Set equal size policy for widgets
            graphicsCenter.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
            tableSum.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
            graphicSum.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
                
            self.tabWidgetAbangku.addTab(tab, str("SUM"))
            for value in unique_values:
                tab = QWidget()
                tab.setObjectName(f"tab_{value}")
                layout = QGridLayout(tab)
                # Set layout properties for equal spacing
                layout.setSpacing(10)
                layout.setContentsMargins(10, 10, 10, 10)
                
                label_14 = QLabel("Land Cover vs Dependent Table")
                layout.addWidget(label_14, 0, 0)
                
                graphicsPerkota = QGraphicsView()
                graphicsPerkota.setObjectName(f"graphics_{value}")
                # Set fixed aspect ratio of 2:1
                graphicsPerkota.setMinimumHeight(200)  # Base height
                graphicsPerkota.setMinimumWidth(400)   # 2x the height
                # Force the aspect ratio to be maintained
                sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
                sizePolicy.setHeightForWidth(True)
                graphicsPerkota.setSizePolicy(sizePolicy)
                graphicsPerkota.heightForWidth = lambda w: w // 2  # Ensure height is half the width
                layout.addWidget(graphicsPerkota, 2, 0, 1, 2)
                
                label_15 = QLabel("Land Cover vs Dependent Graph")
                layout.addWidget(label_15, 3, 0)
                
                graphicsCenter = QgsMapCanvas()
                layout.addWidget(graphicsCenter, 1, 1)
                graphicsCenter.setObjectName(f"graphicsCenter_{value}")
                
                tablePerkota = QTableView()
                layout.addWidget(tablePerkota, 1, 0)
                tablePerkota.setObjectName(f"table_{value}")
            
                
                label_17 = QLabel("Center Shift")
                layout.addWidget(label_17, 0, 1)
                
                # Set symmetrical row and column stretch factors
                layout.setRowStretch(0, 0)
                layout.setRowStretch(1, 3)
                layout.setRowStretch(2, 3)
                layout.setRowStretch(3, 0)
                layout.setColumnStretch(0, 1)
                layout.setColumnStretch(1, 1)
                
                # Set equal size policy for widgets
                graphicsCenter.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
                tablePerkota.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
                graphicsPerkota.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
                
                self.tabWidgetAbangku.addTab(tab, str(value))

                """Populate observedClass with unique classes from the initial land cover layer."""
        self.progressBar.setValue(0)
        self.logWidget.clear()
        self.logWidget.append("Fetching unique classes for observedClass...")
        initialLC = self.initialLC.currentLayer()
        if isinstance(initialLC, QgsRasterLayer):
            unique_classes = get_unique_classes(initialLC.source())
            if hasattr(self, 'observedClass'):
                self.observedClass.clear()
                self.observedClass.addItems([str(cls) for cls in sorted(unique_classes)])
                self.logWidget.append("Unique classes populated in observedClass.")
            else:
                self.logWidget.append("Error: observedClass not found.")
        else:
            self.logWidget.append("Error: Initial land cover layer is not a valid raster layer.")
    
        self.progressBar.setValue(100)
        self.logWidget.append("Dependent tabs created.")

    def Export_LandCover_Insight(self):
        """Export PNG images and CSV data from tabCrosstab to a user-selected folder."""
        self.logWidget.append("Exporting Land Cover Insight data...")
        self.progressBar.setValue(0)

        # Open folder selection dialog
        folder = QFileDialog.getExistingDirectory(self, "Select Folder to Save Land Cover Insight Data")
        if not folder:
            self.logWidget.append("Export canceled: No folder selected.")
            self.progressBar.setValue(0)
            return

        try:
            # Ensure graphics scenes exist
            if self.graphicsLuas.scene() is None or self.graphicsTransition.scene() is None:
                self.logWidget.append("Error: Graphics not yet generated.")
                self.progressBar.setValue(0)
                return

            # Save PNG from graphicsLuas
            luas_item = self.graphicsLuas.scene().items()[0]  # Assuming one item (the pixmap)
            if isinstance(luas_item, QGraphicsPixmapItem):
                luas_pixmap = luas_item.pixmap()
                luas_path = os.path.join(folder, "Area_Statistics.png")
                luas_pixmap.save(luas_path, "PNG")
                self.logWidget.append(f"Saved Area Statistics PNG to {luas_path}")
            self.progressBar.setValue(25)

            # Save PNG from graphicsTransition
            transition_item = self.graphicsTransition.scene().items()[0]  # Assuming one item
            if isinstance(transition_item, QGraphicsPixmapItem):
                transition_pixmap = transition_item.pixmap()
                transition_path = os.path.join(folder, "Transition_Heatmap.png")
                transition_pixmap.save(transition_path, "PNG")
                self.logWidget.append(f"Saved Transition Heatmap PNG to {transition_path}")
            self.progressBar.setValue(50)

            # Export CSV for Area Statistics
            initial_stats = self.task.initial_stats if self.task else {}
            final_stats = self.task.final_stats if self.task else {}
            additional_stats = self.task.additional_stats if self.task else []
            classes = sorted(set(initial_stats.keys()).union(final_stats.keys()))
            data = {
                "Class": [self.classAliases.get(cls, str(cls)) for cls in classes],
                self.initialAlias.text() if self.initialAlias.text() else "Initial": [initial_stats.get(cls, 0) for cls in classes],
            }
            for i, stats in enumerate(additional_stats):
                alias = self.aliasLineEdits[i].text() if i < len(self.aliasLineEdits) else f"LC_{i + 2}"
                data[alias] = [stats.get(cls, 0) for cls in classes]
            data[self.finalAlias.text() if self.finalAlias.text() else "Final"] = [final_stats.get(cls, 0) for cls in classes]
            
            df_area = pd.DataFrame(data)
            area_csv_path = os.path.join(folder, "Area_Statistics.csv")
            df_area.to_csv(area_csv_path, index=False)
            self.logWidget.append(f"Saved Area Statistics CSV to {area_csv_path}")
            self.progressBar.setValue(75)

            # Export CSV for Transition Matrix
            transition_matrix = self.task.transition_matrix if self.task else {}
            classes = sorted(transition_matrix.keys())
            matrix_data = [[transition_matrix[row_cls].get(col_cls, 0) for col_cls in classes] for row_cls in classes]
            df_transition = pd.DataFrame(matrix_data, index=[self.classAliases.get(cls, str(cls)) for cls in classes], 
                                        columns=[self.classAliases.get(cls, str(cls)) for cls in classes])
            transition_csv_path = os.path.join(folder, "Transition_Matrix.csv")
            df_transition.to_csv(transition_csv_path)
            self.logWidget.append(f"Saved Transition Matrix CSV to {transition_csv_path}")
            self.progressBar.setValue(100)

            self.logWidget.append("Export completed successfully.")
        except Exception as e:
            self.logWidget.append(f"Error during export: {str(e)}")
            self.progressBar.setValue(0)

    def on_task_failed(self, error_message):
        """Handle task failure or cancellation."""
        self.logWidget.append(f"Task failed: {error_message}")
        self.progressBar.setValue(0)

    def update_graphics_transition_2(self, percentage_matrix):
        """Update the graphicsTransition_2 table with the percentage matrix."""
        classes = sorted(percentage_matrix.keys())
        data = [[f"{percentage_matrix[row_cls][col_cls] * 100:.2f}%" for col_cls in classes] for row_cls in classes]

        model = TableModel(data, classes)
        self.graphicsTransition_2.setModel(model)
        self.graphicsTransition_2.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.graphicsTransition_2.verticalHeader().setSectionResizeMode(QHeaderView.Stretch)
        for i, cls in enumerate(classes):
            model.setHeaderData(i, Qt.Horizontal, str(cls))
            model.setHeaderData(i, Qt.Vertical, str(cls))