import os
import numpy as np
# from numpy.lib.stride_tricks import sliding_window_view as swv
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

# TODO: class description

class DataGenerator(Sequence):

    # data_path = root path where data files and truth files are located
    # batch_size = # of patches in batch, OR *number of pixels* if randomize_pixels = True
    # randomize_pixels -- set to True if desired to randomize all pixels in all files after each epoch.
    # split_tuple = (train, test, validate). (0, 0, 0) by default for no split, use (80, 10, 10) for usual split
    # split_mode = 'train' by default.  Use 'validate', or 'test' if desired

    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,
                 label_scaler=None, **kwargs): # TODO: scaling y, delete label_scaler if not useful
        
        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 # cat or reg, default is reg
        self.bins = bins

        # canopy height specific
        # TODO: right now CHC depends greatly on the fact that CHM files are named the exact
        # same and have the same directory structure, with the years as the only difference. 
        # this might need to change in the future, but works for now
        self.year1 = year1
        self.year2 = year2

        self.scalers = scalers

        # TODO: scaling y, delete if not useful
        self.scale_y_labels = kwargs.get('scale_y_labels', False)
        self.label_scaler = label_scaler

        self.X, self.y = 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):
        """
        Main data generator flow.
        generate y, generate X, masks for each 
        --> apply masks 
        --> split into train, val, test sets 
        --> scale X

        Parameters
        ----------
        None

        Returns
        -------
        X: List of np matrices each with shape (num_samples, num_features)
            input data
        y: 1D numpy array of shape (num_samples,)
            truth
        """

        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)
            # debugging
            print(f'new upscaled y and y_masks shape: {y.shape} and {y_masks.shape}')

            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)
            # debugging
            print(f'new upscaled X and X_masks shape:')
            [print(x.shape) for x in X]
            [print(x_mask.shape) for x_mask in X_masks]

            print('** Merging new upscaled masks into new master_mask **')
            X_masks, y_masks = self._merge_masks(X, y, X_masks, y_masks)
            # debugging
            print(f'new upscaled X_masks and y_masks post remerged masks shape:')
            [print(x_mask.shape) for x_mask in X_masks]
            print(f'y_masks.shape: {y_masks.shape}')

        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:
            # self.scalers = [] # saving scalers per branch to list, to allow for unscaling
            print('** Scaling input data **')
            # debugging
            print('data shapes: ')
            [print(x.shape) for x in X]
            # X = [self.scale(x) for x in X]
            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])}')
            # trying out scaling y
            # y = self.scale(y.reshape(-1, 1))

        if self.scale_y_labels:
            if self.split_mode == 'train':
                if self.label_scaler is None:
                    all_train_y = y.reshape(-1, 1)
                    self.label_scaler = StandardScaler().fit(all_train_y)
                y = self.label_scaler.transform(y.reshape(-1, 1)).reshape(y.shape)
            else:
                if self.label_scaler is None:
                    raise ValueError(
                        "For split_mode!='train', you MUST pass label_scaler=train_generator.label_scaler"
                    )
                y = self.label_scaler.transform(y.reshape(-1, 1)).reshape(y.shape)

        return X, y

    ######################################################################################################################
    ################# TRUTH (y) PREPROCESSING ############################################################################
    ######################################################################################################################

    # Calculate NDVI truth -- either categorical or regression.  Truth values will be masked by the NDVI data layer mask in
    # apply_filters().
    # TODO: refactor to apply to new data generator paradigm
    def _generate_NDVI_truth(self):
        """
        generates y for the case of using NDVI as truth
        gets the file paths for specific truth
        if using narrow ndvi calculation, gets bands 87 and 53 from HSI image (343 band version)
        if using wide ndvi, gets band 14 from indices file

        Parameters
        ----------
        None

        Returns
        -------
        y: 1D numpy array of shape (num_samples,)
            truth
        y_masks: 1D numpy array of shape (num_samples,)
            mask for truth
        """

        y = []  # use for case of more than 1 branch -- then need to add branch data to list
        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)
                    truth_flattened = truth_array.flatten()
                else:
                    truth_array = src.read(ndvi_truth_file['ndvi']['bands'])  # read all raster values into 2D array
                    truth_flattened = truth_array.flatten()
            truth_mask = self.mask_y(truth_flattened, 'ndvi')
            truth_flattened = truth_flattened.flatten()
            y.append(truth_flattened)
            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

    # Calculate truth values for entire file at once for randomized pixel case
    # and return a flattened 1D array
    def _generate_canopy_truth(self):
        """
        generates y for the case of using canopy height model (CHM) as truth
        gets the file paths for specific truth

        Parameters
        ----------
        None

        Returns
        -------
        y: 1D numpy array of shape (num_samples,)
            truth
        y_masks: 1D numpy array of shape (num_samples,)
            mask for truth
        """

        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).flatten()
            if self.truth == "chc":
                # DTM = Digital Terrain Model elevation (ground interpolated from lidar points classified as ground)
                # DSM = Digital Surface Model elevation (highest surface; e.g. canopy top, building top, bare ground, road) 
                # CHM = Canopy Height Model (DSM – DTM)
                # Truth Label
                # CHMC(A to B)/yr = CHM change from Year A to Year B per year = 
                # {[DSM(B) – DTM(B)] – [DSM(A) – DTM(A)]}/(B – A)

                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).flatten()
                # chc = (ch_year_2 - ch_year_1) / (self.year2 - self.year1)
                # reverting changes
                chc = (ch_year_2 - ch_year_1) / (year2 - year1)
                y.append(chc)

                # debugging
                # print(f'CHC variables: ')
                # print(f'{ch_year_1= }')
                # print(f'{canopy_truth_comparison_file= }')
                # print(f'{ch_year_2= }')
                # print(f'{chc= }')

                y_masks.append(self.mask_y(chc, 'chc', no_data_fill_value=-9999)) # hardcoding for now
                # trying masking chm still
                # y_masks.append(self.mask_y(canopy_height, 'chc'))

            elif self.truth == "chc_maybe":
                # DTM = Digital Terrain Model elevation (ground interpolated from lidar points classified as ground)
                # DSM = Digital Surface Model elevation (highest surface; e.g. canopy top, building top, bare ground, road) 
                # CHM = Canopy Height Model (DSM – DTM)
                # Truth Label
                # CHMC(A to B)/yr = CHM change from Year A to Year B per year = 
                # {[DSM(B) – DTM(B)] – [DSM(A) – DTM(A)]}/(B – A)

                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).flatten()
                # chc = (ch_year_2 - ch_year_1) / (self.year2 - self.year1)
                # reverting changes
                chc = (ch_year_2 - ch_year_1) / (year2 - year1)
                y.append(chc)

                # debugging
                # print(f'CHC variables: ')
                # print(f'{ch_year_1= }')
                # print(f'{canopy_truth_comparison_file= }')
                # print(f'{ch_year_2= }')
                # print(f'{chc= }')

                y_masks.append(self.mask_y(chc, 'chc', no_data_fill_value=-9999)) # hardcoding for now
                # trying masking chm still
                # y_masks.append(self.mask_y(canopy_height, 'chc'))
            else:
                y.append(canopy_height)
                y_masks.append(self.mask_y(canopy_height, 'chm'))
        y = np.concatenate(y)
        y_masks = np.concatenate(y_masks)
        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}')

        # scaling y up to see if that helps
        # print(f'Scaling y * 1000 to see if that helps')
        # y = y * 1000
        return y, y_masks

    ######################################################################################################################
    ################# INPUT (X) PREPROCESSING ############################################################################
    ######################################################################################################################

    def _generate_X(self, branch):
        """
        Generates list of X, with each list item being a branch's inputs of shape (num_samples, num_features)
        Gets the file paths and bands for specific inputs using get_file_paths_and_bands().
        Loops over each branch, then loops over each tile, then for each feature in tile, loads in the features from proper tile.
        Concatenates each feature column-wise, then concatenates each tile row-wise. Does this for each branch.

        Parameters
        ----------
        None

        Returns
        -------
        X: List of np matrices, with each list item being a branch input of shape (num_samples, num_features)
        X_masks: List of np matrices, with each list item being a mask for branch input of shape (num_samples, num_features)
        """

        # tile_features: list of dictionaries - list is tiles, dictionary is features
        # ex: [tile1, tile2] ; tile = { 'feature_name': {'file_path': path/to/file, 'bands': [bands]} }
        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 = hyperspectral.reshape(-1, hyperspectral.shape[-1])
                    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 = savgol.reshape(-1, savgol.shape[-1])
                    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.reshape(-1, 1)
                    brdf_mask = self.mask_X(brdf, 'brdf', no_data_fill_value=src.nodata)
                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.reshape(-1, 1)
                    dtm_mask = self.mask_X(dtm, 'dtm', no_data_fill_value=src.nodata)
                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.reshape(-1, 1)
                    dsm_mask = self.mask_X(dsm, 'dsm', no_data_fill_value=src.nodata)
                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.reshape(-1, 1)
                    slope_mask = self.mask_X(slope, 'slope', no_data_fill_value=src.nodata)
                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.reshape(-1, 1)
                    aspect_mask = self.mask_X(aspect, 'aspect', no_data_fill_value=src.nodata)
                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.reshape(-1, 1)
                    rugosity_mask = self.mask_X(rugosity, 'rugosity', no_data_fill_value=src.nodata)
                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.reshape(-1, 1)
                    chm_mask = self.mask_X(chm, 'chm', no_data_fill_value=src.nodata)
                x.append(chm)
                x_mask.append(chm_mask)


            # hotfix, adding if x
            # TODO: figure out why sometimes x becomes empty
            if len(x) > 0:
                # concatenating tile features column-wise
                x = np.hstack(x)
                x_mask = np.hstack(x_mask)
                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]
        # concatenating tiles row-wise
        X = np.vstack(X)
        X_mask = np.vstack(X_mask)

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

    # TODO: I don't think we need to scale here...the scaler at the end does per feature/band scaling, so it would do it per indice anyways
    def generate_indices(self, src, tile_feature):
        # mcti, ndii, ndli, ndni, ndvi, ndwi, prin, priw = src.read([12,14,15,16,17,18,20,21])
        # removing mcti for now since causing trouble
        # ndii, ndli, ndni, ndvi, ndwi, prin, priw = src.read([14,15,16,17,18,20,21])
        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:
            match band:
                # ndvi
                case 17:
                    ndvi = src.read(band)
                    # ndvi = reshape_as_image(ndvi)
                    print(f'ndvi shape: {ndvi.shape}')
                    ndvi = ndvi.reshape(-1, 1)
                    ndvi_mask = self.mask_indices(ndvi, 'ndvi')
                    indices.append(ndvi)
                    indices_mask.append(ndvi_mask)
                    print(f'ndvi_mask shape: {ndvi_mask.shape}')
                # evi
                case 8:
                    evi = src.read(band)
                    print(f'evi shape: {evi.shape}')
                    evi = evi.reshape(-1, 1)
                    evi_mask = self.mask_indices(evi, 'evi')
                    indices.append(evi)
                    indices_mask.append(evi_mask)
                    print(f'evi_mask shape: {evi_mask.shape}')
                # evi2
                case 9:
                    evi2 = src.read(band)
                    print(f'evi2 shape: {evi2.shape}')
                    evi2 = evi2.reshape(-1, 1)
                    evi2_mask = self.mask_indices(evi2, 'evi2')
                    indices.append(evi2)
                    indices_mask.append(evi2_mask)
                    print(f'evi2_mask shape: {evi2_mask.shape}')
                # savi
                case 22:
                    savi = src.read(band)
                    print(f'savi shape: {savi.shape}')
                    savi = savi.reshape(-1, 1)
                    savi_mask = self.mask_indices(savi, 'savi')
                    indices.append(savi)
                    indices_mask.append(savi_mask)
                    print(f'savi_mask shape: {savi_mask.shape}')
                # all other indices
                case _:
                    indice = src.read(band)
                    print(f'indice shape: {indice.shape}')
                    indice = indice.reshape(-1, 1)
                    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.hstack(indices)
        indices_mask = np.hstack(indices_mask)

        # debugging
        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
                x_mask = nan_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for ndvi specifically')
            # TODO: need to figure out what these masks should be
            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)
                # inf_mask = np.isinf(x)
                # less_than_mask = (x < -1.0)
                # greater_than_mask = (x > 1.0)
                # TODO: hotfix, adding NDVI here. need to split up all indices now,
                # because now ALL indices are masked < 0.5 as well, which is not the most ideal
                # x_mask = nan_mask | less_than_mask | greater_than_mask | inf_mask
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask
                # debugging
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for indices')
            case _:
                raise TypeError('inappropriate input_type for mask_indices()')

        return x_mask # TODO: rename to indices_mask for clarity
        




    ######################################################################################################################
    ################# SCALING ############################################################################################
    ######################################################################################################################

    def split_data(self, data):
        """
        splits data into train, validation, and test sets based on ratio specified in split_tuple in config file

        Parameters
        ----------
        data: the X or y to split (note: should only pass in one branch at a time)

        Returns
        -------
        data: input or truth data, either 1 branches inputs, or whole truth, that has been split into its proper data set 
        """

        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)

    ######################################################################################################################
    ################# MASKING ############################################################################################
    ######################################################################################################################

    def mask_X(self, x, input_type, no_data_fill_value=None):
        """
        Generates mask for specific input feature

        Parameters
        ----------
        x: the input data to mask
        input_type: the type of feature we are masking here. dictates how to mask this data
        no_data_fill_value (optional): the NaN value for this specific tile

        Returns
        -------
        x_mask: the mask for this specific input feature
        """

        match input_type:
            case 'hyperspectral':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                x_mask = nan_mask | fill_mask
                # reverting changes
                # less_than_mask = x < 0
                # x_mask = nan_mask | fill_mask | less_than_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
                # reverting changes
                # less_than_mask = x < -1000
                # greater_than_mask = x > 1000
                # x_mask = nan_mask | fill_mask | less_than_mask | greater_than_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for savgol')
            case 'brdf':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # less_than_mask = (x < 0.0)
                # greater_than_mask = (x > 180)
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for brdf')
            case 'ndvi':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # ndvi_mask = (x < 0.5)
                # x_mask = nan_mask | ndvi_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for ndvi specifically')
            case 'dtm':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # less_than_mask = (x < 0.0)
                # greater_than_mask = (x > 3000)
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for dtm')
            case 'dsm':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # less_than_mask = (x < 0.0)
                # greater_than_mask = (x > 3000)
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for dtm')
            case 'slope':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # less_than_mask = (x < 0.0)
                # greater_than_mask = (x > 30)
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for slope')
            case 'aspect':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # less_than_mask = (x < 0.0)
                # greater_than_mask = (x > 360)
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for aspect')
            case 'rugosity':
                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 rugosity')
            case 'chm':
                nan_mask = np.isnan(x)
                fill_mask = (x == no_data_fill_value)
                # less_than_mask = (x < 0.0)
                # greater_than_mask = (x > 45)
                # x_mask = nan_mask | less_than_mask | greater_than_mask
                x_mask = nan_mask | fill_mask
                print(f'INFO: Masked out {np.count_nonzero(x_mask)} for chm')
            case _:
                raise TypeError('inappropriate input_type for mask_canopy_X()')
        return x_mask
        
    def mask_y(self, y, truth_type='chm', no_data_fill_value=None):
        """
        Generates mask for truth

        Parameters
        ----------
        y: the truth to mask
        truth_type: the type of truth we are masking here. dictates how to mask
        no_data_fill_value (optional): the NaN value for this specific tile

        Returns
        -------
        y_mask: the mask for truth
        """

        match truth_type:
            case 'chm':
                # TODO: make masks part of config file
                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
                # y_mask = nan_mask | greater_than_mask
            case 'chc':
                # TODO: masking +/- 20 meters to see how it performs
                nan_mask = np.isnan(y)
                fill_mask = (y == no_data_fill_value)
                # less_than_mask = (y < 0.5)
                # greater_than_mask = (y > 45)
                # less_than_mask = (y < 20)
                # greater_than_mask = (y > 20)
                # y_mask = nan_mask | less_than_mask | greater_than_mask
                
                # for now, just masking NaNs until I get masking criteria from Dave
                y_mask = nan_mask | fill_mask
            case 'ndvi':
                nan_mask = np.isnan(y)
                # print(f'Count of unmasked pixels for NDVI NaNs = {np.count_nonzero(nan_mask)}')
                # less_than_mask = (y < 0.85)
                # print(f'Count of unmasked pixels for NDVI < 0.85 = {np.count_nonzero(less_than_mask)}')
                # greater_than_mask = (y > 0.95)
                # print(f'Count of unmasked pixels for NDVI > 0.95 masked = {np.count_nonzero(greater_than_mask)}')
                # y_mask = nan_mask | less_than_mask | greater_than_mask
                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[0], dtype=bool)
        for x_mask in X_masks:
            x_mask = np.any(x_mask, axis=1)
            accumulated_x_mask = accumulated_x_mask | x_mask
        X_mask = accumulated_x_mask

        # ensure y_mask is not a scalar
        if np.isscalar(y_mask):
            print("Warning: y_mask is scalar, expanding to match X_mask shape")
            y_mask = np.full(y.shape, False, dtype=bool)

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

        # add check for scalar master_mask
        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)

        # save master_mask as attribute for plotting purposes
        self.master_mask = master_mask
        
        y = np.ma.masked_where(master_mask, y)

        X_masks = [np.broadcast_to(master_mask.reshape(-1,1), np.shape(x)) for x in X]
        y_mask = master_mask

        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)}')

        # debugging
        print('from merge masks...maybe where the issue is?')
        [print(f'X_masks shapes: {x_mask.shape}') for x_mask in X_masks]
        print(f'y_mask.shape: {y_mask.shape}')

        return X_masks, y_mask
    
    def _apply_masks(self, X, y, X_masks, y_mask):
        """
        Cycle through all X and y.  Apply all masks to all layers.
        After this, "compress()" all layers to yield an even rectangular data stack and y stack,
        Create a single mask for all pixel data for all branches and y data, each, then master_mask that contains all masked

        Parameters
        ----------
        X: the list of all branches and input data per branch
        y: all truth
        X_masks: the masks for all branches
        y_mask: the mask for truth

        Returns
        -------
        X: List of masked and compressed input data for all branches
        y: masked truth data
        """
        
        # accumulated_x_mask = np.zeros_like(X_masks[0].shape[0], dtype=bool)
        # for x_mask in X_masks:
        #     x_mask = np.any(x_mask, axis=1)
        #     accumulated_x_mask = accumulated_x_mask | x_mask
        #     # print(f'masked pixels in this branch: {np.count_nonzero(x_mask)}')
        # X_mask = accumulated_x_mask

        # master_mask = np.ma.mask_or(X_mask, y_mask)
        # # save master_mask as attribute for plotting purposes
        # self.master_mask = master_mask

        # add check for scalar master_mask
        if np.isscalar(self.master_mask):
            print("Warning: self.master_mask is scalar, expanding to match input shape")
            self.master_mask = np.full(y.shape, self.master_mask, dtype=bool)

        y = np.ma.masked_where(self.master_mask, y)

        X = [np.ma.masked_array(x, mask=np.broadcast_to(self.master_mask.reshape(-1,1), np.shape(x))) for x in X]
        # X = [np.ma.masked_array(x, mask=x_mask) for x, x_mask in zip(X, X_masks)]

        # debugging
        print(f'debugging here')
        [print(f'X_masks shapes: {x_mask.shape}') for x_mask in X_masks]
        [print(f'X shapes: {x.shape}') for x in X]
        print(f'self.master_mask.shape: {self.master_mask.shape}')
        print(f'y.shape: {y.shape}')
        print(f'y_mask.shape: {y_mask.shape}')

        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 = [np.ma.compress_rowcols(x, 0) for x in X]
        y = np.ma.compressed(y)
        
        [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
    
    ######################################################################################################################
    ################# BATCHING ###########################################################################################
    ######################################################################################################################

    
    def __getitem__(self, index):
        """
        Gets item of shape (batch_size, features) to send to model for training
        An "item" is a group (batch) of "patches" from a desired file.
        Batches are groups of patches of desired dimensions

        Parameters
        ----------
        index: the slice of data to send to model. Will be slice of batch_size*step

        Returns
        -------
        X: List of each sliced branch input
        y: sliced truth
        """

        # uncomment to go back to original if y scaling is not useful
        # x, y = self.__get_pixel_data(index)
        # return tuple(x), y

        x, y = self.__get_pixel_data(index)
        if self.scale_y_labels and self.label_scaler is not None:
            y_flat = y.reshape(-1, 1)
            y_scaled_flat = self.label_scaler.transform(y_flat)
            y = y_scaled_flat.reshape(y.shape)

        return tuple(x), y


    def on_epoch_end(self):
        """
        Stuff to do on each epoch end. For now, will shuffle X and y on feature/column axis if self.randomize_pixels is True

        Parameters
        ----------
        None

        Returns
        -------
        None. slices self.X and self.y in place
        """

        seed = secrets.randbits(69420)
        # Both arrays (the 2D data and 1D truth) are shuffled exactly the same way because
        # they are using the same seed -- set right before the shuffle call each time
        if self.randomize_pixels:
            rng = np.random.default_rng(seed)
            [rng.shuffle(x, axis=0) for x in self.X]  # shuffles the values in each row -- the same way
            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):
        """
        This actually does the computations for the upscaled values.
        It is structured like this to make it easier to work with from
        a multiprocessing standpoint

        Input Params:
        data : masked tiles (ndarray), 
        method : calculation method (str)
        """
        axis_dim = (3,4)
        if method == "median":
            # return np.ma.median(data, axis=(2, 3))
            return np.ma.median(data, axis=axis_dim)
        elif method == "average":
            # return np.ma.average(data, axis=(2, 3))
            return np.ma.average(data, axis=axis_dim)
        elif method == "std":
            # return np.ma.std(data, axis=(2, 3))
            return np.ma.std(data, axis=axis_dim)
        else:
            # TODO some error handling
            return
    
    def upscale(self, data, mask, upscale_params, is_truth=False):
        """
        Window and upscale NEON tiles

        Input Params:
        data (ndarray): Raw data from tiff, expects shape=(pixels, bands)
        mask (ndarray): Good data mask, expects shape=(pixels, bands)
        size_w (int): Window size 10x10 default
        pct_good (float): Percentage of good pixels required for good output
        is_truth (boolean): if passing in truth data, set to True

        Output Params:
        outdata (ndarray): Upscaled and masked array
        out_mask (ndarray): Output Mask, determined by pct_good param
        """

        # debugging
        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)

        # Reshape data to be 1000x1000xnum_bands for x, 1000x1000 for y
        if is_truth:
            data = data.reshape(-1, 1000, 1000) # (tiles, x, y)
            mask = mask.reshape(-1, 1000, 1000) # (tiles, x, y)
            num_tiles = data.shape[0]
        else:
            data = data.reshape(-1, 1000, 1000, data.shape[-1]) # (tiles, x, y, bands)
            mask = mask.reshape(-1, 1000, 1000, mask.shape[-1]) # (tiles, x, y, bands)
            num_tiles = data.shape[0]
            num_bands = data.shape[-1]

        # debugging
        print(f'INFO: num_tiles is: {num_tiles}')
        if not is_truth:
            print(f'INFO: num_bands is: {num_bands}')
        
        # debugging
        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)
        
        # goes over (window1, window2)
        axis_dim = (3,4)
        cutoff = int(size_w*size_w*pct_good)

        # Create views
        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]

            # debugging
            print(f'windows.shape: {windows.shape}')
            print(f'windows_mask.shape: {windows_mask.shape}')
            
            # do we need the tiling? the x_mask comes in as shape (pixels, bands), same as x
            # Generate tiled input mask and use it to mask input array
            # inmask = np.tile(windows_mask[...], (1, 1, 1, 1))
            # windows_masked = np.ma.masked_array(windows, mask=~inmask)
            windows_masked = np.ma.masked_array(windows, mask=~windows_mask)
            out_mask = np.sum(windows_mask, axis=axis_dim) < cutoff
        else:
            # debugging
            print(f'num_bands: {num_bands}')
            print(f'doing windows')
            windows = np.lib.stride_tricks.sliding_window_view(data, (1, size_w, size_w, num_bands))[:, ::size_w, ::size_w, 0, 0]
            print(f'windows.shape: {windows.shape}')
            # debugging
            print(f'doing windows_mask')
            windows_mask = np.lib.stride_tricks.sliding_window_view(mask, (1, size_w, size_w, num_bands))[:, ::size_w, ::size_w, 0, 0]
            print(f'windows_mask.shape: {windows_mask.shape}') # (4, 100, 100, 10, 10, 350) for x, (4, 100, 100, 10, 10) for y

            # Generate tiled input mask and use it to mask input array
            # inmask = np.tile(windows_mask[..., np.newaxis], (1, 1, 1, 1, num_bands))
            # windows_masked = np.ma.masked_array(windows, mask=~inmask)
            windows_masked = np.ma.masked_array(windows, mask=~windows_mask)
            # Generate upscaled output mask
            out_mask = np.sum(windows_mask, axis=axis_dim) < cutoff

        # Calculate upscaled values
        # with Pool(3) as p:
        #     outdata = p.map(calc_upscale, [(windows_masked, "average"),
        #                                 (windows_masked, "median"),
        #                                 (windows_masked, "std")])
        outdata = self.calc_upscale(windows_masked, method, is_truth=is_truth)
        # debugging
        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)
            out_mask = out_mask.reshape(-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 scale(self, X):
        if self.scalers is None:
            self.scalers = []
            for i, x in enumerate(X):
                match self.scale_data_method:
                    case 'minmax':
                        scaler = MinMaxScaler(feature_range=(-1,1))
                    case 'standard':
                        scaler = StandardScaler()
                    case 'robust':
                        scaler = RobustScaler()
                    case _:
                        warnings.warn(f'Scaling method not valid. Please check to make sure it is correct in config. \
                                    Defaulting to minmax scaler with feature_range=(-1,1) for now.')
                        scaler = MinMaxScaler(feature_range=(-1,1))
                X[i] = scaler.fit_transform(x)
                self.scalers.append(scaler)
        else:
            for i, (x, scaler) in enumerate(zip(X, self.scalers)):
                X[i] = scaler.transform(x)
        return X

    def __len__(self):
        """
        Denotes the number of batches per epoch
        Counts the number of possible batches (i.e. "items") that can be made from the total available data
        Number of patches % batch_size = 0, so every sample is seen

        Parameters
        ----------
        None

        Returns
        -------
        Number of steps per epoch
        """
        
        return int(np.ceil(self.y.shape[0] / self.batch_size))

    def __get_pixel_data(self, index):
        """
        Returns a batch of 1D pixel data for the desired index per band per file per branch
        Each 1D batch from each band or file is already vertically stacked into a 2D array.
        Truth is 1D array slice.  Same truth for each row in data stack
        We are using flattened 1D arrays for the data since it is randomized at the end of

        Parameters
        ----------
        index: int
            index of which pixel in array to return

        Returns
        -------
        item_data: 2D numpy array
            input data of batch_size for training
        item_truth: 1D numpy array
            label/ truth of input data
        """
       
        item_data = [x[index * self.batch_size:(index + 1) * self.batch_size, 0:] for x in self.X]
        item_truth = self.y[index * self.batch_size:(index + 1) * self.batch_size]
        return item_data, item_truth  # x slice (2D stack), y slice for the current item
    

    ######################################################################################################################
    ########################## UTILITY FUNCTIONS #########################################################################
    ######################################################################################################################

    # gets file paths and bands depending on the suffixes from config
    def get_file_paths_and_bands(self, branch=None, truth=False):
        """
        Gets the specific file paths and bands for each feature, or the truth.
        Takes the specific tiles being used from file_paths in config,
        then takes either the specific features and their file_endings (or the truth type)
        Then loops through the specific tile directories, and grabs the proper file paths
        from which to load in the data from based on given file_ending in branch. 
        Also takes the specific bands if specified, otherwise will use the default bands.
        Returns the list, with each list item representing a branch,
        and each list item being a dictionary of feature_name (or truth_name), file_path, and band

        Parameters
        ----------
        branch: List of dictionaries of each branch

        Returns
        -------
        feature_file_paths_and_bands: 
            List of dictionaries, 
            items in list representing tiles
            list item being a dictionary of {feature_name: {file_path: /path/to/file, bands: [bands]} }
            or {truth_name: {file_path: /path/to/file, bands: [bands]} }

            Ex:
            [tile1, tile2]
            tile1 = {
                'hyperspectral': {
                    'file_path': /path/to/hyperspectral_file.tif,
                    'bands': None
                },
                'brdf': {
                    'file_path': /path/to/brdf_file.tif,
                    'bands': [17]
                }
            }
        """
  
        if not truth:
            features = branch['features']
        elif truth == "chm" or truth == "chc" or truth == "chc_maybe":
            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):
        """
        Gets mean, standard deviation, min, and max of each branch of data for logging purposes.
        Loops over each branch, uses scaler to inverse_transform if scaling enabled, and saves info.

        Parameters
        ----------
        None

        Returns
        -------
        infos: List of dictionaries
            each list item representing branch of data, 
            each list item being dictionary of {'mean_branch_name': mean, 'std_branch_name': std, etc}
        """

        infos = []
        for branch_num, branch in enumerate(self.branch_inputs):
            branch_name = branch['branch_name']
            # TODO: see if we can just grab the unscaled data and apply this before preprocessing
            # since it's inefficient to scale just to unscale
            # unscale data to get proper info
            X = [self.scalers[branch_num].inverse_transform(x) 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
