"""
ModelEvaluator - Comprehensive model evaluation with plotting and logging.
Moved from KerasModelFactory.plot_metrics() for better separation of concerns.
"""

from meow_ml.src.model_factories.ModelFactory import ModelFactory
from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.layers import Input, Dense, Flatten, Dropout, concatenate, Conv1D, MaxPooling1D, GlobalMaxPooling1D, BatchNormalization
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.optimizers import SGD, Adam
from tensorflow.keras.losses import Huber, MeanSquaredError, MeanAbsoluteError, SparseCategoricalCrossentropy
from tensorflow.keras.utils import set_random_seed
import tensorflow as tf
# from tensorflow.keras.nlp.layers import TransformerEncoder
# from keras_nlp.layers import TransformerEncoder
import mlflow
import matplotlib as mpl
import rasterio
import matplotlib.pyplot as plt
from matplotlib.image import AxesImage
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
import numpy as np
import math
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay, classification_report, r2_score, roc_curve, roc_auc_score, precision_recall_curve, mean_absolute_error, mean_squared_error
from scipy.stats import gaussian_kde
import os
import io
import tempfile
import shutil
import joblib
import warnings


class ModelEvaluator:
    """Handles comprehensive model evaluation with plotting and logging."""
    
    def evaluate_model(self, model, train_generator, validation_generator, test_generator, 
                      history, test_type='reg', scale_method=None):
        """
        Comprehensive model evaluation - moved from KerasModelFactory.plot_metrics()
        
        Parameters
        ----------
        model : keras.Model
            Trained model
        train_generator : DataGenerator
            Training data generator
        validation_generator : DataGenerator
            Validation data generator
        test_generator : DataGenerator
            Test data generator
        history : keras.History
            Training history
        test_type : str
            'reg' for regression, 'cat' for classification
        scale_method : str, optional
            Scaling method used
            
        Returns
        -------
        dict
            Dictionary of stored plots organized by dataset
        """
        print('** Starting comprehensive model evaluation **')
        
        # Initialize plot storage
        stored_plots = {
            'train': {},
            'test': {},
            'validation': {}
        }
        
        # Configure matplotlib settings
        self._configure_plot_settings()
        
        # Log training history
        self.log_loss_plot(history, log_mlflow=True)
        
        # Generate plots based on test type
        if test_type == "cat":
            self._evaluate_classification_model(model, test_generator, history)
        elif test_type == "reg":
            self._evaluate_regression_model(model, train_generator, validation_generator, 
                                          test_generator, stored_plots, scale_method)
        
        # Generate common plots for all datasets
        self._generate_common_plots(model, train_generator, validation_generator, 
                                  test_generator, stored_plots)
        
        # Generate grid plots
        self._generate_grid_plots(model, train_generator, validation_generator, 
                                test_generator, stored_plots)
        
        # Log raw predictions and truth arrays
        self._log_raw_arrays(model, train_generator, validation_generator, test_generator)
        
        # Generate CHC-specific evaluations
        self._evaluate_chc_specific(model, train_generator, validation_generator, test_generator)
        
        print('** Model evaluation complete **')
        return stored_plots
    
    def _configure_plot_settings(self):
        """Configure matplotlib settings for consistent plotting"""
        # You can uncomment and customize these settings as needed
        # plt.rcParams.update({
        #     'font.size': 18,
        #     'axes.titlesize': 24,
        #     'axes.labelsize': 20,
        #     'xtick.labelsize': 20,
        #     'ytick.labelsize': 20,
        #     'legend.fontsize': 20,
        #     'figure.titlesize': 36,
        #     'legend.title_fontsize': 20,
        # })
        pass
    
    def _evaluate_classification_model(self, model, test_generator, history):
        """Handle classification-specific evaluation"""
        self.log_confusion_matrix(model, test_generator, dataset="test", log_mlflow=True)
        self.log_classification_report(model, test_generator, dataset="test", log_mlflow=True)
        self.log_roc_curve(model, test_generator, dataset="test", log_mlflow=True)
        self.log_precision_recall_curve(model, test_generator, dataset="test", log_mlflow=True)
        self.log_accuracy_plot(history, log_mlflow=True)
    
    def _evaluate_regression_model(self, model, train_gen, val_gen, test_gen, stored_plots, scale_method):
        """Handle regression-specific evaluation"""
        # Random sample plots
        self.log_actual_vs_predicted_random_sample(model, test_gen, dataset="test", 
                                                  scale_method=scale_method, log_mlflow=True)
        self.log_actual_vs_predicted_random_sample(model, train_gen, dataset="train", 
                                                  scale_method=scale_method, log_mlflow=True)
        
        # Scatter plots and error intensity plots
        datasets = [(train_gen, "train"), (test_gen, "test"), (val_gen, "validation")]
        for generator, dataset_name in datasets:
            stored_plots[dataset_name]['scatter'] = self.log_actual_vs_predicted_scatter(
                model, generator, dataset=dataset_name, log_mlflow=True)
            stored_plots[dataset_name]['error_intensity'] = self.log_error_intensity_plot(
                model, generator, dataset=dataset_name, log_mlflow=True)
        
        # R2 scores
        self.log_r2_score(model, test_gen, dataset="test", log_mlflow=True)
        self.log_r2_score(model, train_gen, dataset="train", log_mlflow=True)
    
    def _generate_common_plots(self, model, train_gen, val_gen, test_gen, stored_plots):
        """Generate plots common to all model types"""
        datasets = [(train_gen, "train"), (test_gen, "test"), (val_gen, "validation")]
        
        for generator, dataset_name in datasets:
            # Masks, truth, RGB, predictions
            stored_plots[dataset_name]['mask'] = self.log_masks(generator, dataset=dataset_name, log_mlflow=True)
            stored_plots[dataset_name]['y_true'] = self.log_truth(generator, dataset=dataset_name, log_mlflow=True)
            stored_plots[dataset_name]['rgb'] = self.log_original_tiles(generator, dataset=dataset_name, log_mlflow=True)
            stored_plots[dataset_name]['y_pred'] = self.log_pred(model, generator, dataset=dataset_name, log_mlflow=True)
            
            # Log scalers
            self.log_scalers(generator, dataset=dataset_name)
    
    def _generate_grid_plots(self, model, train_gen, val_gen, test_gen, stored_plots):
        """Generate grid subplot combinations"""
        print('calling log all plots to grid')
        datasets = [(train_gen, "train"), (test_gen, "test"), (val_gen, "validation")]
        
        for generator, dataset_name in datasets:
            self.log_all_plots_to_grid(model, generator, dataset=dataset_name, 
                                      plots=stored_plots[dataset_name])
    
    def _log_raw_arrays(self, model, train_gen, val_gen, test_gen):
        """Log raw prediction and truth arrays"""
        datasets = [(train_gen, "train"), (test_gen, "test"), (val_gen, "validation")]
        
        for generator, dataset_name in datasets:
            # Fix the typo from your original code
            actual_dataset_name = dataset_name if dataset_name != "validation" else "validation"
            self.log_y_true_y_pred(model, generator, dataset=actual_dataset_name, log_to_mlflow=True)
    
    def _evaluate_chc_specific(self, model, train_gen, val_gen, test_gen):
        """Handle CHC-specific evaluation (domain-specific)"""
        datasets = [(train_gen, "train"), (test_gen, "test"), (val_gen, "validation")]
        
        for generator, dataset_name in datasets:
            self.log_chc_histogram(model, generator, dataset=dataset_name, log_mlflow=True)
            self.log_chc_statistics_table(model, generator, dataset=dataset_name, log_mlflow=True)
            self.log_chc_scatterplot(model, generator, dataset=dataset_name, log_mlflow=True)
            self.log_chc_metrics_table(model, generator, dataset=dataset_name, log_mlflow=True)
    
    def log_accuracy_plot(self, history, log_mlflow=False):
        # logging and plotting training & validation accuracy values
        fig, ax = plt.subplots()
        ax.plot(history.history['accuracy'])
        ax.plot(history.history['val_accuracy'])
        ax.set_title('Model accuracy')
        ax.set_ylabel('Accuracy')
        ax.set_xlabel('Epoch')
        ax.legend(['Train', 'Validation'], loc='upper left')
        if log_mlflow:
            mlflow.log_figure(fig, "training_validation_accuracy.png")
        
        return [fig]

    def log_loss_plot(self, history, log_mlflow=False):
        # logging and plotting training & validation loss values
        fig, ax = plt.subplots()
        ax.plot(history.history['loss'])
        ax.plot(history.history['val_loss'])
        ax.set_title('Model loss')
        ax.set_ylabel('Loss')
        ax.set_xlabel('Epoch')
        ax.legend(['Train', 'Validation'], loc='upper left')

        if log_mlflow:
            mlflow.log_figure(fig, "training_validation_loss.png")

        return [fig]

    def analyze_data_before_training(self, train_generator, validation_generator, test_generator):
        """
        Comprehensive data analysis before training begins.
        """
        print("** Starting pre-training data analysis **")
        
        # Log basic statistics
        self.log_min_max_mean_std(train_generator, "train")
        
        # Log truth proportions for categorical data
        if (train_generator.test_type == "cat" and 
            validation_generator.test_type == "cat" and 
            test_generator.test_type == "cat"):
            self.log_truth_proportions(test_generator, "test", log_mlflow=True)
        
        print("** Pre-training data analysis complete **")

    def log_min_max_mean_std(self, data_generator, dataset="test"):
        stuff = data_generator.get_mean_std_min_max()
        # [mlflow.log_metrics(x) for x in stuff]
        [mlflow.log_params(x) for x in stuff]

    # logs proportion of truth values in each bin to mlflow
    def log_truth_proportions(self, data_generator, dataset="test", log_mlflow=False):
        bins, counts = np.unique(data_generator.y, return_counts=True)
        proportions = counts / len(data_generator.y)

        for bin, proportion in zip(bins, proportions):
            # print(f"Bin {bin}: {proportion*100:.2f}%")
            mlflow.log_param(f"{data_generator.split_mode} / bin {bin.astype(int)}",
                             f"{proportion*100:.2f}%")

        # logs histogram of binned truth proportions
        fig, ax = plt.subplots()
        ax.bar(bins, proportions, align='center')
        ax.set_title('Distribution of bins in truth')
        ax.set_xlabel('Bins')
        ax.set_ylabel('Proportions')
        if log_mlflow:
            mlflow.log_figure(fig, f"bin_proportions_{dataset}.png")
        
        return [fig]

    # TODO: uncomment if scaling y is useless
    # only for regression tasks, logging plot actual vs. predicted values of random subsample
    # def log_actual_vs_predicted_random_sample(self, model, data_generator, dataset="test", scale_method=None, log_mlflow=False):
    #     print(f'** Plotting Actual vs. Predicted Random Sample for {dataset.capitalize()} Dataset **')
    #     y_true = data_generator.y.flatten()
        
    #     # y_true = y_true / 1000 # unscaling back to original y

    #     num_samples = y_true.shape[0]
    #     y_pred = model.predict(data_generator)
    #     # loss, accuracy = model.evaluate(data_generator)
    #     random_sample_indices = np.random.choice(num_samples, size=min(num_samples, 50), replace=False)
    #     sample_y_true = y_true[random_sample_indices]
    #     sample_y_pred = y_pred[random_sample_indices]

    #     # calculate average values
    #     avg_values = np.average(sample_y_pred)

    #     # Plot actual, predicted, and average values
    #     fig, axs = plt.subplots(figsize=(10, 6))
    #     axs.plot(sample_y_true, color='orange', label='Actual Values')
    #     axs.plot(sample_y_pred, color='lightblue', label='Predicted Values')
    #     axs.axhline(y=avg_values, color='green', label='Average Predicted Values')
    #     axs.set_xlabel('Random Sample')
    #     axs.set_ylabel('Canopy height (meters)')
    #     axs.set_title('Actual, Predicted and Average Values')

    #     # adding r2 score in plot
    #     r2 = r2_score(y_true, y_pred)
    #     score_text = f"r^2: {r2:.2f}"
    #     score_annotation = mpl.text.Annotation(score_text, xy=(0.99, 0.96), xycoords='axes fraction',
    #                                         ha='right', va='center', fontsize=10)
    #     axs.add_artist(score_annotation)

    #     axs.legend(loc='lower right')
    #     if log_mlflow:
    #         mlflow.log_figure(fig, f"actual_vs_predicted_values_{dataset}.png")

    #     return [fig]

    # TODO: uncomment if scaling y is useless
    # def log_actual_vs_predicted_scatter(self, model, data_generator, dataset="test", log_mlflow=False):
    #     print(f'** Plotting Actual vs. Predicted Scatter Plot for {dataset.capitalize()} Dataset **')
    #     y_true = data_generator.y.flatten()
    #     y_pred = model.predict(data_generator).flatten()

    #     # y_true = y_true / 10 # unscaling back to original y
    #     # y_pred = y_pred / 10 # unscaling back to original y

    #     master_mask = data_generator.master_mask
    #     truth_name = data_generator.truth

    #     # Remove NaNs from the entire dataset
    #     valid_mask = ~np.isnan(y_true) & ~np.isnan(y_pred)
    #     y_true_valid = y_true[valid_mask]
    #     y_pred_valid = y_pred[valid_mask]

    #     # Create aggregate scatter plot
    #     fig, ax = plt.subplots(figsize=(10, 10))

    #     # temporarily disabling gaussian kde while I figure out how to run it more efficiently
    #     # try:
    #     #     xy = np.vstack([y_true_valid, y_pred_valid])
    #     #     z = gaussian_kde(xy)(xy)
    #     #     sc = ax.scatter(y_true_valid, y_pred_valid, c=z, cmap='turbo', s=10, edgecolor='none', alpha=0.7)
    #     #     cbar = fig.colorbar(sc, ax=ax)
    #     #     cbar.set_label('Density')
    #     #     cbar.set_ticks([])
            
    #     #     # Store the density information in the figure
    #     #     fig.density_data = {
    #     #         'points': np.column_stack([y_true_valid, y_pred_valid]),
    #     #         'density': z
    #     #     }
    #     # except Exception as e:
    #     #     print(f"KDE failed for aggregate plot: {e}. Proceeding without density.")
    #     #     ax.scatter(y_true_valid, y_pred_valid, color='blue', s=10, alpha=0.7)
    #     ax.scatter(y_true_valid, y_pred_valid, color='blue', s=0.5, alpha=0.7)

    #     min_val = min(y_true_valid.min(), y_pred_valid.min())
    #     max_val = max(y_true_valid.max(), y_pred_valid.max())
    #     ax.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Predictions')
        
    #     # ax.set_xlim(0, 0.5)
    #     # ax.set_ylim(0, 0.5)
    #     ax.set_xlabel(f'{truth_name.upper()} Truth (m)')
    #     ax.set_ylabel(f'{truth_name.upper()} Prediction (m)')
    #     ax.set_title(f'{truth_name.upper()} Actual vs Predicted Values {dataset.capitalize()}')
    #     ax.set_aspect('equal')
    #     ax.legend()
    #     fig.tight_layout()

    #     # Handle the aggregate plot based on mode
    #     # TODO: figure out how to return only aggregate plot if needed
    #     if log_mlflow:
    #         mlflow.log_figure(fig, f"actual_vs_predicted_scatter_{dataset}.png")
    #         plt.close(fig)

    #     plots = []

    #     # Process individual tiles
    #     y_pred_expanded = np.full(master_mask.shape, np.nan)
    #     y_true_expanded = np.full(master_mask.shape, np.nan)
    #     masked_indices = np.where(~master_mask)
    #     # debugging
    #     print(y_pred.shape)
    #     y_pred_expanded[masked_indices] = y_pred.flatten()
    #     y_true_expanded[masked_indices] = y_true

    #     # Determine tile dimensions and reshape data
    #     if data_generator.upscale_params:
    #         original_tile_size = 1000
    #         new_tile_dims = original_tile_size // data_generator.upscale_params.get('window_size', 10)
    #         num_tiles = master_mask.size // (new_tile_dims * new_tile_dims)
    #         y_pred_tiles = y_pred_expanded.reshape(num_tiles, new_tile_dims, new_tile_dims)
    #         y_true_tiles = y_true_expanded.reshape(num_tiles, new_tile_dims, new_tile_dims)
    #     else:
    #         num_tiles = 1 if len(y_pred) <= 1000000 else y_pred_expanded.shape[0] // 1000000
    #         if num_tiles == 1:
    #             y_pred_tiles = [y_pred_expanded]
    #             y_true_tiles = [y_true_expanded]
    #         else:
    #             y_pred_tiles = y_pred_expanded.reshape((num_tiles, 1000, 1000))
    #             y_true_tiles = y_true_expanded.reshape((num_tiles, 1000, 1000))

    #     for i in range(num_tiles):
    #         y_true_tile = y_true_tiles[i].flatten()
    #         y_pred_tile = y_pred_tiles[i].flatten()

    #         valid_mask = ~np.isnan(y_true_tile) & ~np.isnan(y_pred_tile)
    #         y_true_tile = y_true_tile[valid_mask]
    #         y_pred_tile = y_pred_tile[valid_mask]

    #         if len(y_true_tile) < 2 or len(y_pred_tile) < 2:
    #             print(f"Skipping tile {i} due to insufficient valid data points.")
    #             continue

    #         fig, ax = plt.subplots(figsize=(10, 10))
            
    #         # temporarily disabling gaussian kde while I figure out how to run it more efficiently
    #         # try:
    #         #     xy = np.vstack([y_true_tile, y_pred_tile])
    #         #     z = gaussian_kde(xy)(xy)
    #         #     sc = ax.scatter(y_true_tile, y_pred_tile, c=z, cmap='turbo', s=10, edgecolor='none', alpha=0.7)
    #         #     cbar = fig.colorbar(sc, ax=ax)
    #         #     cbar.set_label('Density')
    #         #     cbar.set_ticks([])
                
    #         #     # Store the density information in the figure
    #         #     fig.density_data = {
    #         #         'points': np.column_stack([y_true_tile, y_pred_tile]),
    #         #         'density': z
    #         #     }
    #         # except Exception as e:
    #         #     print(f"KDE failed for tile {i}: {e}. Plotting without density.")
    #         #     ax.scatter(y_true_tile, y_pred_tile, color='blue', s=10, alpha=0.7)
    #         ax.scatter(y_true_tile, y_pred_tile, color='blue', s=10, alpha=0.7)

    #         min_val = min(y_true_tile.min(), y_pred_tile.min())
    #         max_val = max(y_true_tile.max(), y_pred_tile.max())

    #         ax.set_aspect('equal')

    #         ax.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Predictions')

    #         ax.set_xlim(min_val, max_val)
    #         ax.set_ylim(min_val, max_val)
    #         ax.set_xlabel(f'{truth_name.upper()} Truth (m)')
    #         ax.set_ylabel(f'{truth_name.upper()} Prediction (m)')

    #         tile_name = data_generator.file_paths[i].split("/")[1]
    #         ax.set_title(f'{truth_name.upper()} Truth vs Prediction - Tile {tile_name.replace("_", " ")}')

    #         if log_mlflow:
    #             mlflow.log_figure(fig, f'./actual_vs_predicted_scatter/scatter_{dataset}_tile_{tile_name}.png')
    #             plt.close(fig)
    #         plots.append(fig)

    #     return plots

    # TODO: uncomment if scaling y is useless
    # def log_r2_score(self, model, data_generator, dataset="test", log_mlflow=False):
    #     y_true = data_generator.y.flatten()
    #     y_pred = model.predict(data_generator)

    #     # y_true = y_true / 1000 # unscaling back to original y
    #     # y_pred = y_pred / 1000 # unscaling back to original y

    #     r2 = r2_score(y_true, y_pred)
    #     if log_mlflow:
    #         mlflow.log_metric(f"r2_{dataset}", f"{r2:.2f}")
    #     else:
    #         print(r2)

    def log_roc_curve(self, model, data_generator, dataset="test", log_mlflow=False):

        y_true = data_generator.y
        y_pred = model.predict(data_generator)

        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y


        fpr, tpr, thresholds = roc_curve(y_true, y_pred.flatten())

        # plt.figure()
        fig, axs = plt.subplots(figsize=(10, 6))
        lw = 2
        axs.plot(fpr, tpr, color='darkorange',
                lw=lw, label='ROC curve (area = %0.2f)' % roc_auc_score(y_true, y_pred.flatten()))
        axs.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
        axs.set_xlim([0.0, 1.0])
        axs.set_ylim([0.0, 1.05])
        axs.set_xlabel('False Positive Rate')
        axs.set_ylabel('True Positive Rate')
        axs.set_title('ROC Curve')
        axs.legend(loc="lower right")
        if log_mlflow:
            mlflow.log_figure(fig, f"roc_curve_{dataset}.png")

        return [fig]

    def log_precision_recall_curve(self, model, data_generator, dataset="test", log_mlflow=False):
        y_true = data_generator.y
        y_pred = model.predict(data_generator)

        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y

        precision, recall, thresholds = precision_recall_curve(y_true, y_pred.flatten())
        f1_score = 2 * (precision * recall) / (precision + recall)

        fig, axs = plt.subplots(figsize=(10, 6))
        lw = 2
        axs.plot(recall, precision, color='darkorange', lw=lw, label='Precision-recall curve')
        axs.plot([0, 1], [1, 0], color='navy', lw=lw, linestyle='--')
        axs.set_xlim([0.0, 1])
        axs.set_ylim([0.0, 1.05])
        axs.set_xlabel('Recall')
        axs.set_ylabel('Precision')
        axs.set_title('Precision-recall curve example')
        axs.legend(loc="lower left")

        score_text = f"f1: {f1_score.max():.2f}"
        score_annotation = mpl.text.Annotation(score_text, xy=(0.99, 0.96), xycoords='axes fraction',
                                            ha='right', va='center', fontsize=10)
        axs.add_artist(score_annotation)

        axs.legend(loc='lower right')
        if log_mlflow:
            mlflow.log_figure(fig, f"precision_recall_curve{dataset}.png")
        return [fig]

    def log_confusion_matrix(self, model, data_generator, dataset="test", log_mlflow=False):
        # TODO: set y axis title based on truth value from data generator
        # logging and plotting raw and normalized confusion matrix
        y_true = data_generator.y
        y_pred = model.predict(data_generator)

        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y

        y_pred = np.argmax(y_pred, axis=1)
        raw_confusion_matrix = confusion_matrix(y_true, y_pred)
        raw_disp = ConfusionMatrixDisplay(raw_confusion_matrix)
        raw_disp.plot()
        if log_mlflow:
            mlflow.log_figure(raw_disp.figure_, 'raw_confusion_matrix.png')
        normalized_confusion_matrix = confusion_matrix(y_true, y_pred, normalize="true")
        normalized_disp = ConfusionMatrixDisplay(normalized_confusion_matrix)
        normalized_disp.plot()
        if log_mlflow:
            mlflow.log_figure(normalized_disp.figure_, 'normalized_confusion_matrix.png')
        return raw_disp, normalized_disp

    def log_classification_report(self, model, data_generator, dataset="test", log_mlflow=False):
        y_true = data_generator.y
        y_pred = model.predict(data_generator)

        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y

        y_pred_classes = np.argmax(y_pred, axis=1)
        class_report = classification_report(y_true, y_pred_classes)
        if log_mlflow:
            mlflow.log_text(class_report, "classification_report.txt")
        else:
            print(class_report)
        return class_report

    # TODO: uncomment if scaling y is useless
    # TODO: modify for classification error intensity capability as well
    # def log_error_intensity_plot(self, model, data_generator, dataset="test", log_mlflow=False):
    #     """
    #     Plots error intensity heat maps for individual tiles, taking into account masked pixels
    #     and displaying them all at once.

    #     Args:
    #         model: The trained machine learning model.
    #         data_generator: data_generator, contains master_mask, truth, and input data
    #     """

    #     y_true = data_generator.y
    #     y_pred = model.predict(data_generator)

    #     # y_true = y_true / 1000 # unscaling back to original y
    #     # y_pred = y_pred / 1000 # unscaling back to original y

    #     master_mask = data_generator.master_mask

    #     # # unscaling y_pred and y_true to get error in meters
    #     # # debugging
    #     # print(f'self.scalers: {data_generator.scalers}')
    #     # [print(scaler) for scaler in data_generator.scalers]
    #     # y_pred = data_generator.scalers[-1].inverse_transform(y_pred)
    #     # y_true = data_generator.scalers[-1].inverse_transform(y_true)

    #     y_pred_expanded = np.full(master_mask.shape, np.nan)
    #     y_true_expanded = np.full(master_mask.shape, np.nan)

    #     masked_indices = np.where(~master_mask)[0]

    #     y_pred_expanded[masked_indices] = y_pred.flatten()
    #     y_true_expanded[masked_indices] = y_true

    #     # reshaping into individual tiles
    #     if data_generator.upscale_params:
    #         # TODO: make original image size variable in case we do patches (will likely be config in data generator)
    #         original_tile_size = 1000
    #         new_tile_dims = original_tile_size // data_generator.upscale_params.get('window_size', 10)
    #         num_tiles = master_mask.size // (new_tile_dims * new_tile_dims)
    #         y_pred = y_pred_expanded.reshape(num_tiles, new_tile_dims, new_tile_dims)
    #         y_true = y_true_expanded.reshape(num_tiles, new_tile_dims, new_tile_dims)
    #         # debugging
    #         # print(f'y_pred.shape (for upscaling): {y_pred.shape}')
    #         # print(f'y_true.shape (for upscaling): {y_true.shape}')
    #     else:
    #         num_tiles = y_pred_expanded.shape[0] // 1000000
    #         y_pred = y_pred_expanded.reshape((num_tiles, 1000, 1000))
    #         y_true = y_true_expanded.reshape((num_tiles, 1000, 1000))

    #     y_axis_label = "Absolute Error (Meters)" if data_generator.test_type == "reg" else "Absolute Error (off by class)"

    #     plots = []

    #     for i in range(num_tiles):
    #         # errors = np.abs(y_pred[i] - y_true[i])
    #         errors = (y_pred[i] - y_true[i])

    #         mask_color = "black"
    #         cmap = mpl.colormaps.get_cmap('turbo')
    #         # cmap = mpl.colormaps.get_cmap('viridis')
    #         # cmap = mpl.colormaps.get_cmap('cividis')
    #         cmap.set_bad(color=mask_color)

    #         # Set the thresholds for values above and below which you want to change the color
    #         above_threshold_clip = 5
    #         below_threshold_clip = -5

    #         # Clip the extreme values of errors
    #         errors_clipped = np.clip(errors, below_threshold_clip, above_threshold_clip)

    #         # Display error intensity heat map on the corresponding subplot with clipping
    #         fig, ax = plt.subplots(figsize=(10, 10))
    #         cax = ax.imshow(errors_clipped, cmap=cmap, origin='lower', vmin=below_threshold_clip, vmax=above_threshold_clip)
    #         colorbar = fig.colorbar(cax, ax=ax, label='Absolute Error (meters)')
    #         ax.set_xlim([0, errors.shape[1]])
    #         ax.set_ylim([0, errors.shape[0]])
    #         ax.set_xlabel('X')
    #         ax.set_ylabel('Y')
    #         tile_name = data_generator.file_paths[i].split("/")[1]
    #         ax.set_title(f'Error Intensity Map - Tile {tile_name.replace("_", " ")} - {dataset}')
            
    #         plots.append(fig)
            
    #         if log_mlflow:
    #             mlflow.log_figure(fig, f'./error_intensities/error_intensity_map_{dataset}_tile_{tile_name}.png')
    #             # close figure to prevent memory leaks
    #             plt.close(fig)

    #     return plots


    def log_masks(self, data_generator, dataset="test", log_mlflow=False):
        """
        Plots and logs individual tiles of the master mask to MLflow,
        visualizing masked areas for each tile.

        Args:
            data_generator: data_generator, contains master_mask
        """

        if data_generator.upscale_params:
            # TODO: make original image size variable in case we do patches (will likely be config in data generator)
            original_tile_size = 1000
            new_tile_dims = original_tile_size // data_generator.upscale_params.get('window_size', 10)
            master_mask = data_generator.master_mask
            num_tiles = master_mask.size // (new_tile_dims * new_tile_dims)
            master_mask = master_mask.reshape(num_tiles, new_tile_dims, new_tile_dims)
            # debugging
            # print(f'master_mask.shape (for upscaling): {master_mask.shape}')
        else:
            master_mask = data_generator.master_mask
            master_mask = master_mask.reshape(-1, 1000, 1000)
            num_tiles = master_mask.shape[0]

        plots = []

        # Loop through tiles and create separate plots for each
        for i in range(num_tiles):
            # Get the mask for the current tile
            tile_mask = master_mask[i]

            # Create a figure and plot the mask using grayscale colormap
            fig, ax = plt.subplots(figsize=(10, 10))  # Adjust figsize as needed
            ax.imshow(tile_mask, cmap='gray_r', interpolation='nearest', origin='lower')

            # Customize and adjust the tile plot
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_xlim([0, tile_mask.shape[1]])
            ax.set_ylim([0, tile_mask.shape[0]])

            # Log the figure to MLflow
            tile_name = data_generator.file_paths[i].split("/")[1]
            ax.set_title(f'Mask Tile {tile_name.replace("_", " ")} - {dataset}')
            if log_mlflow:
                mlflow.log_figure(fig, f'./masks/mask_{dataset}_tile_{tile_name}.png')
                # Close the figure to prevent memory leaks
                plt.close(fig)
            plots.append(fig)

        return plots

    # TODO: uncomment if scaling y is useless
    # def log_pred(self, model, data_generator, dataset="test", log_mlflow=False):
    #     """
    #     Plots and logs individual tiles of the prediction to MLflow,
    #     visualizing truth areas for each tile.

    #     Args:
    #     data_generator: data_generator, contains y (truth)
    #     """

    #     # re-expand y with masked values
    #     y_pred = model.predict(data_generator)
    #     y_pred = y_pred.flatten()
    #     master_mask = data_generator.master_mask
    #     y_pred_expanded = np.full(master_mask.shape, np.nan)
    #     masked_indices = np.where(~master_mask)[0]

    #     y_pred_expanded[masked_indices] = y_pred
    #     y_pred = y_pred_expanded

    #     if data_generator.upscale_params:
    #         # TODO: make original image size variable in case we do patches (will likely be config in data generator)
    #         original_tile_size = 1000
    #         new_tile_dims = original_tile_size // data_generator.upscale_params.get('window_size', 10)
    #         num_tiles = y_pred.size // (new_tile_dims * new_tile_dims)
    #         y_pred = y_pred.reshape(num_tiles, new_tile_dims, new_tile_dims)
    #     else:
    #         y_pred = y_pred.reshape(-1, 1000, 1000)
    #         num_tiles = y_pred.shape[0]

    #     mask_color = "black"
    #     # cmap = mpl.colormaps.get_cmap('viridis')
    #     cmap = mpl.colormaps.get_cmap('gray_r')
    #     cmap.set_bad(color=mask_color)

    #     plots = []

    #     # Loop through tiles and create separate plots for each
    #     for i in range(num_tiles):
    #         # Get the y_pred for the current tile
    #         y_pred_tile = y_pred[i]

    #         # debugging
    #         # print('FOR LOGGING MIN AND MAX VALUES FOR Y_PRED')
    #         # print(f'Tile {i} min value: {np.nanmin(y_pred_tile)}, max value: {np.nanmax(y_pred_tile)}')

    #         # Create a figure and plot the y_pred using grayscale colormap
    #         fig, ax = plt.subplots(figsize=(10, 10))
    #         cax = ax.imshow(y_pred_tile, cmap=cmap, interpolation='nearest', origin='lower', vmin=0, vmax=40)
    #         # TODO: cap vmax at a good one. vmin good at 0
    #         fig.colorbar(cax, ax=ax)

    #         # Customize and adjust the tile plot
    #         ax.set_xlabel('X')
    #         ax.set_ylabel('Y')
    #         ax.set_xlim([0, y_pred_tile.shape[1]])
    #         ax.set_ylim([0, y_pred_tile.shape[0]])

    #         # Remove axis ticks
    #         ax.set_xticks([])
    #         ax.set_yticks([])

    #         # Log the figure to MLflow
    #         tile_name = data_generator.file_paths[i].split("/")[1]
    #         truth_name = data_generator.truth
    #         ax.set_title(f'Predicted {truth_name.upper()} Tile {tile_name.replace("_", " ")} - {"10 m" if data_generator.upscale_params else "1 m"} - {dataset}')
    #         if log_mlflow:
    #             mlflow.log_figure(fig, f'./y_pred/pred_{truth_name.lower()}_{dataset}_tile_{tile_name}.png')
    #             # Close the figure to prevent memory leaks
    #             plt.close(fig)
    #         plots.append(fig)

    #     return plots

    # TODO: uncomment if scaling y is useless
    # def log_truth(self, data_generator, dataset="test", log_mlflow=False):
    #     """
    #     Plots and logs individual tiles of the truth to MLflow,
    #     visualizing truth areas for each tile.

    #     Args:
    #     data_generator: data_generator, contains y (truth)
    #     """

    #     # re-expand y with masked values
    #     y_true = data_generator.y

    #     # y_true = y_true / 1000 # unscaling back to original y

    #     master_mask = data_generator.master_mask
    #     y_true_expanded = np.full(master_mask.shape, np.nan)
    #     masked_indices = np.where(~master_mask)[0]

    #     y_true_expanded[masked_indices] = y_true
    #     y_true = y_true_expanded

    #     if data_generator.upscale_params:
    #         # TODO: make original image size variable in case we do patches (will likely be config in data generator)
    #         original_tile_size = 1000
    #         new_tile_dims = original_tile_size // data_generator.upscale_params.get('window_size', 10)
    #         num_tiles = y_true.size // (new_tile_dims * new_tile_dims)
    #         y_true = y_true.reshape(num_tiles, new_tile_dims, new_tile_dims)
    #     else:
    #         y_true = y_true.reshape(-1, 1000, 1000)
    #         num_tiles = y_true.shape[0]

    #     mask_color = "black"
    #     # cmap = mpl.colormaps.get_cmap('viridis')
    #     cmap = mpl.colormaps.get_cmap('gray_r')
    #     cmap.set_bad(color=mask_color)

    #     plots = []

    #     # Loop through tiles and create separate plots for each
    #     for i in range(num_tiles):
    #         # Get the y_true for the current tile
    #         y_true_tile = y_true[i]

    #         # Create a figure and plot the y_true using grayscale colormap
    #         fig, ax = plt.subplots(figsize=(10, 10))
    #         ax.imshow(y_true_tile, cmap=cmap, interpolation='nearest', origin='lower')

    #         # Customize and adjust the tile plot
    #         ax.set_xlabel('X')
    #         ax.set_ylabel('Y')
    #         ax.set_xlim([0, y_true_tile.shape[1]])
    #         ax.set_ylim([0, y_true_tile.shape[0]])

    #         # Remove axis ticks
    #         ax.set_xticks([])
    #         ax.set_yticks([])

    #         # Log the figure to MLflow
    #         tile_name = data_generator.file_paths[i].split("/")[1]
    #         truth_name = data_generator.truth
    #         ax.set_title(f'True {truth_name.upper()} Tile {tile_name.replace("_", " ")} - {"10 m" if data_generator.upscale_params else "1 m"} - {dataset}')
    #         if log_mlflow:
    #             mlflow.log_figure(fig, f'./y_true/true_{truth_name.lower()}_{dataset}_tile_{tile_name}.png')
    #             # Close the figure to prevent memory leaks
    #             plt.close(fig)
            
    #         plots.append(fig)

    #     return plots
            

    # TODO: log original images out
    def log_original_tiles(self, data_generator, dataset="test", log_mlflow=False, year2=False):
        plots = []
        for tile_directory in data_generator.file_paths:
            tile_directory = os.path.join(data_generator.data_path, tile_directory)
            if year2:
                tile_directory.replace(str(data_generator.year1), str(data_generator.year2))
            tile_path = None
            if os.path.isdir(tile_directory):
                for file in os.listdir(tile_directory):
                    if file.endswith('reflectance_excluded_bands.tif'):
                        tile_path = os.path.join(tile_directory, file)

            if tile_path:
                # for now, going to just grab the hsi image because even if only doing on indices, the hsi will be what is displayed
                rgb = rasterio.open(tile_path)

                # r,g,b is 58,34,19
                r = rgb.read(58)
                g = rgb.read(34)
                b = rgb.read(19)
                
                # Mask
                r[r<0] = 0
                g[g<0] = 0
                b[b<0] = 0

                rgb_image = np.zeros((1000, 1000, 3))

                # Scale by max
                rgb_image[:, :, 0] = r/np.max(r)
                rgb_image[:, :, 1] = g/np.max(g)
                rgb_image[:, :, 2] = b/np.max(b)

                fig, ax = plt.subplots(figsize=(10, 10))
                ax.imshow(rgb_image, origin='lower')
                tile_name = tile_path.split("/")[-1]
                # upscaling specific title
                # ax.set_title(f'CHM Tile {tile_name.replace("_", " ")} - {"10 m" if data_generator.upscale_params else "1 m"} - {dataset}')
                ax.set_title(f'RGB Tile {tile_name.replace("_", " ")} - "1 m" - {dataset}')
                if log_mlflow:
                    mlflow.log_figure(fig, f'./rgb/rgb_{dataset}_tile_{tile_name}.png')
                    # Close the figure to prevent memory leaks
                    plt.close(fig)
                
                plots.append(fig)
                
            else:
                print('tile path not found')

        return plots

    def log_all_plots_to_grid(self, model, data_generator, dataset="test", plots={}):
        # TODO: plot_types can be replaced by key names in plots dict
        plot_types = ['RGB', 'Masks', 'True', 'Predicted', 'Error Intensity', 'Actual vs Predicted Scatter']
        # all_plot_lists = [rgbs, masks, trues, preds, error_intensities, actual_vs_pred_scatter]
        all_plot_lists = [
            plots['rgb'],
            plots['mask'],
            plots['y_true'],
            plots['y_pred'],
            plots['error_intensity'],
            plots['scatter']
        ]

        num_tiles = len(all_plot_lists[0])

        for tile_index in range(num_tiles):
            fig = plt.figure(figsize=(30, 36))
            gs = GridSpec(3, 2, figure=fig, wspace=0.3, hspace=0.15)

            for i, (plot_type, plot_list) in enumerate(zip(plot_types, all_plot_lists)):
                plot = plot_list[tile_index]
                outer_grid = GridSpecFromSubplotSpec(1, 2, subplot_spec=gs[i // 2, i % 2], width_ratios=[20, 1])
                ax = fig.add_subplot(outer_grid[0])

                if isinstance(plot, plt.Figure):
                    original_ax = plot.axes[0]

                    if plot_type == 'Actual vs Predicted Scatter':
                        if hasattr(plot, 'density_data'):
                            points = plot.density_data['points']
                            density = plot.density_data['density']
                            
                            # Recreate scatter plot with original data
                            sc = ax.scatter(points[:, 0], points[:, 1], 
                                        c=density, cmap='turbo',
                                        s=0.5, alpha=0.7)
                            
                            # Add colorbar
                            cbar_ax = fig.add_subplot(outer_grid[1])
                            cbar = fig.colorbar(sc, cax=cbar_ax)
                            cbar.set_label('Density', fontsize=16)
                            cbar.set_ticks([])

                            min_val = min(points[:, 0].min(), points[:, 1].min())
                            max_val = max(points[:, 0].max(), points[:, 1].max())
                            ax.set_xlim(min_val, max_val)
                            ax.set_ylim(min_val, max_val)

                            ax.plot([min_val, max_val], [min_val, max_val], 'r--', 
                                label='Perfect Predictions')
                            ax.tick_params(axis='both', labelsize=20)
                        else:
                            # If no density data, plot without density
                            offsets = original_ax.collections[0].get_offsets()
                            ax.scatter(offsets[:, 0], offsets[:, 1], 
                                    color='blue', s=0.5, alpha=0.7)

                            # Compute min and max values from the offsets
                            min_val = min(offsets[:, 0].min(), offsets[:, 1].min())
                            max_val = max(offsets[:, 0].max(), offsets[:, 1].max())

                            # Plot the perfect prediction line
                            ax.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Predictions')
                            ax.tick_params(axis='both', labelsize=20)
                    else:
                        for artist in original_ax.get_children():
                            if isinstance(artist, AxesImage):
                                im = ax.imshow(artist.get_array(), cmap=artist.get_cmap(),
                                            vmin=artist.get_clim()[0], vmax=artist.get_clim()[1],
                                            aspect='auto')
                                
                                if plot_type in ['True', 'Predicted', 'Error Intensity']:
                                    cbar_ax = fig.add_subplot(outer_grid[1])
                                    cbar = fig.colorbar(im, cax=cbar_ax)
                                    cbar.ax.tick_params(labelsize=20)

                    # Transfer axis properties
                    ax.set_xlabel(original_ax.get_xlabel(), fontsize=18)
                    ax.set_ylabel(original_ax.get_ylabel(), fontsize=18)
                    ax.set_xlim(original_ax.get_xlim())
                    ax.set_ylim(original_ax.get_ylim())
                    ax.set_title(original_ax.get_title(), fontsize=20) # TODO: set title to remove all of the excess stuff given already in the grid title
                    ax.tick_params(axis='both', labelsize=20)

                elif plot_type in ['Masks', 'RGB']:
                    ax.imshow(plot, cmap='binary' if plot_type == 'Masks' else None, aspect='equal')
                    ax.set_xlim([0, plot.shape[1]])
                    ax.set_ylim([plot.shape[0], 0])
                    ax.axis('off')

            tile_name = data_generator.file_paths[tile_index].split("/")[1]
            plt.suptitle(f'Tile {tile_name} - {dataset.capitalize()}', fontsize=24)
            mlflow.log_figure(fig, f'all_plots/{tile_name}_{dataset.capitalize()}.png')
            plt.close(fig)

        print(f"Logged {num_tiles} tile grids to MLflow.")

    def log_y_true_y_pred(self, model, data_generator, dataset="test", log_to_mlflow=False):
        y_true = data_generator.y
        y_pred = model.predict(data_generator)

        for name, array in [("y_true", y_true), ("y_pred", y_pred)]:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.npy') as temp_file:
                np.save(temp_file, array)
                temp_file.flush()
                temp_file_path = temp_file.name

            desired_filename = f"{name}_{dataset}.npy"
            temp_dest_path = os.path.join(tempfile.gettempdir(), desired_filename)
            shutil.copy(temp_file_path, temp_dest_path)

            mlflow.log_artifact(temp_dest_path)

            os.unlink(temp_file_path)
            os.unlink(temp_dest_path)

        print(f"Logged y_true and y_pred for {dataset} dataset to MLflow.")

    def log_scalers(self, data_generator, dataset="test"):
        with tempfile.TemporaryDirectory() as temp_dir:
            for i, scaler in enumerate(data_generator.scalers):
                # Define the custom file name
                file_name = f'scaler_{dataset}_branch_{i}.pkl'
                file_path = os.path.join(temp_dir, file_name)
                
                # Serialize the scaler to the custom-named file
                joblib.dump(scaler, file_path)
                
                # Log the artifact (temporary file)
                mlflow.log_artifact(file_path, artifact_path='scalers')

    def enable_determinism(self):
        # TODO: add ability to custmoize seed for this
        set_random_seed(69420)
        # If using TensorFlow, this will make GPU ops as deterministic as possible,
        # but it will affect the overall performance, so be mindful of that.
        tf.config.experimental.enable_op_determinism()


    # TODO: DELETE if scaling y is useless
    def log_actual_vs_predicted_random_sample(
        self,
        model,
        data_generator,
        dataset="test",
        scale_method=None,
        log_mlflow=False
    ):
        """
        Plots a small random sample of actual vs. predicted values (in meters) for a regression dataset.
        """
        print(f'** Plotting Actual vs. Predicted Random Sample for {dataset.capitalize()} Dataset **')

        # ➊ Pull out the scaled y_true and scaled y_pred
        y_true_scaled = data_generator.y.flatten()                   # shape: (N,)
        y_pred_scaled = model.predict(data_generator).flatten()      # shape: (N,)

        # ➋ Inverse‐scale both back to “meters”
        y_true_raw = data_generator.label_scaler.inverse_transform(
            y_true_scaled.reshape(-1, 1)
        ).ravel()
        y_pred_raw = data_generator.label_scaler.inverse_transform(
            y_pred_scaled.reshape(-1, 1)
        ).ravel()

        num_samples = y_true_raw.shape[0]
        random_idxs = np.random.choice(num_samples, size=min(num_samples, 50), replace=False)

        sample_y_true = y_true_raw[random_idxs]
        sample_y_pred = y_pred_raw[random_idxs]

        # ➌ Calculate average of the predictions (in meters)
        avg_pred = np.average(sample_y_pred)

        # ➍ Plot
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(sample_y_true, color='orange', label='Actual (m)')
        ax.plot(sample_y_pred, color='lightblue', label='Predicted (m)')
        ax.axhline(y=avg_pred, color='green', label='Average Predicted (m)')
        ax.set_xlabel('Random Sample Index')
        ax.set_ylabel('Canopy Height Change (meters)')
        ax.set_title(f'Actual vs. Predicted Values ({dataset.capitalize()})')

        # ➎ Compute & annotate R² on the *entire* set (in meters)
        r2 = r2_score(y_true_raw, y_pred_raw)
        annotation = mpl.text.Annotation(
            f"R² = {r2:.2f}",
            xy=(0.99, 0.96),
            xycoords='axes fraction',
            ha='right',
            va='center',
            fontsize=10
        )
        ax.add_artist(annotation)

        ax.legend(loc='lower right')
        fig.tight_layout()

        if log_mlflow:
            mlflow.log_figure(fig, f"actual_vs_predicted_random_{dataset}.png")
            plt.close(fig)

        return [fig]

    def log_actual_vs_predicted_scatter(
        self,
        model,
        data_generator,
        dataset="test",
        log_mlflow=False
    ):
        """
        Plots an aggregate scatter of all actual vs. predicted values (in meters), 
        plus tile‐by‐tile scatter.  Returns a list of figure objects.
        """
        print(f'** Plotting Actual vs. Predicted Scatter Plot for {dataset.capitalize()} Dataset **')

        # ➊ Pull scaled truths & preds
        y_true_scaled = data_generator.y.flatten()                # shape: (N,)
        y_pred_scaled = model.predict(data_generator).flatten()   # shape: (N,)

        # ➋ Inverse‐scale both
        y_true_raw = data_generator.label_scaler.inverse_transform(
            y_true_scaled.reshape(-1, 1)
        ).ravel()
        y_pred_raw = data_generator.label_scaler.inverse_transform(
            y_pred_scaled.reshape(-1, 1)
        ).ravel()

        # ➌ Filter out NaNs (masked areas)
        valid_mask = (~np.isnan(y_true_raw)) & (~np.isnan(y_pred_raw))
        y_true_valid = y_true_raw[valid_mask]
        y_pred_valid = y_pred_raw[valid_mask]

        # ➍ First: an aggregate scatter plot
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.scatter(y_true_valid, y_pred_valid, color='blue', s=0.5, alpha=0.7)

        min_val = min(y_true_valid.min(), y_pred_valid.min())
        max_val = max(y_true_valid.max(), y_pred_valid.max())
        ax.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Prediction')
        ax.set_xlabel(f'{data_generator.truth.upper()} Truth (m)')
        ax.set_ylabel(f'{data_generator.truth.upper()} Prediction (m)')
        ax.set_title(f'{data_generator.truth.upper()} Actual vs. Predicted ({dataset.capitalize()})')
        ax.set_aspect('equal', 'box')
        ax.legend()
        fig.tight_layout()

        if log_mlflow:
            mlflow.log_figure(fig, f"scatter_all_{dataset}.png")
            plt.close(fig)

        plots = [fig]

        # ➎ Now: expand back into tile-shaped arrays, tile by tile
        master_mask = data_generator.master_mask  # boolean mask (True=masked pixel)
        y_pred_expanded = np.full(master_mask.shape, np.nan)
        y_true_expanded = np.full(master_mask.shape, np.nan)

        # Place raw (unmasked) values back into each location
        masked_idxs = np.where(~master_mask)
        y_pred_expanded[masked_idxs] = y_pred_raw
        y_true_expanded[masked_idxs] = y_true_raw

        # Determine how many tiles & their dimension
        if data_generator.upscale_params:
            original_size = 1000
            new_dim = original_size // data_generator.upscale_params.get('window_size', 10)
            num_tiles = master_mask.size // (new_dim * new_dim)
            y_pred_tiles = y_pred_expanded.reshape(num_tiles, new_dim, new_dim)
            y_true_tiles = y_true_expanded.reshape(num_tiles, new_dim, new_dim)
        else:
            # assume 1000×1000 patches
            num_tiles = y_pred_expanded.shape[0] // 1000000
            y_pred_tiles = y_pred_expanded.reshape((num_tiles, 1000, 1000))
            y_true_tiles = y_true_expanded.reshape((num_tiles, 1000, 1000))

        # ➏ Loop through each tile & plot scatter
        for i in range(num_tiles):
            tile_y_true = y_true_tiles[i].flatten()
            tile_y_pred = y_pred_tiles[i].flatten()

            valid_mask_tile = (~np.isnan(tile_y_true)) & (~np.isnan(tile_y_pred))
            tile_true_valid = tile_y_true[valid_mask_tile]
            tile_pred_valid = tile_y_pred[valid_mask_tile]

            if len(tile_true_valid) < 2:
                print(f"Skipping tile {i} due to insufficient valid data.")
                continue

            fig_tile, ax_tile = plt.subplots(figsize=(10, 10))
            ax_tile.scatter(tile_true_valid, tile_pred_valid, color='blue', s=10, alpha=0.7)

            mn = min(tile_true_valid.min(), tile_pred_valid.min())
            mx = max(tile_true_valid.max(), tile_pred_valid.max())
            ax_tile.plot([mn, mx], [mn, mx], 'r--', label='Perfect Prediction')
            ax_tile.set_xlabel(f'{data_generator.truth.upper()} Truth (m)')
            ax_tile.set_ylabel(f'{data_generator.truth.upper()} Prediction (m)')
            tile_name = data_generator.file_paths[i].split("/")[-1].replace("_", " ")
            ax_tile.set_title(f'{data_generator.truth.upper()} Tile {tile_name} ({dataset.capitalize()})')
            ax_tile.set_aspect('equal', 'box')
            ax_tile.legend()
            fig_tile.tight_layout()

            if log_mlflow:
                mlflow.log_figure(fig_tile, f"scatter_tile_{dataset}_{tile_name}.png")
                plt.close(fig_tile)

            plots.append(fig_tile)

        return plots

    def log_r2_score(self, model, data_generator, dataset="test", log_mlflow=False):
        """
        Computes and logs R² on the unscaled (meter) values.
        """
        # ➊ Pull scaled y_true and scaled y_pred
        y_true_scaled = data_generator.y.flatten()
        y_pred_scaled = model.predict(data_generator).flatten()

        # ➋ Inverse‐scale to meters
        y_true_raw = data_generator.label_scaler.inverse_transform(
            y_true_scaled.reshape(-1, 1)
        ).ravel()
        y_pred_raw = data_generator.label_scaler.inverse_transform(
            y_pred_scaled.reshape(-1, 1)
        ).ravel()

        # ➌ Compute R²
        r2 = r2_score(y_true_raw, y_pred_raw)

        if log_mlflow:
            mlflow.log_metric(f"r2_{dataset}", r2)
        else:
            print(f"R² ({dataset}): {r2:.4f}")

    def log_error_intensity_plot(self, model, data_generator, dataset="test", log_mlflow=False):
        """
        Plots error‐intensity heatmaps in meters for each tile.  
        """
        # ➊ Pull scaled y_true and scaled y_pred
        y_true_scaled = data_generator.y.flatten()
        y_pred_scaled = model.predict(data_generator).flatten()

        # ➋ Inverse‐scale both back to meters
        y_true_raw = data_generator.label_scaler.inverse_transform(
            y_true_scaled.reshape(-1, 1)
        ).ravel()
        y_pred_raw = data_generator.label_scaler.inverse_transform(
            y_pred_scaled.reshape(-1, 1)
        ).ravel()

        master_mask = data_generator.master_mask

        # Expand them back into full‐tile shapes
        y_pred_expanded = np.full(master_mask.shape, np.nan)
        y_true_expanded = np.full(master_mask.shape, np.nan)
        masked_idxs = np.where(~master_mask)

        y_pred_expanded[masked_idxs] = y_pred_raw
        y_true_expanded[masked_idxs] = y_true_raw

        # Determine tiling
        if data_generator.upscale_params:
            orig = 1000
            new_dim = orig // data_generator.upscale_params.get('window_size', 10)
            num_tiles = master_mask.size // (new_dim * new_dim)
            y_pred_tiles = y_pred_expanded.reshape(num_tiles, new_dim, new_dim)
            y_true_tiles = y_true_expanded.reshape(num_tiles, new_dim, new_dim)
        else:
            num_tiles = y_pred_expanded.shape[0] // 1000000
            y_pred_tiles = y_pred_expanded.reshape((num_tiles, 1000, 1000))
            y_true_tiles = y_true_expanded.reshape((num_tiles, 1000, 1000))

        plots = []
        for i in range(num_tiles):
            errors = y_pred_tiles[i] - y_true_tiles[i]  # error (m), can be negative/positive

            # Clip extremes for visualization
            above_clip, below_clip = 5.0, -5.0  # meters
            errors_clipped = np.clip(errors, below_clip, above_clip)

            fig, ax = plt.subplots(figsize=(10, 10))
            cmap = mpl.colormaps.get_cmap('turbo')
            cmap.set_bad(color='black')

            cax = ax.imshow(
                errors_clipped,
                cmap=cmap,
                origin='lower',
                vmin=below_clip,
                vmax=above_clip
            )
            cbar = fig.colorbar(cax, ax=ax, label='Error (meters)')
            ax.set_xlim([0, errors_clipped.shape[1]])
            ax.set_ylim([0, errors_clipped.shape[0]])
            ax.set_xlabel('X')
            ax.set_ylabel('Y')

            tile_name = data_generator.file_paths[i].split("/")[-1].replace("_", " ")
            ax.set_title(f'Error Intensity (m) - Tile {tile_name} ({dataset.capitalize()})')

            plots.append(fig)
            if log_mlflow:
                mlflow.log_figure(fig, f"error_intensity_{dataset}_tile_{tile_name}.png")
                plt.close(fig)

        return plots

    def log_pred(self, model, data_generator, dataset="test", log_mlflow=False):
        """
        Plots & logs predicted CHC (in meters) for each tile.
        """
        y_pred_scaled = model.predict(data_generator).flatten()
        master_mask = data_generator.master_mask

        # Inverse‐scale to raw meters
        y_pred_raw = data_generator.label_scaler.inverse_transform(
            y_pred_scaled.reshape(-1, 1)
        ).ravel()

        # Expand back into full mask shape
        y_pred_expanded = np.full(master_mask.shape, np.nan)
        masked_idxs = np.where(~master_mask)
        y_pred_expanded[masked_idxs] = y_pred_raw

        # Tile‐reshape
        if data_generator.upscale_params:
            orig = 1000
            new_dim = orig // data_generator.upscale_params.get('window_size', 10)
            num_tiles = y_pred_expanded.size // (new_dim * new_dim)
            y_pred_tiles = y_pred_expanded.reshape(num_tiles, new_dim, new_dim)
        else:
            y_pred_tiles = y_pred_expanded.reshape(-1, 1000, 1000)
            num_tiles = y_pred_tiles.shape[0]

        plots = []
        for i in range(num_tiles):
            tile_pred = y_pred_tiles[i]

            fig, ax = plt.subplots(figsize=(10, 10))
            cmap = mpl.colormaps.get_cmap('gray_r')
            cmap.set_bad(color='black')

            cax = ax.imshow(
                tile_pred,
                cmap=cmap,
                interpolation='nearest',
                origin='lower',
                # vmin=0,
                # vmax=40
            )
            fig.colorbar(cax, ax=ax, label='CHC (m)')
            ax.set_xlim([0, tile_pred.shape[1]])
            ax.set_ylim([0, tile_pred.shape[0]])
            ax.set_xticks([])
            ax.set_yticks([])

            tile_name = data_generator.file_paths[i].split("/")[-1].replace("_", " ")
            truth_name = data_generator.truth
            ax.set_title(
                f'Predicted {truth_name.upper()} (m) - Tile {tile_name} ({dataset.capitalize()})'
            )

            plots.append(fig)
            if log_mlflow:
                mlflow.log_figure(fig, f"pred_{truth_name.lower()}_{dataset}_tile_{tile_name}.png")
                plt.close(fig)

        return plots

    def log_truth(self, data_generator, dataset="test", log_mlflow=False):
        """
        Plots & logs true CHC (in meters) for each tile.
        """
        y_true_scaled = data_generator.y.flatten()

        # Inverse‐scale
        y_true_raw = data_generator.label_scaler.inverse_transform(
            y_true_scaled.reshape(-1, 1)
        ).ravel()

        master_mask = data_generator.master_mask
        y_true_expanded = np.full(master_mask.shape, np.nan)
        masked_idxs = np.where(~master_mask)
        y_true_expanded[masked_idxs] = y_true_raw

        # Tile‐reshape
        if data_generator.upscale_params:
            orig = 1000
            new_dim = orig // data_generator.upscale_params.get('window_size', 10)
            num_tiles = y_true_expanded.size // (new_dim * new_dim)
            y_true_tiles = y_true_expanded.reshape(num_tiles, new_dim, new_dim)
        else:
            y_true_tiles = y_true_expanded.reshape(-1, 1000, 1000)
            num_tiles = y_true_tiles.shape[0]

        plots = []
        for i in range(num_tiles):
            tile_true = y_true_tiles[i]

            fig, ax = plt.subplots(figsize=(10, 10))
            cmap = mpl.colormaps.get_cmap('gray_r')
            cmap.set_bad(color='black')

            ax.imshow(tile_true, cmap=cmap, interpolation='nearest', origin='lower')
            ax.set_xlim([0, tile_true.shape[1]])
            ax.set_ylim([0, tile_true.shape[0]])
            ax.set_xticks([])
            ax.set_yticks([])

            tile_name = data_generator.file_paths[i].split("/")[-1].replace("_", " ")
            truth_name = data_generator.truth
            ax.set_title(f'True {truth_name.upper()} (m) - Tile {tile_name} ({dataset.capitalize()})')

            plots.append(fig)
            if log_mlflow:
                mlflow.log_figure(fig, f"truth_{truth_name.lower()}_{dataset}_tile_{tile_name}.png")
                plt.close(fig)

        return plots



    ################################ NEW FUNCTIONS FOR DAVE #####################################################
    ########################## IMPORTANT: THIS IS SEPARATE FROM THE SCALING Y UNCOMMENT/COMMENT STUFF ABOVE ######

    def log_chc_histogram(self, model, data_generator, dataset="test", log_mlflow=False):
        """Make histogram plots of CHC overlaying truth and predicted distributions"""
        
        y_true = data_generator.y
        y_pred = model.predict(data_generator)
        
        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y
        
        fig, axs = plt.subplots(figsize=(10, 6))
        
        # Plot histograms with overlay
        axs.hist(y_true.flatten(), bins=50, alpha=0.7, color='black', 
                label='Truth', density=True, edgecolor='black', linewidth=1)
        axs.hist(y_pred.flatten(), bins=50, alpha=0.7, color='grey', 
                label='Predicted', density=True, edgecolor='grey', linewidth=1)
        
        axs.set_xlabel('CHC Values')
        axs.set_ylabel('Density')
        axs.set_title('CHC Distribution: Truth vs Predicted')
        axs.legend()
        axs.grid(True, alpha=0.3)
        
        if log_mlflow:
            mlflow.log_figure(fig, f"chc_histogram_{dataset}.png")
        
        return [fig]

    def log_chc_statistics_table(self, model, data_generator, dataset="test", log_mlflow=False):
        """Compute min, max, median, mean and st dev for truth and predicted distributions"""
        
        y_true = data_generator.y
        y_pred = model.predict(data_generator)
        
        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y
        
        # Calculate statistics
        stats_data = {
            'Statistic': ['Min', 'Max', 'Median', 'Mean', 'Std Dev'],
            'Truth': [
                np.min(y_true),
                np.max(y_true),
                np.median(y_true),
                np.mean(y_true),
                np.std(y_true)
            ],
            'Predicted': [
                np.min(y_pred),
                np.max(y_pred),
                np.median(y_pred),
                np.mean(y_pred),
                np.std(y_pred)
            ]
        }
        
        # Create table visualization
        fig, axs = plt.subplots(figsize=(10, 6))
        axs.axis('tight')
        axs.axis('off')
        
        table = axs.table(cellText=[[f"{val:.4f}" if isinstance(val, (int, float)) else val 
                                    for val in row] for row in zip(*stats_data.values())],
                        colLabels=list(stats_data.keys()),
                        cellLoc='center',
                        loc='center')
        table.auto_set_font_size(False)
        table.set_fontsize(12)
        table.scale(1.2, 1.5)
        
        axs.set_title('CHC Statistics: Truth vs Predicted', pad=20)
        
        if log_mlflow:
            mlflow.log_figure(fig, f"chc_statistics_table_{dataset}.png")
            # Also log as text for easy access
            stats_text = "\n".join([f"{stat}: Truth={truth:.4f}, Predicted={pred:.4f}" 
                                for stat, truth, pred in zip(stats_data['Statistic'], 
                                                            stats_data['Truth'], 
                                                            stats_data['Predicted'])])
            mlflow.log_text(stats_text, f"chc_statistics_{dataset}.txt")
        
        return [fig], stats_data

    def log_chc_scatterplot(self, model, data_generator, dataset="test", log_mlflow=False):
        """Make scatterplot of truth CHC vs predicted CHC with equal axes"""
        
        y_true = data_generator.y
        y_pred = model.predict(data_generator)
        
        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y
        
        fig, axs = plt.subplots(figsize=(10, 8))
        
        # Determine common axis range
        min_val = min(np.min(y_true), np.min(y_pred))
        max_val = max(np.max(y_true), np.max(y_pred))
        
        # Create scatterplot with smallest possible points
        axs.scatter(y_true.flatten(), y_pred.flatten(), s=1, alpha=0.6, color='blue')
        
        # Add perfect prediction line
        axs.plot([min_val, max_val], [min_val, max_val], 'r--', lw=2, 
                label='Perfect Prediction')
        
        axs.set_xlim([min_val, max_val])
        axs.set_ylim([min_val, max_val])
        axs.set_xlabel('Truth CHC')
        axs.set_ylabel('Predicted CHC')
        axs.set_title('CHC: Truth vs Predicted')
        axs.legend()
        axs.grid(True, alpha=0.3)
        
        # Ensure equal aspect ratio
        axs.set_aspect('equal', adjustable='box')
        
        if log_mlflow:
            mlflow.log_figure(fig, f"chc_scatterplot_{dataset}.png")
        
        return [fig]

    def log_chc_metrics_table(self, model, data_generator, dataset="test", log_mlflow=False):
        """Compute MSE, RMSE, MAE and R^2 for CHC truth vs predicted"""
        
        y_true = data_generator.y
        y_pred = model.predict(data_generator)
        
        # y_true = y_true / 1000 # unscaling back to original y
        # y_pred = y_pred / 1000 # unscaling back to original y
        
        # Calculate regression metrics
        mse = mean_squared_error(y_true, y_pred)
        rmse = np.sqrt(mse)
        mae = mean_absolute_error(y_true, y_pred)
        r2 = r2_score(y_true, y_pred)
        
        metrics_data = {
            'Metric': ['MSE', 'RMSE', 'MAE', 'R²'],
            'Value': [mse, rmse, mae, r2]
        }
        
        # Create table visualization
        fig, axs = plt.subplots(figsize=(8, 4))
        axs.axis('tight')
        axs.axis('off')
        
        table = axs.table(cellText=[[metric, f"{value:.6f}"] 
                                for metric, value in zip(metrics_data['Metric'], 
                                                        metrics_data['Value'])],
                        colLabels=['Metric', 'Value'],
                        cellLoc='center',
                        loc='center')
        table.auto_set_font_size(False)
        table.set_fontsize(12)
        table.scale(1.2, 1.5)
        
        axs.set_title('CHC Regression Metrics', pad=20)
        
        if log_mlflow:
            mlflow.log_figure(fig, f"chc_metrics_table_{dataset}.png")
            # Log individual metrics
            mlflow.log_metric(f"chc_mse_{dataset}", mse)
            mlflow.log_metric(f"chc_rmse_{dataset}", rmse)
            mlflow.log_metric(f"chc_mae_{dataset}", mae)
            mlflow.log_metric(f"chc_r2_{dataset}", r2)
            # Also log as text
            metrics_text = "\n".join([f"{metric}: {value:.6f}" 
                                    for metric, value in zip(metrics_data['Metric'], 
                                                            metrics_data['Value'])])
            mlflow.log_text(metrics_text, f"chc_metrics_{dataset}.txt")
        
        return [fig], metrics_data
