import os
import numpy as np
import secrets
import warnings
from tensorflow.keras.utils import Sequence
import rasterio
from rasterio.plot import reshape_as_image
from sklearn.preprocessing import MinMaxScaler, StandardScaler, RobustScaler
from sklearn.model_selection import train_test_split
from multiprocessing import Pool

import time

class TwoDDataGenerator(Sequence):
    def __init__(self, branch_inputs, data_path, file_paths=None, truth="chm", batch_size=1,
                 full_year=False, scale_data_method=None, randomize_pixels=False, split_tuple=(0, 0, 0),
                 split_mode='train', split_method='tiles', features=['brdf, indices, dtm, slope, aspect'], 
                 test_type='reg', bins=None, year1=None, year2=None, upscale_params=False, scalers=None):
        
        self.branch_inputs = branch_inputs
        self.data_path = data_path
        self.file_paths = file_paths
        self.truth = truth
        self.split_tuple = split_tuple
        self.split_mode = split_mode
        self.split_method = split_method
        self.batch_size = batch_size
        self.full_year = full_year
        self.randomize_pixels = randomize_pixels
        self.scale_data_method = scale_data_method
        self.upscale_params = upscale_params

        self.test_type = test_type
        self.bins = bins

        self.year1 = year1
        self.year2 = year2

        self.scalers = scalers

        self.X, self.y, self.shape = self.main()

    def _resolve_years_for_file(self, file_path):
        if isinstance(self.year1, int) and isinstance(self.year2, int):
            return self.year1, self.year2

        if isinstance(self.year1, dict) and isinstance(self.year2, dict):
            path_lower = file_path.lower()
            year2_lookup = {str(key).lower(): value for key, value in self.year2.items()}
            for site_key, year1_value in self.year1.items():
                normalized_key = str(site_key).lower()
                if normalized_key in path_lower:
                    if normalized_key not in year2_lookup:
                        raise ValueError(
                            f"Missing year2 entry for site '{site_key}' in config."
                        )
                    return year1_value, year2_lookup[normalized_key]

            raise ValueError(
                "Could not resolve year1/year2 for canopy truth file. "
                f"Add matching keys to year1/year2 for path: {file_path}"
            )

        raise TypeError(
            "year1 and year2 must both be integers or both be dictionaries keyed by site name."
        )

    def main(self):
        if self.split_method == "tiles":
            self.file_paths = self.split_data(self.file_paths)
            print(f'using following tiles for {self.split_mode} set: {self.file_paths}')

        print('** Loading and preprocessing truth data **')
        y, y_masks = self._generate_NDVI_truth() if self.truth == "narrow_ndvi" or self.truth == "wide_ndvi" else self._generate_canopy_truth()
        X = []
        X_masks = []

        print('** Loading and preprocessing X data for all branches **')
        for branch_num, branch in enumerate(self.branch_inputs):
            x, x_masks = self._generate_X(branch)
            X.append(x)
            X_masks.append(x_masks)

        print('** Merging masks into master_mask **')
        X_masks, y_masks = self._merge_masks(X, y, X_masks, y_masks)

        if self.upscale_params:
            print(f'** Upscaling y using {self.upscale_params.get("calc_method", "median")} **')
            y, y_masks = self.upscale(y, ~y_masks, self.upscale_params, is_truth=True)
            print(f'** Upscaling x using {self.upscale_params.get("calc_method", "median")} **')
            for i, (x, x_mask) in enumerate(zip(X, X_masks)):
                X[i], X_masks[i] = self.upscale(x, ~x_mask, self.upscale_params, is_truth=False)
            print('** Merging new upscaled masks into new master_mask **')
            X_masks, y_masks = self._merge_masks(X, y, X_masks, y_masks)

        print('** Applying masks to X and y **')
        X, y = self._apply_masks(X, y, X_masks, y_masks)

        if self.split_method == "pixels":
            if self.split_tuple != (0, 0, 0):
                print('** Splitting X and y data into train, val, and test sets **')
                X = [self.split_data(x) for x in X]
                y = self.split_data(y)
        
        if self.scale_data_method:
            print('** Scaling input data **')
            X = self.scale(X)
            print(f'INFO: np.mean(X) post scaling: {np.mean(X[0])}')
            print(f'INFO: np.std(X) post scaling: {np.std(X[0])}')
            print(f'INFO: np.max(X) post scaling: {np.max(X[0])}')
            print(f'INFO: np.min(X) post scaling: {np.min(X[0])}')

        shape = X[0].shape[1:]  # Store the shape of a single image
        return X, y, shape

    def _generate_NDVI_truth(self):
        y = []
        y_masks = []
        ndvi_truth_files = self.get_file_paths_and_bands(truth=self.truth)
        for ndvi_truth_file in ndvi_truth_files:
            with rasterio.open(ndvi_truth_file['ndvi']['file_path']) as src:
                if self.truth == "narrow_ndvi":
                    band_53, band_87 = src.read(ndvi_truth_file['ndvi']['bands'])
                    truth_array = (band_87 - band_53) / (band_87 + band_53)
                else:
                    truth_array = src.read(ndvi_truth_file['ndvi']['bands'])
            truth_mask = self.mask_y(truth_array, 'ndvi')
            y.append(truth_array)
            y_masks.append(truth_mask)
        
        y = np.concatenate(y, axis=0)
        y_masks = np.concatenate(y_masks, axis=0)
        if self.test_type == 'cat':
            y = np.digitize(y, self.bins)
        print(f'INFO: Resulting shapes of y and y_mask: {y.shape} , {y_masks.shape}')
        return y, y_masks

    def _generate_canopy_truth(self):
        canopy_truth_files = self.get_file_paths_and_bands(truth=self.truth)
        y = []
        y_masks = []
        for canopy_truth_file in canopy_truth_files:
            with rasterio.open(canopy_truth_file['canopy_height']['file_path']) as src:
                canopy_height = src.read(1)
            if self.truth == "chc":
                year1, year2 = self._resolve_years_for_file(
                    canopy_truth_file['canopy_height']['file_path']
                )

                ch_year_1 = canopy_height
                canopy_truth_comparison_file = canopy_truth_file['canopy_height']['file_path'].replace(str(year1), str(year2))
                with rasterio.open(canopy_truth_comparison_file) as src:
                    ch_year_2 = src.read(1)
                chc = ch_year_1 - ch_year_2 / (year2 - year1)
                y.append(chc)
                y_masks.append(self.mask_y(chc, 'chc'))
            else:
                y.append(canopy_height)
                y_masks.append(self.mask_y(canopy_height, 'chm'))
        y = np.stack(y, axis=0)
        y_masks = np.stack(y_masks, axis=0)
        if self.test_type == 'cat':
            y = np.digitize(y, self.bins)
        print(f'INFO: Resulting shapes of y and y_mask: {y.shape} , {y_masks.shape}')
        return y, y_masks

    def _generate_X(self, branch):
        print(f'** Loading in and generating masks for branch {branch.get("branch_name", None)} **')
        tile_features = self.get_file_paths_and_bands(branch)
        X = []
        X_mask = []
        for tile_feature in tile_features:
            x = []
            x_mask = []
            if 'hyperspectral' in tile_feature:
                with rasterio.open(tile_feature['hyperspectral']['file_path']) as src:
                    print(f'taking hyperspectral from file: {tile_feature["hyperspectral"]["file_path"]}')
                    bands = tile_feature['hyperspectral']['bands'] if tile_feature['hyperspectral']['bands'] else None
                    hyperspectral = src.read(bands)
                    hyperspectral = reshape_as_image(hyperspectral)
                    hyperspectral_mask = self.mask_X(hyperspectral, 'hyperspectral', no_data_fill_value=src.nodata)
                x.append(hyperspectral)
                x_mask.append(hyperspectral_mask)

            if 'savgol' in tile_feature:
                with rasterio.open(tile_feature['savgol']['file_path']) as src:
                    print(f'taking savgol from file: {tile_feature["savgol"]["file_path"]}')
                    bands = tile_feature['savgol']['bands'] if tile_feature['savgol']['bands'] else None
                    savgol = src.read(bands)
                    savgol = reshape_as_image(savgol)
                    savgol_mask = self.mask_X(savgol, 'savgol', no_data_fill_value=src.nodata)
                x.append(savgol)
                x_mask.append(savgol_mask)

            if 'brdf' in tile_feature:
                with rasterio.open(tile_feature['brdf']['file_path']) as src:
                    print(f'taking brdf from file: {tile_feature["brdf"]["file_path"]}')
                    brdf = src.read(17)
                    brdf = brdf[np.newaxis, :, :]
                brdf_mask = self.mask_X(brdf, 'brdf')
                x.append(brdf)
                x_mask.append(brdf_mask)

            if 'indices' in tile_feature:
                with rasterio.open(tile_feature['indices']['file_path']) as src:
                    print(f'taking indices from file: {tile_feature["indices"]["file_path"]}')
                    indices, indices_masks = self.generate_indices(src, tile_feature)

                x.append(indices)
                x_mask.append(indices_masks)

            if 'dtm' in tile_feature:
                with rasterio.open(tile_feature['dtm']['file_path']) as src:
                    print(f'taking dtm from file: {tile_feature["dtm"]["file_path"]}')
                    dtm = src.read(1)
                    dtm = dtm[np.newaxis, :, :]
                dtm_mask = self.mask_X(dtm, 'dtm')
                x.append(dtm)
                x_mask.append(dtm_mask)
            
            if 'dsm' in tile_feature:
                with rasterio.open(tile_feature['dsm']['file_path']) as src:
                    print(f'taking dsm from file: {tile_feature["dsm"]["file_path"]}')
                    dsm = src.read(1)
                    dsm = dsm[np.newaxis, :, :]
                dsm_mask = self.mask_X(dsm, 'dsm')
                x.append(dsm)
                x_mask.append(dsm_mask)

            if 'slope' in tile_feature:
                with rasterio.open(tile_feature['slope']['file_path']) as src:
                    print(f'taking slope from file: {tile_feature["slope"]["file_path"]}')
                    slope = src.read(1)
                    slope = slope[np.newaxis, :, :]
                slope_mask = self.mask_X(slope, 'slope')
                x.append(slope)
                x_mask.append(slope_mask)
                
            if 'aspect' in tile_feature:
                with rasterio.open(tile_feature['aspect']['file_path']) as src:
                    print(f'taking aspect from file: {tile_feature["aspect"]["file_path"]}')
                    aspect = src.read(1)
                    aspect = aspect[np.newaxis, :, :]
                aspect_mask = self.mask_X(aspect, 'aspect')
                x.append(aspect)
                x_mask.append(aspect_mask)

            if 'rugosity' in tile_feature:
                with rasterio.open(tile_feature['rugosity']['file_path']) as src:
                    print(f'taking rugosity from file: {tile_feature["rugosity"]["file_path"]}')
                    rugosity = src.read(1)
                    rugosity = rugosity[np.newaxis, :, :]
                rugosity_mask = self.mask_X(rugosity, 'rugosity')
                x.append(rugosity)
                x_mask.append(rugosity_mask)

            if 'chm' in tile_feature:
                with rasterio.open(tile_feature['chm']['file_path']) as src:
                    print(f'taking chm from file: {tile_feature["chm"]["file_path"]}')
                    chm = src.read(1)
                    chm = chm[np.newaxis, :, :]
                chm_mask = self.mask_X(chm, 'chm')
                x.append(chm)
                x_mask.append(chm_mask)

            if len(x) > 0:
                x = np.concatenate(x, axis=0)
                x_mask = np.concatenate(x_mask, axis=0)
                X.append(x)
                X_mask.append(x_mask)
            else:
                print(f'BAD TILE: {tile_feature}. Skipping over. Check if this directory has the file you need.')
                [print(i) for i in tile_feature]
                [print(f'BAD FILE: {i["file_path"]}') for i in tile_feature]
        
        X = np.stack(X, axis=0)
        X_mask = np.stack(X_mask, axis=0)

        print(f'INFO: Resulting shapes of X and X_mask for branch {branch["branch_name"]}: {X.shape} , {X_mask.shape}')
        return X, X_mask

    def generate_indices(self, src, tile_feature):
        bands = tile_feature['indices']['bands'] if tile_feature['indices']['bands'] else [14,15,16,17,18,20,21,9,13,22]
        indices = []
        indices_mask = []
        for band in bands:
            indice = src.read(band)
            print(f'indice shape: {indice.shape}')
            indice_mask = self.mask_indices(indice, 'indice')
            indices.append(indice)
            indices_mask.append(indice_mask)
            print(f'indice_ mask shape: {indice_mask.shape}')
        
        indices = np.stack(indices, axis=0)
        indices_mask = np.stack(indices_mask, axis=0)

        print(f'final shape indices from generate_indices: {indices.shape}')
        print(f'final shape indices_mask from generate_indices: {indices_mask.shape}')
        return indices, indices_mask

    def mask_indices(self, x, input_type, no_data_fill_value=None):
        match input_type:
            case 'ndvi':
                nan_mask = np.isnan(x)
                ndvi_mask = (x < 0.5)
                x_mask = nan_mask | ndvi_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for ndvi specifically')
            case 'evi' | 'evi2' | 'savi':
                nan_mask = np.isnan(x)
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for evi specifically')
            case 'indice':
                nan_mask = np.isnan(x)
                less_than_mask = (x < -1.0)
                greater_than_mask = (x > 1.0)
                x_mask = nan_mask | less_than_mask | greater_than_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for indices')
            case _:
                raise TypeError('inappropriate input_type for mask_indices()')

        return x_mask

    def scale(self, X):
        if self.scalers is None:
            self.scalers = []
            for i, x in enumerate(X):
                scaler = self._get_scaler()
                shape = x.shape
                x_flat = x.reshape(-1, shape[-1])
                X[i] = scaler.fit_transform(x_flat).reshape(shape)
                self.scalers.append(scaler)
        else:
            for i, (x, scaler) in enumerate(zip(X, self.scalers)):
                shape = x.shape
                x_flat = x.reshape(-1, shape[-1])
                X[i] = scaler.transform(x_flat).reshape(shape)
        return X

    def _get_scaler(self):
        match self.scale_data_method:
            case 'minmax':
                return MinMaxScaler(feature_range=(-1,1))
            case 'standard':
                return StandardScaler()
            case 'robust':
                return RobustScaler()
            case _:
                warnings.warn('Scaling method not valid. Defaulting to minmax scaler with feature_range=(-1,1).')
                return MinMaxScaler(feature_range=(-1,1))

    def split_data(self, data):
        if sum(self.split_tuple[:3]) != 100:
            raise ValueError("split_tuple must sum to 100.")
        train_ratio, val_ratio, test_ratio = self.split_tuple[0] / 100.0, self.split_tuple[1] / 100.0, self.split_tuple[2] / 100.0
        data_train, data_test = train_test_split(data, test_size=test_ratio, random_state=420)
        data_train, data_val = train_test_split(data_train, test_size=val_ratio / (train_ratio + val_ratio), random_state=420)
        
        match self.split_mode:
            case 'train':
                print(f'INFO: Resulting train split: {len(data_train) if self.split_method == "tiles" else data_train.shape}')
                return data_train
            case 'validate':
                print(f'INFO: Resulting validation split: {len(data_val) if self.split_method == "tiles" else data_val.shape}')
                return data_val
            case 'test':
                print(f'INFO: Resulting validation split: {len(data_test) if self.split_method == "tiles" else data_test.shape}')
                return data_test
            case _:
                print("ERROR -- Split mode must be: 'train', 'validate', or 'test'")
                exit(-1)

    def mask_X(self, x, input_type, no_data_fill_value=None):
        match input_type:
            case 'hyperspectral':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for hyperspectral')
            case 'savgol':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for savgol')
            case 'brdf':
                nan_mask = np.isnan(x)
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for brdf')
            case 'ndvi':
                nan_mask = np.isnan(x)
                ndvi_mask = (x < 0.5)
                x_mask = nan_mask | ndvi_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for ndvi specifically')
            case 'dtm' | 'dsm':
                nan_mask = np.isnan(x)
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for {input_type}')
            case 'slope':
                nan_mask = np.isnan(x)
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for slope')
            case 'aspect':
                nan_mask = np.isnan(x)
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for aspect')
            case 'rugosity':
                nan_mask = np.isnan(x)
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for rugosity')
            case 'chm':
                nan_mask = np.isnan(x)
                less_than_mask = (x < 0.0)
                greater_than_mask = (x > 45)
                x_mask = nan_mask | less_than_mask | greater_than_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for chm')
            case _:
                raise TypeError('inappropriate input_type for mask_X()')
        return x_mask

    def mask_y(self, y, truth_type='chm', no_data_fill_value=None):
        match truth_type:
            case 'chm':
                nan_mask = np.isnan(y)
                less_than_mask = (y < 0.5)
                greater_than_mask = (y > 45)
                y_mask = nan_mask | less_than_mask | greater_than_mask
            case 'chc':
                nan_mask = np.isnan(y)
                less_than_mask = (y < 0.5)
                greater_than_mask = (y > 45)
                y_mask = nan_mask | less_than_mask | greater_than_mask
            case 'ndvi':
                nan_mask = np.isnan(y)
                y_mask = nan_mask
        return y_mask

    def _merge_masks(self, X, y, X_masks, y_mask):
        accumulated_x_mask = np.zeros(X_masks[0].shape[1:3], dtype=bool)
        for x_mask in X_masks:
            x_mask = np.any(x_mask, axis=(0, 3))
            accumulated_x_mask = accumulated_x_mask | x_mask
        X_mask = accumulated_x_mask

        if np.isscalar(y_mask):
            print("Warning: y_mask is scalar, expanding to match X_mask shape")
            y_mask = np.full(y.shape[1:3], False, dtype=bool)

        master_mask = np.ma.mask_or(X_mask, y_mask, shrink=False)

        if np.isscalar(master_mask):
            print("Warning: master_mask is scalar, expanding to match input shape")
            master_mask = np.full(X_mask.shape, master_mask, dtype=bool)

        self.master_mask = master_mask
        
        y = np.ma.masked_array(y, mask=np.broadcast_to(master_mask, y.shape))

        X_masks = [np.broadcast_to(master_mask[np.newaxis, :, :, np.newaxis], x.shape) for x in X]
        y_mask = np.broadcast_to(master_mask, y.shape)

        print(f'Count of masked pixels based on input X filtering: {np.count_nonzero(X_mask)}')
        print(f'Count of masked y based on y filtering: {np.count_nonzero(y_mask)}')
        print(f'Count of total masked: {np.count_nonzero(master_mask)}')

        return X_masks, y_mask
    
    def _apply_masks(self, X, y, X_masks, y_mask):
        y = np.ma.masked_array(y, mask=y_mask)
        X = [np.ma.masked_array(x, mask=x_mask) for x, x_mask in zip(X, X_masks)]
        
        print(f'Count of masked pixels based on input X filtering: {np.count_nonzero(X_masks[0])}')
        print(f'Count of masked y based on y filtering: {np.count_nonzero(y_mask)}')
        print(f'Count of total masked: {np.count_nonzero(self.master_mask)}')
    
        X = [x.compressed().reshape(-1, x.shape[-1]) for x in X]
        y = y.compressed()
        
        [print(f'INFO: Resulting shapes of X and y after masks applied: {X[i].shape} , {y.shape}') 
         for i in range(len(X))]
        return X, y

    def __getitem__(self, index):
        x, y = self.__get_pixel_data(index)
        return tuple(x), y

    def on_epoch_end(self):
        seed = secrets.randbits(69420)
        if self.randomize_pixels:
            rng = np.random.default_rng(seed)
            [rng.shuffle(x, axis=0) for x in self.X]
            rng = np.random.default_rng(seed)
            rng.shuffle(self.y)

    def get_truth_vals(self):
        return self.y

    def get_training_data(self):
        return self.X
    
    def calc_upscale(self, data, method, is_truth=False):
        axis_dim = (3,4) if not is_truth else (2,3)
        if method == "median":
            return np.ma.median(data, axis=axis_dim)
        elif method == "average":
            return np.ma.average(data, axis=axis_dim)
        elif method == "std":
            return np.ma.std(data, axis=axis_dim)
        else:
            return None
    
    def upscale(self, data, mask, upscale_params, is_truth=False):
        print(f"INFO: original shape of {'x' if not is_truth else 'y'}:")
        print(data.shape)
        print(f"INFO: original mask shape of {'x' if not is_truth else 'y'}:")
        print(mask.shape)

        method = upscale_params.get('calc_method', 'median')
        size_w = upscale_params.get('window_size', 10)
        pct_good = upscale_params.get('pct_good', 0.5)

        if is_truth:
            data = data.reshape(-1, 1000, 1000)
            mask = mask.reshape(-1, 1000, 1000)
        else:
            data = data.reshape(-1, 1000, 1000, data.shape[-1])
            mask = mask.reshape(-1, 1000, 1000, mask.shape[-1])
        
        num_tiles = data.shape[0]
        num_bands = data.shape[-1] if not is_truth else 1
        
        print(f'INFO: num_tiles is: {num_tiles}')
        if not is_truth:
            print(f'INFO: num_bands is: {num_bands}')
        
        print(f"INFO: modified shape of {'x' if not is_truth else 'y'}:")
        print(data.shape)
        print(f"INFO: modified mask shape of {'x' if not is_truth else 'y'}:")
        print(mask.shape)
        
        axis_dim = (3,4) if not is_truth else (2,3)
        cutoff = int(size_w*size_w*pct_good)

        if is_truth:
            windows = np.lib.stride_tricks.sliding_window_view(data, (1, size_w, size_w))[:, ::size_w, ::size_w, 0]
            windows_mask = np.lib.stride_tricks.sliding_window_view(mask, (1, size_w, size_w))[:, ::size_w, ::size_w, 0]
            windows_masked = np.ma.masked_array(windows, mask=~windows_mask)
            out_mask = np.sum(windows_mask, axis=axis_dim) < cutoff
        else:
            windows = np.lib.stride_tricks.sliding_window_view(data, (1, size_w, size_w, num_bands))[:, ::size_w, ::size_w, 0, 0]
            windows_mask = np.lib.stride_tricks.sliding_window_view(mask, (1, size_w, size_w, num_bands))[:, ::size_w, ::size_w, 0, 0]
            windows_masked = np.ma.masked_array(windows, mask=~windows_mask)
            out_mask = np.sum(windows_mask, axis=axis_dim) < cutoff

        outdata = self.calc_upscale(windows_masked, method, is_truth=is_truth)
        print(f'windows_masked.shape: {windows_masked.shape}')
        print(f'outdata.shape BEFORE reshaping: {outdata.shape}')
        print(f'out_mask.shape BEFORE reshaping: {out_mask.shape}')

        if is_truth:
            outdata = outdata.reshape(-1, 1)
            out_mask = out_mask.reshape(-1, 1)
        else:
            outdata = outdata.reshape(-1, num_bands)
            out_mask = out_mask.reshape(-1, num_bands)

        print(f'outdata.shape AFTER: {outdata.shape}')
        print(f'out_mask.shape AFTER: {out_mask.shape}')

        return outdata, out_mask

    def __len__(self):
        return int(np.ceil(self.y.shape[0] / self.batch_size))

    def __get_pixel_data(self, index):
        start = index * self.batch_size
        end = (index + 1) * self.batch_size
        item_data = [x[start:end] for x in self.X]
        item_truth = self.y[start:end]
        return item_data, item_truth

    def get_file_paths_and_bands(self, branch=None, truth=False):
        if not truth:
            features = branch['features']
        elif truth == "chm" or truth == "chc":
            features = {'canopy_height': {'file_ending':'_CHM.tif', 'bands': None}}
        elif truth == "narrow_ndvi":
            features = {'ndvi': {'file_ending':'_reflectance_excluded_bands.tif', 'bands': [53,87]}}
        elif truth == "wide_ndvi":
                features = {'ndvi': {'file_ending':'_reflectance_indices.tif', 'bands': [14]}}
        else:
            raise TypeError("Truth type not recognized for get_file_paths_and_bands() function")

        feature_file_paths_and_bands = []
        for tile_directory in self.file_paths:
            tile_path = os.path.join(self.data_path, tile_directory)
            if os.path.isdir(tile_path):
                feature_file = {}
                for feature_name, feature_info in features.items():
                    for file in os.listdir(tile_path):
                        if file.endswith(feature_info['file_ending']):
                            feature_file[feature_name] = {
                                'file_path': os.path.join(tile_path, file),
                                'bands': feature_info.get('bands', None)
                            }
                feature_file_paths_and_bands.append(feature_file)
        return feature_file_paths_and_bands
    
    def get_mean_std_min_max(self):
        infos = []
        for branch_num, branch in enumerate(self.branch_inputs):
            branch_name = branch['branch_name']
            X = [self.scalers[branch_num].inverse_transform(x.reshape(-1, x.shape[-1])).reshape(x.shape) for branch_num, x in enumerate(self.X)] if self.scale_data_method else self.X
            mean = np.mean(X[branch_num])
            std = np.std(X[branch_num])
            min = np.min(X[branch_num])
            max = np.max(X[branch_num])
            infos.append({f"mean_{branch_name}": mean, f"std_{branch_name}": std, f"min_{branch_name}": min, f"max_{branch_name}": max})
        return infos
