import pandas as pd
import numpy as np
import glob, os
import xarray as xr
import pickle
import gc
import logging
from yaml import dump, Dumper
from skbio.diversity.alpha import shannon


from .analyze import area_of_applicability

logger = logging.getLogger("abil")

class AbilPostProcessor:
    """
    Post processing of SDM
    """
    def __init__(self, X_train, y_train, X_predict, model_config, statistic="mean", datatype=None):
        """
        A class for initializing and setting up a model with configuration, input data, and parameters.

        Parameters
        ----------
        X_train : {array-like, sparse matrix} of shape (n_samples, n_features)
            Training features used for model fitting.
        y : array-like of shape (n_samples,) or (n_samples, n_outputs)
            Target values used for model fitting.
        X_predict : {array-like, sparse matrix} of shape (n_samples, n_features)
            Features to predict on (e.g., environmental data).
        model_config : dict
            Dictionary containing model configuration parameters such as:
            - seed: int, random seed for reproducibility
            - path_out: str, output path for saving results
            - verbose: int, verbosity level (0-3)
            - cv: int, number of cross-validation folds
            - ensemble_config: dict, configuration for ensemble models
        pi : str
            The prediction interval identifier, defaulting to "50".
        datatype : str, optional
            The datatype of the predictions. This is used to access conversion factors (e.g. pic or poc) 
            and is appended to the file names during export            

        Attributes
        ----------
        path_out : str
            The output directory path where the model results will be saved.
        ds : xarray.Dataset
            The dataset containing the merged data from NetCDF files.
        traits : pd.DataFrame
            DataFrame containing the target trait information loaded from a CSV file.
        root : str
            Root directory as specified in the model configuration.
        d : pd.DataFrame
            DataFrame representation of the dataset after conversion and cleaning.
        targets : pd.Series
            The target values from the trait data that are present in the dataset columns.
        model_config : dict
            The model configuration dictionary containing paths, parameters, and other settings.
        pi : str
            The input parameter identifier, defaulting to "50".
        model_type : str
            The type of model being used, determined from the ensemble configuration (either "zir" or "reg").
        extension : str
            The file extension used for saving the model, based on the model type (e.g., "_zir.sav").
        datatype: str
            The datatype of the data being processed (e.g. "pg poc") which is appended to the data exports (optional)

        Methods
        -------
        merge_netcdf(path_in):
            Merges multiple NetCDF files from the specified directory into a single dataset.
        """

        def merge_netcdf(path_in, statistic):
            """
            Merges multiple NetCDF files from the specified directory into a single dataset.

            This function uses `xarray.open_mfdataset` to load all NetCDF files in the given directory 
            (matching the pattern "*.nc") and combines them into one xarray.Dataset. The function 
            logs status messages indicating the start and completion of the merging process.

            Parameters
            ----------
            path_in : str
                The path to the directory containing the NetCDF files to be merged.
            statistic : str
                The name of the statistic variable to extract from each dataset.

            Returns
            -------
            xarray.Dataset
                The merged dataset containing the combined data from all the NetCDF files in the directory.
                The variable names in the merged dataset are derived from the 'target' values in each file.
            """
            logger.info(f"merging... {path_in}")

            datasets = []
            
            for file in os.listdir(path_in):
                if file.endswith(".nc"):
                    ds = xr.open_dataset(os.path.join(path_in, file))
                    if statistic in ds:
                        # Extract the target name
                        target_name = ds['target'].values.item()  # Assuming target is a single value
                        
                        # Select the statistic and rename it to the target name
                        ds_subset = ds[[statistic]].rename({statistic: target_name})
                        datasets.append(ds_subset)
                    else:
                        logger.info(f"Statistic '{statistic}' not found in {file}")

            # Merge datasets by variables, keeping same coordinates
            merged_ds = xr.merge(datasets, compat='override')  # 'override' skips conflicts

            logger.info("finished merging NetCDF files")
            return merged_ds

        self.path_out = os.path.join(model_config['root'], model_config['path_out'], model_config['run_name'], "posts/")
        self.ds = merge_netcdf(os.path.join(model_config['root'], model_config['path_out'], model_config['run_name'], "predictions", "ens"), statistic)
        self.traits = pd.read_csv(os.path.join(model_config['root'], model_config['targets']))

        self.root  =  model_config['root'] 
        self.statistic = statistic

        self.d = self.ds.to_dataframe()
        self.unique_targets = np.unique(self.d.columns.values).tolist()

        self.d = self.d.dropna()
        self.targets = self.unique_targets

        self.model_config = model_config

        self.y_train = y_train
        self.X_train = X_train
        self.X_predict = X_predict
   
        # Export model_config to a YAML file
        self.export_model_config()
        if self.model_config['ensemble_config']['classifier'] and not self.model_config['ensemble_config']['regressor']:
            raise ValueError("classifiers are not supported")
        elif self.model_config['ensemble_config']['classifier'] and self.model_config['ensemble_config']['regressor']:
            self.model_type = "zir"
        if self.model_config['ensemble_config']['regressor'] and not self.model_config['ensemble_config']['classifier']:
            self.model_type = "reg"

        self.extension = "_" + self.model_type + ".sav"

        self.merge_parameters()
        self.merge_performance()

        if datatype:
            self.datatype = "_" + datatype
        else:
            self.datatype = ""
        
        
    def export_model_config(self):
        """
        Export the model_config dictionary to a YAML file in self.path_out.
        
        Raises
        ------
        Exception
            If an error occurs during the directory creation or file writing process, an exception
            is caught and an error message is logger.info(ed.

        Notes
        -----
        The YAML file is saved as "model_config.yml" in the `self.path_out` directory.
        """
        try:
            os.makedirs(self.path_out, exist_ok=True)  # Ensure the output directory exists
            yml_file_path = os.path.join(self.path_out, "model_config.yml")
            
            # Write the model_config dictionary to a YAML file
            with open(yml_file_path, 'w') as yml_file:
                dump(self.model_config, yml_file, Dumper=Dumper, default_flow_style=False)
            
            logger.info(f"Model configuration exported to: {yml_file_path}")
        except Exception as e:
            logger.info(f"Error exporting model_config to YAML: {e}")   

    def merge_performance(self):
        """
        Merges the performance data of multiple models as specified in the model configuration.

        Notes
        -----
        The function relies on the `merge_performance_single_model` method to merge individual model 
        performance data, and this is done for each model in the list, including the ensemble model.
        """    

        models = [value for key, value in self.model_config['ensemble_config'].items() if key.startswith("m")]
        logger.info(f"models included in merge performance! {models}")
        models.append("ens")
        for model in models:
            self.merge_performance_single_model(model)

       
    def merge_performance_single_model(self, model):
        """
        Merges performance metrics for a single model and saves the results to a CSV file.

        Parameters
        ----------
        model : str
            The name of the model for which performance metrics are being calculated and merged. 
            The model's performance data is expected to be stored in a `pickle` file in the "scoring" 
            directory under the model name and target name.

        Returns
        -------
        None

        Raises
        ------
        ValueError
            If the model configuration includes a classifier but not a regressor, an error is raised 
            since classifiers are not supported for performance merging.

        Notes
        -----
        The method calculates several performance metrics for each target column in the dataset:
            - R2: Coefficient of determination.
            - RMSE: Root Mean Squared Error.
            - MAE: Mean Absolute Error.
            - rRMSE: Relative Root Mean Squared Error.
            - rMAE: Relative Mean Absolute Error.

        The performance metrics for each target are aggregated into a DataFrame, which is then saved 
        as a CSV file in the "posts/performance" directory for the specified model.
        """
        
        all_performance = []

        for i in range(len(self.unique_targets)):
            
            target = self.unique_targets[i]
            target_no_space = target.replace(' ', '_')
            with open(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "scoring", model, target_no_space) + self.extension, 'rb') as file:

                m = pickle.load(file)
            
            if self.model_config['ensemble_config']['classifier'] and not self.model_config['ensemble_config']['regressor']:
                raise ValueError("classifiers are not supported")
            else:
                mean = np.mean(self.d[self.d.columns[i]])
                R2 = np.mean(m['test_R2'])
                RMSE = -1*np.mean(m['test_RMSE'])
                MAE = -1*np.mean(m['test_MAE'])
                rRMSE = -1*np.mean(m['test_RMSE'])/mean
                rMAE = -1*np.mean(m['test_MAE'])/mean            
                performance = pd.DataFrame({'target':[target], 'R2':[R2], 'RMSE':[RMSE], 'MAE':[MAE],
                                            'rRMSE':[rRMSE], 'rMAE':[rMAE]})
                all_performance.append(performance)

        all_performance = pd.concat(all_performance)
        #make new dir if needed
        os.makedirs(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "posts/performance"), exist_ok=True)
        all_performance.to_csv(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "posts/performance", model) + "_performance.csv", index=False)

        logger.info("finished merging performance")

    def merge_parameters(self):
        """
        Merges model parameters for multiple models as specified in the model configuration.

        Notes
        -----
        The method operates by iterating over each model in the ensemble configuration, collecting 
        model parameters using `merge_parameters_single_model`, and saving the results to a CSV file 
        in the "posts/parameters" directory.
        """

        models = [value for key, value in self.model_config['ensemble_config'].items() if key.startswith("m")]
        for model in models:
            self.merge_parameters_single_model(model)

    def merge_parameters_single_model(self, model):
        """
        Merges and saves model parameters for a single model.

        This method extracts the hyperparameters of a specified model (e.g., "rf", "xgb", "knn") from 
        serialized files stored as pickle objects. The method supports different model types, including 
        regression ("reg"), classification ("clf"), and ensemble ("zir") models. The extracted parameters 
        are stored in a DataFrame and then saved to a CSV file.

        The function also handles the creation of the necessary directories to save the resulting CSV file 
        if they do not already exist.

        Parameters
        ----------
        model : str
            The name of the model for which parameters are being merged. Expected models include 
            "rf" (Random Forest), "xgb" (XGBoost), and "knn" (K-Nearest Neighbors).

        Raises
        ------
        ValueError
            If the model configuration includes classifiers but not regressors, an error is raised 
            since classifiers are not supported for parameter merging.

        Notes
        -----
        The method processes the parameters of each model for regression and ensemble models, 
        extracting hyperparameters such as `n_estimators`, `max_depth`, `learning_rate`, and others.
        The parameters for each target are aggregated into a DataFrame and saved as a CSV file in the 
        "posts/parameters" directory.
        """
        
        all_parameters = []

        for i in range(len(self.unique_targets)):
            
            target = self.unique_targets[i]
            logger.info(f"the target is: {target}")
            target_no_space = target.replace(' ', '_')

            with open(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "model", model, target_no_space) + self.extension, 'rb') as file:

                m = pickle.load(file)

            if self.model_type == "reg":

                if model == "rf":
                    max_depth = m.regressor_.named_steps.estimator.max_depth
                    max_features = m.regressor_.named_steps.estimator.max_features
                    max_samples = m.regressor_.named_steps.estimator.max_samples
                    min_samples_leaf = m.regressor_.named_steps.estimator.min_samples_leaf
                    n_estimators = m.regressor_.named_steps.estimator.n_estimators
                    parameters = pd.DataFrame({'target':[target], 'n_estimators':[n_estimators], 'max_features':[max_features], 'max_depth':[max_depth], 
                                            'min_samples_leaf':[min_samples_leaf], 'max_samples':[max_samples]
                                            })
                    all_parameters.append(parameters)
                elif model == "xgb":
                    learning_rate = m.regressor_.named_steps.estimator.learning_rate
                    n_estimators = m.regressor_.named_steps.estimator.n_estimators
                    max_depth = m.regressor_.named_steps.estimator.max_depth
                    subsample = m.regressor_.named_steps.estimator.subsample
                    colsample_bytree = m.regressor_.named_steps.estimator.colsample_bytree
                    gamma = m.regressor_.named_steps.estimator.gamma
                    alpha = m.regressor_.named_steps.estimator.reg_alpha
                    parameters = pd.DataFrame({'target':[target], 'learning_rate':[learning_rate], 'n_estimators':[n_estimators], 
                                            'max_depth':[max_depth], 'subsample':[subsample], 'colsample_bytree':[colsample_bytree],
                                            'learning_rate':[learning_rate], 'gamma':[gamma], 'alpha':[alpha]                                           
                                            })
                    all_parameters.append(parameters)
                elif model == "knn":
                    max_samples = m.regressor_.named_steps.estimator.max_samples
                    max_features = m.regressor_.named_steps.estimator.max_features
                    leaf_size = m.regressor_.named_steps.estimator.estimator.leaf_size
                    n_neighbors = m.regressor_.named_steps.estimator.estimator.n_neighbors
                    p = m.regressor_.named_steps.estimator.estimator.p
                    weights = m.regressor_.named_steps.estimator.estimator.weights
                    parameters = pd.DataFrame({'target':[target], 'max_samples':[max_samples], 'max_features':[max_features],
                                            'leaf_size':[leaf_size], 'n_neighbors':[n_neighbors], 'p':[p], 'weights':[weights]
                                            })
                    all_parameters.append(parameters) 

            elif self.model_type == "clf":
                raise ValueError("classifiers are not supported")

            elif self.model_type == "zir":
                if model == "rf":
                    max_depth_reg = m.regressor_.regressor.named_steps.estimator.max_depth
                    max_features_reg = m.regressor_.regressor.named_steps.estimator.max_features
                    max_samples_reg = m.regressor_.regressor.named_steps.estimator.max_samples
                    min_samples_leaf_reg = m.regressor_.regressor.named_steps.estimator.min_samples_leaf
                    n_estimators_reg = m.regressor_.regressor.named_steps.estimator.n_estimators

                    n_estimators_clf = m.classifier.named_steps.estimator.n_estimators
                    max_features_clf = m.classifier.named_steps.estimator.max_features
                    max_depth_clf = m.classifier.named_steps.estimator.max_depth
                    min_samples_leaf_clf = m.classifier.named_steps.estimator.min_samples_leaf
                    max_samples_clf = m.classifier.named_steps.estimator.max_samples

                    parameters = pd.DataFrame({'target':[target], 'reg_n_estimators':[n_estimators_reg], 
                                            'reg_max_features':[max_features_reg], 'reg_max_depth':[max_depth_reg], 
                                            'reg_min_samples_leaf':[min_samples_leaf_reg], 'reg_max_samples':[max_samples_reg],
                                            'clf_n_estimators':[n_estimators_clf], 
                                            'clf_max_features':[max_features_clf], 'clf_max_depth':[max_depth_clf], 
                                            'clf_min_samples_leaf':[min_samples_leaf_clf], 'clf_max_samples':[max_samples_clf]
                                            })
                    all_parameters.append(parameters)

                elif model == "xgb":
                    learning_rate_reg = m.regressor_.regressor.named_steps.estimator.learning_rate
                    n_estimators_reg = m.regressor_.regressor.named_steps.estimator.n_estimators
                    max_depth_reg = m.regressor_.regressor.named_steps.estimator.max_depth
                    subsample_reg = m.regressor_.regressor.named_steps.estimator.subsample
                    colsample_bytree_reg = m.regressor_.regressor.named_steps.estimator.colsample_bytree
                    gamma_reg = m.regressor_.regressor.named_steps.estimator.gamma
                    alpha_reg = m.regressor_.regressor.named_steps.estimator.reg_alpha

                    learning_rate_clf = m.classifier.named_steps.estimator.learning_rate
                    n_estimators_clf = m.classifier.named_steps.estimator.n_estimators
                    max_depth_clf = m.classifier.named_steps.estimator.max_depth
                    subsample_clf = m.classifier.named_steps.estimator.subsample
                    colsample_bytree_clf = m.classifier.named_steps.estimator.colsample_bytree
                    gamma_clf = m.classifier.named_steps.estimator.gamma
                    alpha_clf = m.classifier.named_steps.estimator.reg_alpha


                    parameters = pd.DataFrame({'target':[target], 'reg_learning_rate':[learning_rate_reg], 'reg_n_estimators':[n_estimators_reg], 
                                            'reg_max_depth':[max_depth_reg], 'reg_subsample':[subsample_reg], 'reg_colsample_bytree':[colsample_bytree_reg],
                                            'reg_learning_rate':[learning_rate_reg], 'reg_gamma':[gamma_reg], 'reg_alpha':[alpha_reg],
                                            'clf_learning_rate':[learning_rate_clf], 'clf_n_estimators':[n_estimators_clf], 
                                            'clf_max_depth':[max_depth_clf], 'clf_subsample':[subsample_clf], 'clf_colsample_bytree':[colsample_bytree_clf],
                                            'clf_learning_rate':[learning_rate_clf], 'clf_gamma':[gamma_clf], 'clf_alpha':[alpha_clf]                                           
                                            })
                    all_parameters.append(parameters)

                elif model == "knn":
                    max_samples_reg = m.regressor_.regressor.named_steps.estimator.max_samples
                    max_features_reg = m.regressor_.regressor.named_steps.estimator.max_features
                    leaf_size_reg = m.regressor_.regressor.named_steps.estimator.estimator.leaf_size
                    n_neighbors_reg = m.regressor_.regressor.named_steps.estimator.estimator.n_neighbors
                    p_reg = m.regressor_.regressor.named_steps.estimator.estimator.p
                    weights_reg = m.regressor_.regressor.named_steps.estimator.estimator.weights

                    max_samples_clf = m.classifier.named_steps.estimator.max_samples
                    max_features_clf = m.classifier.named_steps.estimator.max_features
                    leaf_size_clf = m.classifier.named_steps.estimator.estimator.leaf_size
                    n_neighbors_clf = m.classifier.named_steps.estimator.estimator.n_neighbors
                    p_clf = m.classifier.named_steps.estimator.estimator.p
                    weights_clf = m.classifier.named_steps.estimator.estimator.weights



                    parameters = pd.DataFrame({'target':[target], 'reg_max_samples':[max_samples_reg], 'reg_max_features':[max_features_reg],
                                            'reg_leaf_size':[leaf_size_reg], 'reg_n_neighbors':[n_neighbors_reg], 
                                            'reg_p':[p_reg], 'reg_weights':[weights_reg],
                                            'clf_max_samples':[max_samples_clf], 'clf_max_features':[max_features_clf],
                                            'clf_leaf_size':[leaf_size_clf], 'clf_n_neighbors':[n_neighbors_clf], 
                                            'clf_p':[p_clf], 'clf_weights':[weights_clf]
                                            })
                    all_parameters.append(parameters) 

        all_parameters= pd.concat(all_parameters)
        #make new dir if needed
        os.makedirs(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "posts/parameters"), exist_ok=True)
        all_parameters.to_csv(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "posts/parameters", model) + "_parameters.csv", index=False)

        
        logger.info("finished merging parameters")

    def estimate_carbon(self, variable):

        """
        Estimate carbon content for each target based on a specified variable.

        This method calculates the carbon content for each target by scaling the data in `self.d` 
        with the values of the specified variable from the `traits` DataFrame. The results are 
        stored back in `self.d`.

        Parameters
        ----------
        variable : str
            The name of the column in the `traits` DataFrame containing the carbon content values 
            to be used for scaling the target data.
        """

        w = self.traits.query('Target in @self.targets')
        var = w[variable].to_numpy()
        logger.info(f"{variable} values:\n\t{var}")
        self.d = self.d.apply(lambda row : (row[self.targets]* var), axis = 1)
        logger.info(f"{variable} finished estimating.")

    def def_groups(self, dict):
        """
        Define groups of species based on a provided dictionary.

        Parameters
        ----------
        dict : dict
            A dictionary where keys represent group names, and values are lists of species or 
            column names to be grouped under each key.
            
        Notes
        -----
        - The method renames columns in `self.d` based on the provided dictionary and then sums 
        their values to create grouped columns.
        - The resulting grouped data is concatenated to the original `self.d`.
        """     

        df = self.d[self.targets]
        df = (df.rename(columns=dict)
            .groupby(level=0, axis=1, dropna=False)).sum( min_count=1)
        self.d = pd.concat([self.d, df], axis=1)
        logger.info("finished defining groups")

    def cwm(self, variable):
        """
        Calculate community weighted mean values for a given parameter. 

        Parameters
        ----------

        variable : string
            variable that is used to estimate cwm.

        """

        w = self.traits.query('Target in @self.targets')
        var = w[variable].to_numpy()
        var_name = 'cwm ' + variable
        self.d[var_name] = self.d.apply(lambda row : np.average(var, weights=row[self.targets]), axis = 1)
        logger.info(f"finished calculating CWM {variable}")

    def diversity(self):
        """
        Estimates Shannon diversity using scikit-bio.
        """
        self.d['shannon'] = self.d.apply(shannon, axis=1)
        logger.info("finished calculating shannon diversity")

    def total(self):
        """
        Sum target rows to estimate total.

        Notes
        ----------
        Useful for estimating total species abundances or varable sum if targets are continuous.
        Total is estimated based on the target list defined in model_config. 

        """

        self.d['total'] = self.d[self.targets].sum( axis='columns')
        self.d['total_log'] = np.log(self.d['total'])
        logger.info("finished calculating total")

    def process_resampled_runs(self):
        """
        Take mean of target rows.
        Take the standard deviation of the target rows.
        Calculate the 2.5th and 97.5th percentiles of target rows.

        Notes
        -----
        Useful when running resampled targets of the same initial target.
        Mean is estimated based on the target list defined in model_config.

        """

        self.d['mean'] = self.d[self.targets].mean(axis='columns')
        logger.info('finished calculating mean')
    
        self.d['stdev'] = self.d[self.targets].std(axis='columns')
        logger.info('finished calculating standard deviation')

        self.d['prctile_2.5'] = self.d[self.targets].quantile(0.025, axis='columns')
        self.d['prctile_97.5'] = self.d[self.targets].quantile(0.975, axis='columns')

        logger.info('finished calculating 2.5th and 97.5th percentiles')

    def integration(self, *args, **kwargs):
        return self.integration_class(self, *args, **kwargs)

    class integration:
        def __init__(self, parent, 
                     resolution_lat=1.0, resolution_lon=1.0, depth_w=5, 
                     vol_conversion=1, magnitude_conversion=1, molar_mass=1, rate=False):
            """
            Parameters
            ----------
            resolution_lat : float
                Latitude resolution in degrees, default is 1.0 degree.
            
            resolution_lon : float
                Longitude resolution in degrees, default is 1.0 degree.
            
            depth_w : float
                Bin depth in meters, default is 5m.

            vol_conversion : float
                Conversion to m^3, e.g., l to m^3 would be 1e3, default is 1 (no conversion).
            
            magnitude_conversion : float
                Prefix conversion, e.g., umol to Pmol would be 1e-21, default is 1 (no conversion).
            
            molar_mass : float
                Conversion from mol to grams, default is 1 (no conversion). Optional: 12.01 (carbon).
            
            rate : bool
                If input data is in rate per day, integrates over each month to provide an annual rate (yr^-1).            
            """


            self.parent = parent
            self.resolution_lat = resolution_lat
            self.resolution_lon = resolution_lon
            self.depth_w = depth_w
            self.vol_conversion = vol_conversion
            self.magnitude_conversion = magnitude_conversion
            self.molar_mass = molar_mass
            self.rate = rate
            self.calculate_volume()

        def calculate_volume(self):
            """
            Calculate the volume for each cell and add it as a new field to the dataset.

            Works with any subset of ('lat','lon','depth'). If only 'lat' is present,
            uses the zonal band area (integrated over all longitudes).

            Examples
            --------
            >>> m = AbilPostProcessor(model_config)
            >>> int = m.Integration(m, resolution_lat=1.0, resolution_lon=1.0, depth_w=5, vol_conversion=1, magnitude_conversion=1e-21, molar_mass=12.01, rate=True)
            >>> print("Volume calculated:", int.ds['volume'].values)
            """            
            ds = self.parent.d.to_xarray()
            resolution_lat = float(self.resolution_lat)
            resolution_lon = float(self.resolution_lon)
            depth_w = float(self.depth_w)

            earth_radius = 6371000.0  # Earth's radius in meters

            # area factor (lat/lon)
            has_lat = ('lat' in ds.dims or 'lat' in ds.coords)
            has_lon = ('lon' in ds.dims or 'lon' in ds.coords)

            if has_lat and has_lon:
                # lat & lon present: per-cell horizontal area using uniform resolutions
                lat = ds['lat'].values
                lon = ds['lon'].values

                dphi = np.deg2rad(resolution_lat)
                dlambda = np.deg2rad(resolution_lon)

                # vectorized latitude band contribution
                phi = np.deg2rad(lat)
                lat_band = (earth_radius ** 2) * (np.sin(phi + dphi / 2.0) - np.sin(phi - dphi / 2.0))  # (lat,)

                # longitude width (same for all longitudes for a given Δλ)
                lon_width = np.full(lon.shape, dlambda)  # (lon,)

                # outer product -> (lat, lon)
                area = np.outer(lat_band, lon_width)
                area_da = xr.DataArray(area, dims=('lat', 'lon'))

            elif has_lat and not has_lon:
                # lat only: use full 2π longitude span (zonal band area per latitude cell)
                lat = ds['lat'].values
                dphi = np.deg2rad(resolution_lat)
                phi = np.deg2rad(lat)
                # A_band = R^2 * (sin(phi+Δ/2) - sin(phi-Δ/2)) * 2π
                area_band = (earth_radius ** 2) * (np.sin(phi + dphi / 2.0) - np.sin(phi - dphi / 2.0)) * (2.0 * np.pi)
                area_da = xr.DataArray(area_band, dims=('lat',))
            else:
                # no horizontal dimensions -> multiply by 1
                area_da = xr.DataArray(1.0)

            # depth factor
            if 'depth' in ds.dims or 'depth' in ds.coords:
                nz = ds.sizes.get('depth', len(ds['depth']))
                dz_da = xr.DataArray(np.full((nz,), depth_w), dims=('depth',))
            else:
                dz_da = xr.DataArray(1.0)

            # estimate volume
            volume = area_da * dz_da
            volume.name = 'volume'
            ds['volume'] = volume
            self.parent.d = ds.to_dataframe()

        def integrate_total(self, variable='total', monthly=False, subset_depth=None):
            """
            Estimates global integrated values for a single target. Returns the depth integrated annual total.
            
            Works with any subset/ordering of ('lat','lon','depth','time').

            Parameters
            ----------
            variable : str
                The field to be integrated. Default is 'total' from PIC or POC Abil output.

            monthly : bool
                Whether or not to calculate a monthly average value instead of an annual total. Default is False.

            subset_depth : float
                Depth in meters from surface to which integral should be calculated. Default is None. Ex. 100 for top 100m integral.

            Examples
            --------
            >>> m = AbilPostProcessor(model_config)
            >>> int = m.Integration(m, resolution_lat=1.0, resolution_lon=1.0, depth_w=5, vol_conversion=1, magnitude_conversion=1e-21, molar_mass=12.01, rate=True)
            >>> result = integration.integrate_total(variable='Calcification')
            >>> print("Final integrated total:", result.values)
            """
            logger.info("Initiate integrated_total")
            ds = self.parent.d.to_xarray()
            vol_conversion = self.vol_conversion
            magnitude_conversion = self.magnitude_conversion
            molar_mass = self.molar_mass
            rate = self.rate

            # Average number of days for each month (accounting for leap years)
            days_per_month_full = np.array([31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])
            
            if variable not in ds:
                raise KeyError(f"Variable '{variable}' not found in dataset.")

            # Optional depth subset if depth exists
            if subset_depth is not None and ('depth' in ds.dims or 'depth' in ds.coords):
                ds = ds.sel(depth=slice(0, subset_depth))

            var = ds[variable]
            vol = ds['volume']

            # Ensure volume is broadcastable to var
            vol = vol.broadcast_like(var)

            has_time = ('time' in var.dims)

            # If time exists, map each time step to a month index 1..12
            if has_time:
                available_time = ds['time'].values
                if np.issubdtype(available_time.dtype, np.datetime64):
                    months_idx = pd.to_datetime(available_time).month.astype(int)
                else:
                    months_idx = available_time.astype(int)
                days_per_month = days_per_month_full[months_idx - 1]

            if rate and has_time:
                if monthly:
                    # Calculate monthly total (separately for each time step) with month-day weighting
                    total = []
                    for i in range(var.sizes['time']):
                        vol_i = vol.isel(time=i) if ('time' in vol.dims) else vol
                        monthly_total = (var.isel(time=i) * vol_i * days_per_month[i]).sum(
                            dim=[d for d in var.dims if d != 'time']
                        )
                        monthly_total = (monthly_total * molar_mass) * vol_conversion * magnitude_conversion
                        total.append(monthly_total)
                    total = xr.concat(total, dim="month")
                    logger.info(f"All monthly totals: {total.values}")
                else:
                    # Annual total with month-day weighting
                    weight = xr.DataArray(days_per_month, dims=('time',))
                    total = (var * vol * weight).sum(dim=list(var.dims))
                    total = (total * molar_mass) * vol_conversion * magnitude_conversion
                    logger.info(f"Final integrated total: {total.values}")
            else:
                if monthly and has_time:
                    # Monthly totals without rate weighting
                    total = []
                    for i in range(var.sizes['time']):
                        vol_i = vol.isel(time=i) if ('time' in vol.dims) else vol
                        monthly_total = (var.isel(time=i) * vol_i).sum(
                            dim=[d for d in var.dims if d != 'time']
                        )
                        monthly_total = (monthly_total * molar_mass) * vol_conversion * magnitude_conversion
                        total.append(monthly_total)
                    total = xr.concat(total, dim="month")
                    logger.info(f"All monthly totals: {total.values}")
                else:
                    # Integrate over whatever dims exist
                    total = (var * vol).sum(dim=list(var.dims))
                    total = (total * molar_mass) * vol_conversion * magnitude_conversion
                    logger.info(f"Final integrated total: {total.values}")
            return total

        def integrated_totals(self, targets=None, monthly=False, subset_depth=None, 
                         export=True, model="ens"):
            """
            Estimates global integrated values for all targets.

            Considers latitude and depth bin size.

            Parameters
            ----------
            targets : an np.array of str, optional
                An np.array of target variable names to include in the merge. If None, the default 
                targets from `self.targets` are used (default is None).

            monthly : bool
                Whether or not to calculate a monthly average value instead of an annual total. Default is False.

            subset_depth : float
                Depth in meters from surface to which integral should be calculated. Default is None. Ex. 100 for top 100m integral.

            export : bool
                Whether of not to export integrated totals as .csv. Default is True.

            model : str
                The model version to be integrated. Default is "ens". Other options include {"rf", "xgb", "knn"}.
            """
            ds = self.parent.d.to_xarray()
            if targets is None:
                targets = self.targets
            if "total" in ds:
                targets = np.append(targets, 'total')
            if "mean" in ds:
                targets = np.append(targets, 'mean')
            if "stdev" in ds:
                targets = np.append(targets, 'stdev')
            if "prctile_2.5" in ds:
                targets = np.append(targets, 'prctile_2.5')
            if "prctile_97.5" in ds:
                targets = np.append(targets, 'prctile_97.5')
            totals = []

            for target in targets:
                try:
                    logger.info(f"Processing target: {target}")
                    total = self.integrate_total(variable=target, monthly=monthly, subset_depth=subset_depth)
                    total_df = pd.DataFrame({'total': [total.values], 'variable': target})
                    totals.append(total_df)
                except Exception as e:
                    logger.info(f"Some targets do not have predictions! Missing: {target}")
                    logger.info(f"Error: {e}")
            totals = pd.concat(totals)

            if export:
                depth_str = f"_depth_{subset_depth}m" if subset_depth else ""
                month_str = "_monthly_int" if monthly else ""
                #make new dir if needed
                os.makedirs(os.path.join(self.parent.root, self.parent.model_config['path_out'], self.parent.model_config['run_name'], "posts/integrated_totals"), exist_ok=True)


                path_out = self.parent.model_config['path_out']
                run_name = self.parent.model_config['run_name']
                #pi = self.parent.pi
                statistic = self.parent.statistic
                datatype = self.parent.datatype

                # Build the full file path
                output_dir = os.path.join(self.parent.root, path_out, run_name, "posts/integrated_totals")
                filename = f"{model}_integrated_totals_{statistic}{depth_str}{month_str}{datatype}.csv"
                file_path = os.path.join(output_dir, filename)

                # Write to CSV
                totals.to_csv(file_path, index=False)

                logger.info(f"Exported totals")


    def estimate_applicability(self, targets=None, threshold='tukey', return_all=False, drop_zeros=False):
        """
        Estimate the area of applicability for the data using a strategy similar to Meyer & Pebesma 2022).
        This calculates the importance-weighted feature distances from test to train points,
        and then defines the "applicable" test sites as those closer than some threshold
        distance.
        A value of 0 indicates the point is within the Area of Applicability, 
        while a value of 1 indicates the point is outside the Area of Applicability.
        Note: if using pseudo-absences in y_train and  X_train, mask out where y_train = 0 to calculate
        the AOA for the original dataset.
        
        Parameters
        ----------
        targets : an np.array of str, optional
            An np.array of target variable names to include in the merge. If None, the default 
            targets from `self.targets` are used (default is None).
        drop_zeros : bool, optional
            Wether or not to exclude rows where y_train values are equal to 0.
            (default is False)
        """
        if targets is None:
            targets = self.targets

        # create empty dataframe with the same index as X_predict
        aoa_dataset = pd.DataFrame(index=self.X_predict.index)

        # estimate the aoa for each target:
        for i in range(len(targets)):

            target = targets[i]
            target_no_space = target.replace(' ', '_')

            # load the voting regressor model object for each target:
            with open(os.path.join(self.root, self.model_config['path_out'], self.model_config['run_name'], "model", "ens", target_no_space) + self.extension, 'rb') as file:
                m = pickle.load(file)

            if drop_zeros:
                if isinstance(self.y_train, pd.Series):
                    y_train = self.y_train.where(self.y_train > 0)
                else:
                    y_train = self.y_train.copy()
                    y_train.loc[y_train[target] <= 0, target] = np.nan
                    y_train = y_train[target]
            else:
                if isinstance(self.y_train, pd.Series):
                    y_train = self.y_train
                else:
                    y_train = self.y_train[target]

            if return_all is True:
                aoa, di_test, lpd_test, cutpoint, test_to_train_d = area_of_applicability(
                    X_test=self.X_predict,
                    X_train=self.X_train,
                    y_train= y_train,
                    model=m,
                    threshold=threshold,
                    return_all=return_all
                )
                aoa_dataset[f"{target}_aoa"] = aoa
                aoa_dataset[f"{target}_di"] = di_test
                aoa_dataset[f"{target}_lpd"] = lpd_test
                aoa_dataset[f"{target}_cutpoint"] = cutpoint

                encoding = {
                    f"{target}_aoa": {"zlib": True, "complevel": 4, "dtype": "float32", "_FillValue": np.float32(np.nan)},
                    f"{target}_di": {"zlib": True, "complevel": 4, "dtype": "float64", "_FillValue": np.float64(np.nan)},
                    f"{target}_lpd": {"zlib": True, "complevel": 4, "dtype": "float64", "_FillValue": np.float64(np.nan)},
                    f"{target}_cutpoint": {"zlib": True, "complevel": 4, "dtype": "float64", "_FillValue": np.float64(np.nan)},
                }
                
            elif return_all is False:
                aoa, di_test, cutpoint = area_of_applicability(
                    X_test=self.X_predict,
                    X_train=self.X_train,
                    y_train= y_train,
                    model=m,
                    threshold=threshold,
                    return_all=return_all
                )
                aoa_dataset[f"{target}_aoa"] = aoa
                aoa_dataset[f"{target}_di"] = di_test
                aoa_dataset[f"{target}_cutpoint"] = cutpoint

                encoding = {
                    f"{target}_aoa": {"zlib": True, "complevel": 4, "dtype": "float32", "_FillValue": np.float32(np.nan)},
                    f"{target}_di": {"zlib": True, "complevel": 4, "dtype": "float64", "_FillValue": np.float64(np.nan)},
                    f"{target}_cutpoint": {"zlib": True, "complevel": 4, "dtype": "float64", "_FillValue": np.float64(np.nan)}
                }                

            else:
                raise ValueError(f"return_all requires a boolean input, but recieved {return_all}")

        # convert df to xarray ds:
        aoa_dataset = aoa_dataset.to_xarray()

        # add metadata:
        aoa_dataset['lat'].attrs['units'] = 'degrees_north'
        aoa_dataset['lat'].attrs['long_name'] = 'latitude'

        aoa_dataset['lon'].attrs['units'] = 'degrees_east'
        aoa_dataset['lon'].attrs['long_name'] = 'longitude'

        aoa_dataset['depth'].attrs['units'] = 'm'
        aoa_dataset['depth'].attrs['positive'] = 'down'

        # export aoa to netcdf:
        aoa_dataset.to_netcdf(os.path.join(self.path_out, "aoa.nc"), encoding=encoding)

    def merge_env(self):
        """
        Merge model output with environmental data.
        This method aligns and merges the predicted values (model output) with the existing 
        environmental dataset stored in `self.d`. The merged data replaces `self.d`.
        Returns
        -------
        None
        """

        X_predict = self.X_predict.to_xarray()
        ds = self.d.to_xarray()
        aligned_datasets = xr.align(ds,X_predict, join="inner")
        ds = xr.merge(aligned_datasets)
        if 'FID' in ds:
            ds['FID'] = ds['FID'].where(ds['FID'] != '', np.nan)
        self.d = ds.to_dataframe()
        self.d = self.d.dropna()
        
    def export_ds(self, file_name, 
                  author=None, description=None):
        """
        Export the processed dataset to a NetCDF file.

        This method saves the processed dataset (`self.d`) to a NetCDF file in the location 
        defined by `self.path_out`, with optional metadata such as author and description.

        Parameters
        ----------
        file_name : str 
            The name of the NetCDF file (without extension). 
        author : str, optional
            The name of the author to include in NetCDF metadata (default is None).
        description : str, optional
            A description or title to include in the NetCDF metadata (default is None).

        Notes
        -----
        - The export location is defined in the `model_config.yml` file and is stored in `self.path_out`.
        - The method sets metadata attributes such as conventions, creator name, and units for 
        latitude, longitude, and depth.
        - Missing directories in the export path are created if necessary.
        - The file is saved with a suffix that includes the `pi` value (e.g., `_PI50.nc`).
        """
    
        #make new dir if needed
        os.makedirs(self.path_out, exist_ok=True)


        logger.info("export_ds")
        logger.info(f"dataframe: {self.d.head()}")
        ds = self.d.to_xarray()

        if description is not None:
            ds.attrs['description'] = description
        ds.attrs['Conventions'] = 'CF-1.5'
        if author is not None:
            ds.attrs['creator_name'] = author
        try:
            ds['lat'].attrs['units'] = 'degrees_north'
            ds['lat'].attrs['long_name'] = 'latitude'
        except KeyError:
            pass
        try:
            ds['lon'].attrs['units'] = 'degrees_east'
            ds['lon'].attrs['long_name'] = 'longitude'
        except KeyError:
            pass
        try:
            ds['depth'].attrs['units'] = 'm'
            ds['depth'].attrs['positive'] = 'down'
        except KeyError:
            pass
        # TODO: to add loop defining units of variables

        logger.info(self.d.head())
        fname = os.path.join(self.path_out, file_name) + "_" + self.statistic + self.datatype + ".nc"
        ds.to_netcdf(fname)
        logger.info(f"exported ds to: {fname}")
        #add nice metadata


    def export_csv(self, file_name):
        """
        Export the processed dataset to a csv file.

        This method saves the processed dataset (`self.d`) to a csv file in the location 
        defined by `self.path_out`, with optional metadata such as author and description.

        Parameters
        ----------
        file_name : str 
            The name of the csv file (without extension). 
        
        Notes
        -----
        - The export location is defined in the `model_config.yml` file and is stored in `self.path_out`.
        - Missing directories in the export path are created if necessary.
        - The file is saved with a suffix that includes the `pi` value (e.g., `_PI50.nc`).
        """
    
        os.makedirs(self.path_out, exist_ok=True)
    
        logger.info(self.d.head())
        fname = os.path.join(self.path_out, file_name) + "_" + self.statistic + self.datatype + ".csv"
        self.d.to_csv(fname)

        logger.info(f"exported d to: {fname}")

    def merge_obs(self, file_name, targets=None, index_cols=['lat', 'lon', 'depth', 'time']):
        """
        Merge model output with observational data and calculate residuals.

        This function integrates model predictions with observational data based on 
        spatial and temporal indices, calculates residuals, and exports the merged dataset.

        Parameters
        ----------
        file_name : str
            The base name of the output file to save the merged dataset.
        targets : an np.array of str, optional
            An np.array of target variable names to include in the merge. If None, the default 
            targets from `self.targets` are used (default is None).
        index_cols : an np.array of str, optional
            A list of indices which denote the data dimensions. (default is ['lat', 'lon', 'depth', 'time'])

        Notes
        -----
        - The function matches the observational data with model predictions based on the 
        index_cols which default to `['lat', 'lon', 'depth', 'time']`.
        - Residuals are calculated as `observed - predicted` for each target variable.
        - Columns included in the output are the original targets, their modeled values 
        (suffixed with `_mod`), and their residuals (suffixed with `_resid`).
        - The merged dataset is saved as a CSV file with a suffix `_PI` followed by the 
        `pi` value, appended to the output file name.
        - Observational data is loaded from the path defined in `self.model_config['training']`.

        Raises
        ------
        FileNotFoundError
            If the observational dataset file cannot be found at the specified location.
        """
        # Select and rename the target columns for d
        if targets is None:
            targets = self.targets
        d = self.d[targets]

        mod_columns = {target: target + '_mod' for target in targets}
        d = d.rename(mod_columns, axis=1)
        d.reset_index(inplace=True)
        d.set_index(index_cols, inplace=True)        

        # Read the training targets from the training.csv file defined in model_config
        try:
            df2_path = os.path.join(self.root, self.model_config['training'])

            df2 = pd.read_csv(df2_path)
        except:
            raise FileNotFoundError(f"Dataset not found at {df2_path}")
        

        df2.set_index(index_cols, inplace=True)
        df2['dummy'] = 1

        out = pd.concat([df2, d], axis=1)
        out = out[out['dummy'] == 1].drop(['dummy'], axis=1)

        # Calculate residuals
        for target in targets:
            out[target + '_resid'] = out[target] - out[target + '_mod']

        # Define the columns to keep in the final DataFrame
        keep_columns = list(targets) + list(mod_columns.values()) + [target + '_resid' for target in targets]

        out = out[keep_columns]
        file_name = f"{file_name}_obs"
        logger.info(out.head())
        fname = os.path.join(self.path_out, file_name)  + self.datatype +  ".csv"
        out.to_csv(fname)

        logger.info(f"exported d to: {fname}")

        logger.info('training merged with predictions')