# -*- coding: utf-8 -*-
"""
Sup3r batch_handling module.
@author: bbenton
"""
import logging
import numpy as np
from datetime import datetime as dt
import os
import pickle
from concurrent.futures import (as_completed, ThreadPoolExecutor)

from rex.utilities import log_mem

from sup3r.utilities.utilities import (estimate_max_workers,
                                       weighted_time_sampler,
                                       spatial_coarsening,
                                       temporal_coarsening,
                                       nsrdb_reduce_daily_data,
                                       uniform_box_sampler,
                                       uniform_time_sampler,
                                       nn_fill_array)
from sup3r.preprocessing.data_handling import DataHandlerDCforH5

np.random.seed(42)

logger = logging.getLogger(__name__)


class Batch:
    """Batch of low_res and high_res data"""

    def __init__(self, low_res, high_res):
        """Stores low and high res data

        Parameters
        ----------
        low_res : np.ndarray
            4D | 5D array
            (batch_size, spatial_1, spatial_2, features)
            (batch_size, spatial_1, spatial_2, temporal, features)
        high_res : np.ndarray
            4D | 5D array
            (batch_size, spatial_1, spatial_2, features)
            (batch_size, spatial_1, spatial_2, temporal, features)
        """
        self._low_res = low_res
        self._high_res = high_res

    def __len__(self):
        """Get the number of observations in this batch."""
        return len(self._low_res)

    @property
    def shape(self):
        """Get the (low_res_shape, high_res_shape) shapes."""
        return (self._low_res.shape, self._high_res.shape)

    @property
    def low_res(self):
        """Get the low-resolution data for the batch."""
        return self._low_res

    @property
    def high_res(self):
        """Get the high-resolution data for the batch."""
        return self._high_res

    @staticmethod
    def reduce_features(high_res, output_features_ind=None):
        """Remove any feature channels that are only intended for the low-res
        training input.

        Parameters
        ----------
        high_res : np.ndarray
            4D | 5D array
            (batch_size, spatial_1, spatial_2, features)
            (batch_size, spatial_1, spatial_2, temporal, features)
        output_features_ind : list | np.ndarray | None
            List/array of feature channel indices that are used for generative
            output, without any feature indices used only for training.
        """
        if output_features_ind is None:
            return high_res
        else:
            return high_res[..., output_features_ind]

    # pylint: disable=W0613
    @classmethod
    def get_coarse_batch(cls, high_res,
                         s_enhance, t_enhance=1,
                         temporal_coarsening_method='subsample',
                         output_features_ind=None,
                         output_features=None):
        """Coarsen high res data and return Batch with high res and
        low res data

        Parameters
        ----------
        high_res : np.ndarray
            4D | 5D array
            (batch_size, spatial_1, spatial_2, features)
            (batch_size, spatial_1, spatial_2, temporal, features)
        s_enhance : int
            Factor by which to coarsen spatial dimensions of the high
            resolution data
        t_enhance : int
            Factor by which to coarsen temporal dimension of the high
            resolution data
        temporal_coarsening_method : str
            Method to use for temporal coarsening. Can be subsample, average,
            or total
        output_features_ind : list | np.ndarray | None
            List/array of feature channel indices that are used for generative
            output, without any feature indices used only for training.
        output_features : list
            List of Generative model output feature names

        Returns
        -------
        Batch
            Batch instance with low and high res data
        """
        low_res = spatial_coarsening(high_res, s_enhance)

        if t_enhance != 1:
            low_res = temporal_coarsening(low_res, t_enhance,
                                          temporal_coarsening_method)

        high_res = cls.reduce_features(high_res, output_features_ind)
        batch = cls(low_res, high_res)

        return batch


class ValidationData:
    """Iterator for validation data"""

    # Classes to use for handling an individual batch obj.
    BATCH_CLASS = Batch

    def __init__(self, data_handlers, batch_size=8, s_enhance=3, t_enhance=1,
                 temporal_coarsening_method='subsample',
                 output_features_ind=None,
                 output_features=None):
        """
        Parameters
        ----------
        handlers : list[DataHandler]
            List of DataHandler instances
        batch_size : int
            Size of validation data batches
        s_enhance : int
            Factor by which to coarsen spatial dimensions of the high
            resolution data
        t_enhance : int
            Factor by which to coarsen temporal dimension of the high
            resolution data
        temporal_coarsening_method : str
            [subsample, average, total]
            Subsample will take every t_enhance-th time step, average will
            average over t_enhance time steps, total will sum over t_enhance
            time steps
        output_features_ind : list | np.ndarray | None
            List/array of feature channel indices that are used for generative
            output, without any feature indices used only for training.
        output_features : list
            List of Generative model output feature names
        """

        handler_shapes = np.array(
            [d.sample_shape for d in data_handlers])
        assert np.all(handler_shapes[0] == handler_shapes)

        self.handlers = data_handlers
        self.batch_size = batch_size
        self.sample_shape = handler_shapes[0]
        self.val_indices = self._get_val_indices()
        self.max = np.ceil(
            len(self.val_indices) / (batch_size))
        self.s_enhance = s_enhance
        self.t_enhance = t_enhance
        self._remaining_observations = len(self.val_indices)
        self.temporal_coarsening_method = temporal_coarsening_method
        self._i = 0
        self.output_features_ind = output_features_ind
        self.output_features = output_features

    def _get_val_indices(self):
        """List of dicts to index each validation data observation across all
        handlers

        Returns
        -------
        val_indices : list[dict]
            List of dicts with handler_index and tuple_index. The tuple index
            is used to get validation data observation with
            data[tuple_index]"""

        val_indices = []
        for i, h in enumerate(self.handlers):
            for _ in range(h.val_data.shape[2]):
                spatial_slice = uniform_box_sampler(h.val_data,
                                                    self.sample_shape[:2])
                temporal_slice = uniform_time_sampler(h.val_data,
                                                      self.sample_shape[2])
                tuple_index = tuple(spatial_slice + [temporal_slice]
                                    + [np.arange(h.val_data.shape[-1])])
                val_indices.append({'handler_index': i,
                                    'tuple_index': tuple_index})
        return val_indices

    def any(self):
        """Return True if any validation data exists"""
        return any(self.val_indices)

    @property
    def shape(self):
        """Shape of full validation dataset across all handlers

        Returns
        -------
        shape : tuple
            (spatial_1, spatial_2, temporal, features)
            With temporal extent equal to the sum across all data handlers time
            dimension
        """
        time_steps = 0
        for h in self.handlers:
            time_steps += h.val_data.shape[2]
        return (self.handlers[0].val_data.shape[0],
                self.handlers[0].val_data.shape[1],
                time_steps,
                self.handlers[0].val_data.shape[3])

    def __iter__(self):
        self._i = 0
        self._remaining_observations = len(self.val_indices)
        return self

    def __len__(self):
        """
        Returns
        -------
        len : int
            Number of total batches
        """
        return int(self.max)

    def __next__(self):
        """Get validation data batch

        Returns
        -------
        batch : Batch
            validation data batch with low and high res data each with
            n_observations = batch_size
        """
        if self._remaining_observations > 0:
            if self._remaining_observations > self.batch_size:
                high_res = np.zeros((self.batch_size, self.sample_shape[0],
                                     self.sample_shape[1],
                                     self.sample_shape[2],
                                     self.handlers[0].shape[-1]),
                                    dtype=np.float32)
            else:
                high_res = np.zeros((self._remaining_observations,
                                     self.sample_shape[0],
                                     self.sample_shape[1],
                                     self.sample_shape[2],
                                     self.handlers[0].shape[-1]),
                                    dtype=np.float32)
            for i in range(high_res.shape[0]):
                val_index = self.val_indices[self._i + i]
                high_res[i, ...] = self.handlers[
                    val_index['handler_index']].val_data[
                    val_index['tuple_index']]
                self._remaining_observations -= 1

            if self.sample_shape[2] == 1:
                high_res = high_res[..., 0, :]
            batch = self.BATCH_CLASS.get_coarse_batch(
                high_res, self.s_enhance,
                t_enhance=self.t_enhance,
                temporal_coarsening_method=self.temporal_coarsening_method,
                output_features_ind=self.output_features_ind,
                output_features=self.output_features)
            self._i += 1
            return batch
        else:
            raise StopIteration


class BatchHandler:
    """Sup3r base batch handling class"""

    # Classes to use for handling an individual batch obj.
    VAL_CLASS = ValidationData
    BATCH_CLASS = Batch
    DATA_HANDLER_CLASS = None

    def __init__(self, data_handlers, batch_size=8, s_enhance=3, t_enhance=4,
                 means=None, stds=None, norm=True, n_batches=10,
                 temporal_coarsening_method='subsample', stdevs_file=None,
                 means_file=None, norm_workers=None, load_workers=None,
                 overwrite_stats=False):
        """
        Parameters
        ----------
        data_handlers : list[DataHandler]
            List of DataHandler instances
        batch_size : int
            Number of observations in a batch
        s_enhance : int
            Factor by which to coarsen spatial dimensions of the high
            resolution data to generate low res data
        t_enhance : int
            Factor by which to coarsen temporal dimension of the high
            resolution data to generate low res data
        means : np.ndarray
            dimensions (features)
            array of means for all features with same ordering as data
            features. If not None and norm is True these will be used for
            normalization
        stds : np.ndarray
            dimensions (features)
            array of means for all features with same ordering as data
            features.  If not None and norm is True these will be used form
            normalization
        norm : bool
            Whether to normalize the data or not
        n_batches : int
            Number of batches in an epoch, this sets the iteration limit for
            this object.
        temporal_coarsening_method : str
            [subsample, average, total]
            Subsample will take every t_enhance-th time step, average will
            average over t_enhance time steps, total will sum over t_enhance
            time steps
        stdevs_file : str | None
            Path to stdevs data or where to save data after calling get_stats
        means_file : str | None
            Path to means data or where to save data after calling get_stats
        load_workers : int | None
            Max number of workers to use for parallel data loading. If None
            the max number of available workers will be used.
        norm_workers : int | None
            Max number of workers to use for parallel data normalization. If
            None the max number of available workers will be used.
        overwrite_stats : bool
            Whether to overwrite stats cache files.
        """

        handler_shapes = np.array(
            [d.sample_shape for d in data_handlers])
        assert np.all(handler_shapes[0] == handler_shapes)

        now = dt.now()
        self.data_handlers = data_handlers
        self.parallel_load(load_workers)
        logger.debug(f'Finished loading data of shape {self.shape} '
                     f'for BatchHandler in {dt.now() - now}.')
        log_mem(logger, log_level='INFO')

        self._i = 0
        self.low_res = None
        self.high_res = None
        self.batch_size = batch_size
        self._val_data = None
        self.s_enhance = s_enhance
        self.t_enhance = t_enhance
        self.sample_shape = handler_shapes[0]
        self.means = means
        self.stds = stds
        self.n_batches = n_batches
        self.temporal_coarsening_method = temporal_coarsening_method
        self.current_batch_indices = None
        self.current_handler_index = None
        self.stdevs_file = stdevs_file
        self.means_file = means_file
        self.overwrite_stats = overwrite_stats
        self.norm_workers = norm_workers
        self.load_workers = load_workers

        if norm:
            self.means, self.stds = self.check_cached_stats()
            self.normalize(self.means, self.stds, norm_workers)

        logger.debug('Getting validation data for BatchHandler.')
        self.val_data = self.VAL_CLASS(
            data_handlers, batch_size=batch_size,
            s_enhance=s_enhance, t_enhance=t_enhance,
            temporal_coarsening_method=temporal_coarsening_method,
            output_features_ind=self.output_features_ind,
            output_features=self.output_features)

        logger.info('Finished initializing BatchHandler.')
        log_mem(logger, log_level='INFO')

    def parallel_normalization(self, max_workers=None):
        """Normalize data in all data handlers in parallel.

        Parameters
        ----------
        max_workers : int | None
            Max number of workers to use for parallel data normalization. If
            None the max number of available workers will be used.
        """
        proc_mem = len(self.data_handlers[0].features)
        proc_mem *= 2 * self.data_handlers[0].feature_mem
        logger.info(f'Normalizing {len(self.data_handlers)} data handlers.')
        max_workers = estimate_max_workers(max_workers, proc_mem,
                                           len(self.data_handlers))
        if max_workers == 1:
            for d in self.data_handlers:
                d.normalize(self.means, self.stds)
        else:
            with ThreadPoolExecutor(max_workers=max_workers) as exe:
                futures = {}
                now = dt.now()
                for i, d in enumerate(self.data_handlers):
                    future = exe.submit(d.normalize, self.means, self.stds)
                    futures[future] = i

                logger.info(f'Started normalizing {len(self.data_handlers)} '
                            f'data handlers in {dt.now() - now}.')

                for i, _ in enumerate(as_completed(futures)):
                    future.result()
                    logger.debug(f'{i + 1} out of {len(futures)} data handlers'
                                 ' normalized.')

    def parallel_load(self, max_workers=None):
        """Load data handler data in parallel

        Parameters
        ----------
        max_workers : int | None
            Max number of workers to use for parallel data loading. If None
            the max number of available workers will be used.
        """
        proc_mem = len(self.data_handlers[0].features)
        proc_mem *= 2 * self.data_handlers[0].feature_mem
        logger.info(f'Loading {len(self.data_handlers)} data handlers.')
        max_workers = estimate_max_workers(max_workers, proc_mem,
                                           len(self.data_handlers))
        if max_workers == 1:
            for d in self.data_handlers:
                d.load_cached_data()
        else:
            with ThreadPoolExecutor(max_workers=max_workers) as exe:
                futures = {}
                now = dt.now()
                for i, d in enumerate(self.data_handlers):
                    if d.data is None:
                        future = exe.submit(d.load_cached_data)
                        futures[future] = i

                logger.info(f'Started loading all {len(self.data_handlers)} '
                            f'data handlers in {dt.now() - now}.')

                for i, future in enumerate(as_completed(futures)):
                    future.result()
                    logger.debug(f'{i + 1} out of {len(futures)} handlers '
                                 'loaded.')

    def parallel_stats(self, max_workers=None):
        """Get standard deviations and means for training features in parallel

        Parameters
        ----------
        max_workers : int | None
            Max number of workers to use for parallel stats computation. If
            None the max number of available workers will be used.
        """
        proc_mem = 2 * self.data_handlers[0].feature_mem
        logger.info(f'Calculating stats for {len(self.training_features)} '
                    'features.')
        max_workers = estimate_max_workers(max_workers, proc_mem,
                                           len(self.training_features))
        if max_workers == 1:
            for f in self.training_features:
                self.get_stats_for_feature(f)
        else:
            with ThreadPoolExecutor(max_workers=max_workers) as exe:
                futures = {}
                now = dt.now()
                for i, f in enumerate(self.training_features):
                    future = exe.submit(self.get_stats_for_feature, f)
                    futures[future] = i

                logger.info('Started calculating stats for '
                            f'{len(self.training_features)} features in '
                            f'{dt.now() - now}.')

                for i, future in enumerate(as_completed(futures)):
                    future.result()
                    logger.debug(f'{i + 1} out of '
                                 f'{len(self.training_features)} stats '
                                 'calculated.')

    def __len__(self):
        """Use user input of n_batches to specify length

        Returns
        -------
        self.n_batches : int
            Number of batches possible to iterate over
        """
        return self.n_batches

    @property
    def training_features(self):
        """Get the ordered list of feature names held in this object's
        data handlers"""
        return self.data_handlers[0].features

    @property
    def output_features(self):
        """Get the ordered list of feature names held in this object's
        data handlers"""
        return self.data_handlers[0].output_features

    @property
    def output_features_ind(self):
        """Get the feature channel indices that should be used for the
        generated output features"""
        if self.training_features == self.output_features:
            return None
        else:
            out = [i for i, feature in enumerate(self.training_features)
                   if feature in self.output_features]
            return out

    @property
    def shape(self):
        """Shape of full dataset across all handlers

        Returns
        -------
        shape : tuple
            (spatial_1, spatial_2, temporal, features)
            With temporal extent equal to the sum across all data handlers time
            dimension
        """
        time_steps = np.sum([h.shape[-2] for h in self.data_handlers])
        return (self.data_handlers[0].shape[0], self.data_handlers[0].shape[1],
                time_steps, self.data_handlers[0].shape[-1])

    def check_cached_stats(self):
        """Get standard deviations and means for all data features from cache
        files if available.

        Returns
        -------
        means : ndarray
            Array of means for each feature
        stds : ndarray
            Array of stdevs for each feature
        """
        stdevs_check = (self.stdevs_file is not None
                        and not self.overwrite_stats)
        stdevs_check = stdevs_check and os.path.exists(self.stdevs_file)
        means_check = (self.means_file is not None
                       and not self.overwrite_stats)
        means_check = means_check and os.path.exists(self.means_file)
        if stdevs_check and means_check:
            logger.info(f'Loading stdevs from {self.stdevs_file}')
            with open(self.stdevs_file, 'rb') as fh:
                self.stds = pickle.load(fh)
            logger.info(f'Loading means from {self.means_file}')
            with open(self.means_file, 'rb') as fh:
                self.means = pickle.load(fh)

            msg = ('The training features and cached statistics are '
                   'incompatible. Number of training features is '
                   f'{len(self.training_features)} and number of stats is'
                   f' {len(self.stds)}')
            check = len(self.means) == len(self.training_features)
            check = check and (len(self.stds) == len(self.training_features))
            assert check, msg
        return self.means, self.stds

    def cache_stats(self):
        """Saved stdevs and means to cache files if files are not None
        """

        if self.stdevs_file is not None:
            logger.info(f'Saving stdevs to {self.stdevs_file}')
            with open(self.stdevs_file, 'wb') as fh:
                pickle.dump(self.stds, fh)
        if self.means_file is not None:
            logger.info(f'Saving means to {self.means_file}')
            with open(self.means_file, 'wb') as fh:
                pickle.dump(self.means, fh)

    def get_stats(self, max_workers=None):
        """Get standard deviations and means for all data features"""

        self.means = np.zeros((self.shape[-1]), dtype=np.float32)
        self.stds = np.zeros((self.shape[-1]), dtype=np.float32)

        now = dt.now()
        logger.info('Calculating stdevs/means.')
        self.parallel_stats(max_workers=max_workers)
        logger.info(f'Finished calculating stats in {dt.now() - now}.')
        self.cache_stats()

    def get_handler_mean(self, feature_idx, handler_idx):
        """Get feature mean for a given handler

        Parameters
        ----------
        feature_idx : int
            Index of feature to get mean for
        handler_idx : int
            Index of data handler to get mean for

        Returns
        -------
        float
            Feature mean
        """
        return np.nanmean(
            self.data_handlers[handler_idx].data[..., feature_idx])

    def get_handler_stdev(self, feature_idx, handler_idx, mean):
        """Get feature stdev for a given handler

        Parameters
        ----------
        feature_idx : int
            Index of feature to get stdev for
        handler_idx : int
            Index of data handler to get stdev for

        Returns
        -------
        float
            Feature stdev
        """
        istd = self.data_handlers[handler_idx].data[..., feature_idx] - mean
        istd = istd**2
        return np.nanmean(istd)

    def get_stats_for_feature(self, feature):
        """Get standard deviation and mean for requested feature

        Parameters
        ----------
        feature : str
            Feature to get stats for
        """
        idx = self.training_features.index(feature)
        logger.debug(f'Calculating stdev/mean for {feature}')
        self.means[idx] = self.get_means_for_feature(feature)
        self.stds[idx] = self.get_stdevs_for_feature(feature)

    def get_means_for_feature(self, feature):
        """Get mean for requested feature

        Parameters
        ----------
        feature : str
            Feature to get mean for
        """
        idx = self.training_features.index(feature)
        logger.debug(f'Calculating mean for {feature}')
        proc_mem = 2 * self.data_handlers[0].feature_mem
        max_workers = estimate_max_workers(self.norm_workers, proc_mem,
                                           len(self.data_handlers))

        if max_workers == 1:
            for didx, _ in enumerate(self.data_handlers):
                self.means[idx] += self.get_handler_mean(idx, didx)
        else:
            with ThreadPoolExecutor(max_workers=max_workers) as exe:
                futures = {}
                now = dt.now()
                for didx, _ in enumerate(self.data_handlers):
                    future = exe.submit(self.get_handler_mean, idx, didx)
                    futures[future] = didx

                logger.info('Started calculating means for '
                            f'{len(self.data_handlers)} data_handlers in '
                            f'{dt.now() - now}.')

                for i, future in enumerate(as_completed(futures)):
                    self.means[idx] += future.result()
                    logger.debug(f'{i + 1} out of {len(self.data_handlers)} '
                                 'means calculated.')
        self.means[idx] /= len(self.data_handlers)
        return self.means[idx]

    def get_stdevs_for_feature(self, feature):
        """Get stdev for requested feature

        Parameters
        ----------
        feature : str
            Feature to get stdev for
        """
        idx = self.training_features.index(feature)
        logger.debug(f'Calculating stdev for {feature}')
        proc_mem = 2 * self.data_handlers[0].feature_mem
        max_workers = estimate_max_workers(self.norm_workers, proc_mem,
                                           len(self.data_handlers))
        if max_workers == 1:
            for didx, _ in enumerate(self.data_handlers):
                self.stds[idx] += self.get_handler_stdev(idx, didx,
                                                         self.means[idx])
        else:
            with ThreadPoolExecutor(max_workers=max_workers) as exe:
                futures = {}
                now = dt.now()
                for didx, _ in enumerate(self.data_handlers):
                    future = exe.submit(self.get_handler_stdev, idx, didx,
                                        self.means[idx])
                    futures[future] = didx

                logger.info('Started calculating stdevs for '
                            f'{len(self.data_handlers)} data_handlers in '
                            f'{dt.now() - now}.')

                for i, future in enumerate(as_completed(futures)):
                    self.stds[idx] += future.result()
                    logger.debug(f'{i + 1} out of {len(self.data_handlers)} '
                                 'stdevs calculated.')
        self.stds[idx] /= len(self.data_handlers)
        self.stds[idx] = np.sqrt(self.stds[idx])
        return self.stds[idx]

    def normalize(self, means=None, stds=None, norm_workers=None):
        """Compute means and stds for each feature across all datasets and
        normalize each data handler dataset.  Checks if input means and stds
        are different from stored means and stds and renormalizes if they are

        Parameters
        ----------
        means : ndarray | None
            Array of means for each feature. If None it will be calculated.
        stds : ndarray | None
            Array of stdevs for each feature. If None it will be calculated.
        norm_workers : int | None
            Max number of workers to use for parallel data normalization. If
            None the max number of available workers will be used.
        """
        if means is None or stds is None:
            self.get_stats(max_workers=norm_workers)
        elif means is not None and stds is not None:
            if (not np.array_equal(means, self.means)
                    or not np.array_equal(stds, self.stds)):
                self.unnormalize()
            self.means = means
            self.stds = stds
        now = dt.now()
        logger.info('Normalizing data in each data handler.')
        self.parallel_normalization(norm_workers)
        logger.info('Finished normalizing data in all data handlers in '
                    f'{dt.now() - now}.')

    def unnormalize(self):
        """Remove normalization from stored means and stds"""
        for d in self.data_handlers:
            d.unnormalize(self.means, self.stds)

    def __iter__(self):
        self._i = 0
        return self

    def __next__(self):
        """Get the next iterator output.

        Returns
        -------
        batch : Batch
            Batch object with batch.low_res and batch.high_res attributes
            with the appropriate coarsening.
        """
        self.current_batch_indices = []
        if self._i < self.n_batches:
            handler_index = np.random.randint(0, len(self.data_handlers))
            self.current_handler_index = handler_index
            handler = self.data_handlers[handler_index]
            high_res = np.zeros((self.batch_size, self.sample_shape[0],
                                 self.sample_shape[1], self.sample_shape[2],
                                 self.shape[-1]), dtype=np.float32)

            for i in range(self.batch_size):
                high_res[i, ...] = handler.get_next()
                self.current_batch_indices.append(handler.current_obs_index)

            batch = self.BATCH_CLASS.get_coarse_batch(
                high_res, self.s_enhance, t_enhance=self.t_enhance,
                temporal_coarsening_method=self.temporal_coarsening_method,
                output_features_ind=self.output_features_ind,
                output_features=self.output_features)

            self._i += 1
            return batch
        else:
            raise StopIteration


class BatchHandlerCC(BatchHandler):
    """Batch handling class for climate change data with daily averages as the
    coarse dataset."""

    # Classes to use for handling an individual batch obj.
    VAL_CLASS = ValidationData
    BATCH_CLASS = Batch

    def __init__(self, *args, sub_daily_shape=None, **kwargs):
        """
        Parameters
        ----------
        *args : list
            Same positional args as BatchHandler
        sub_daily_shape : int
            Number of hours to use in the high res sample output. This is the
            shape of the temporal dimension of the high res batch observation.
            This time window will be sampled for the daylight hours on the
            middle day of the data handler observation.
        **kwargs : dict
            Same keyword args as BatchHandler
        """
        super().__init__(*args, **kwargs)
        self.sub_daily_shape = sub_daily_shape

    def __next__(self):
        """Get the next iterator output.

        Returns
        -------
        batch : Batch
            Batch object with batch.low_res and batch.high_res attributes
            with the appropriate coarsening.
        """
        self.current_batch_indices = []
        if self._i < self.n_batches:
            handler_index = np.random.randint(0, len(self.data_handlers))
            self.current_handler_index = handler_index
            handler = self.data_handlers[handler_index]

            low_res = None
            high_res = None

            for i in range(self.batch_size):
                obs_hourly, obs_daily_avg = handler.get_next()
                self.current_batch_indices.append(handler.current_obs_index)

                obs_hourly = self.BATCH_CLASS.reduce_features(
                    obs_hourly, self.output_features_ind)

                if low_res is None:
                    lr_shape = (self.batch_size,) + obs_daily_avg.shape
                    hr_shape = (self.batch_size,) + obs_hourly.shape
                    low_res = np.zeros(lr_shape, dtype=np.float32)
                    high_res = np.zeros(hr_shape, dtype=np.float32)

                low_res[i] = obs_daily_avg
                high_res[i] = obs_hourly

            high_res = self.reduce_high_res_sub_daily(high_res)
            low_res = spatial_coarsening(low_res, self.s_enhance)

            if (self.output_features is not None
                    and 'clearsky_ratio' in self.output_features):
                i_cs = self.output_features.index('clearsky_ratio')
                if np.isnan(high_res[..., i_cs]).any():
                    high_res[..., i_cs] = nn_fill_array(high_res[..., i_cs])

            batch = self.BATCH_CLASS(low_res, high_res)

            self._i += 1
            return batch
        else:
            raise StopIteration

    def reduce_high_res_sub_daily(self, high_res):
        """Take an hourly high-res observation and reduce the temporal axis
        down to the self.sub_daily_shape using only daylight hours on the
        center day.

        Parameters
        ----------
        high_res : np.ndarray
            5D array with dimensions (n_obs, spatial_1, spatial_2, temporal,
            n_features) where temporal >= 24 (set by the data handler).

        Returns
        -------
        high_res : np.ndarray
            5D array with dimensions (n_obs, spatial_1, spatial_2, temporal,
            n_features) where temporal has been reduced down to the integer
            self.sub_daily_shape. For example if the input temporal shape is 72
            (3 days) and sub_daily_shape=9, the center daylight 9 hours from
            the second day will be returned in the output array.
        """

        if self.sub_daily_shape is not None:
            n_days = int(high_res.shape[3] / 24)
            if n_days > 1:
                ind = np.arange(high_res.shape[3])
                day_slices = np.array_split(ind, n_days)
                day_slices = [slice(x[0], x[-1] + 1) for x in day_slices]
                assert n_days % 2 == 1, 'Need odd days'
                i_mid = int((n_days - 1) / 2)
                high_res = high_res[:, :, :, day_slices[i_mid], :]

            high_res = nsrdb_reduce_daily_data(high_res, self.sub_daily_shape)

        return high_res


class SpatialBatchHandlerCC(BatchHandler):
    """Batch handling class for climate change data with daily averages as the
    coarse dataset with only spatial samples, e.g. the batch tensor shape is
    (n_obs, spatial_1, spatial_2, features)
    """

    # Classes to use for handling an individual batch obj.
    VAL_CLASS = ValidationData
    BATCH_CLASS = Batch

    def __next__(self):
        """Get the next iterator output.

        Returns
        -------
        batch : Batch
            Batch object with batch.low_res and batch.high_res attributes
            with the appropriate coarsening.
        """

        self.current_batch_indices = []
        if self._i < self.n_batches:
            handler_index = np.random.randint(0, len(self.data_handlers))
            self.current_handler_index = handler_index
            handler = self.data_handlers[handler_index]

            high_res = None

            for i in range(self.batch_size):
                _, obs_daily_avg = handler.get_next()
                self.current_batch_indices.append(handler.current_obs_index)

                obs_daily_avg = self.BATCH_CLASS.reduce_features(
                    obs_daily_avg, self.output_features_ind)

                if high_res is None:
                    hr_shape = (self.batch_size,) + obs_daily_avg.shape
                    high_res = np.zeros(hr_shape, dtype=np.float32)

                    msg = ('SpatialBatchHandlerCC can only use n_temporal==1 '
                           'but received HR shape {} with n_temporal={}.'
                           .format(hr_shape, hr_shape[3]))
                    assert hr_shape[3] == 1, msg

                high_res[i] = obs_daily_avg

            low_res = spatial_coarsening(high_res, self.s_enhance)
            low_res = low_res[:, :, :, 0, :]
            high_res = high_res[:, :, :, 0, :]

            if (self.output_features is not None
                    and 'clearsky_ratio' in self.output_features):
                i_cs = self.output_features.index('clearsky_ratio')
                if np.isnan(high_res[..., i_cs]).any():
                    high_res[..., i_cs] = nn_fill_array(high_res[..., i_cs])

            batch = self.BATCH_CLASS(low_res, high_res)

            self._i += 1
            return batch
        else:
            raise StopIteration


class SpatialBatchHandler(BatchHandler):
    """Sup3r spatial batch handling class"""

    def __next__(self):
        if self._i < self.n_batches:
            handler_index = np.random.randint(
                0, len(self.data_handlers))
            handler = self.data_handlers[handler_index]
            high_res = np.zeros((self.batch_size, self.sample_shape[0],
                                 self.sample_shape[1], self.shape[-1]),
                                dtype=np.float32)
            for i in range(self.batch_size):
                high_res[i, ...] = handler.get_next()[..., 0, :]

            batch = self.BATCH_CLASS.get_coarse_batch(
                high_res, self.s_enhance,
                output_features_ind=self.output_features_ind)

            self._i += 1
            return batch
        else:
            raise StopIteration


class ValidationDataDC(ValidationData):
    """Iterator for data-centric validation data"""

    N_TIME_BINS = 12

    def _get_val_indices(self):
        """List of dicts to index each validation data observation across all
        handlers

        Returns
        -------
        val_indices : list[dict]
            List of dicts with handler_index and tuple_index. The tuple index
            is used to get validation data observation with
            data[tuple_index]"""

        val_indices = {}
        for t in range(self.N_TIME_BINS):
            val_indices[t] = []
            h_idx = np.random.choice(np.arange(len(self.handlers)))
            h = self.handlers[h_idx]
            for _ in range(self.batch_size):
                spatial_slice = uniform_box_sampler(
                    h.data, self.sample_shape[:2])
                weights = np.zeros(self.N_TIME_BINS)
                weights[t] = 1
                temporal_slice = weighted_time_sampler(
                    h.data, self.sample_shape[2], weights)
                tuple_index = tuple(spatial_slice + [temporal_slice]
                                    + [np.arange(h.data.shape[-1])])
                val_indices[t].append({'handler_index': h_idx,
                                       'tuple_index': tuple_index})
        return val_indices

    def __next__(self):
        if self._i < len(self.val_indices.keys()):
            high_res = np.zeros((self.batch_size, self.sample_shape[0],
                                 self.sample_shape[1],
                                 self.sample_shape[2],
                                 self.handlers[0].shape[-1]),
                                dtype=np.float32)
            val_indices = self.val_indices[self._i]
            for i, idx in enumerate(val_indices):
                high_res[i, ...] = self.handlers[
                    idx['handler_index']].data[idx['tuple_index']]

            batch = self.BATCH_CLASS.get_coarse_batch(
                high_res, self.s_enhance, t_enhance=self.t_enhance,
                temporal_coarsening_method=self.temporal_coarsening_method,
                output_features_ind=self.output_features_ind,
                output_features=self.output_features)
            self._i += 1
            return batch
        else:
            raise StopIteration


class BatchHandlerDC(BatchHandler):
    """Data-centric batch handler"""

    VAL_CLASS = ValidationDataDC
    BATCH_CLASS = Batch
    DATA_HANDLER_CLASS = DataHandlerDCforH5

    def __init__(self, *args, **kwargs):
        """
        Parameters
        ----------
        *args : list
            Same positional args as BatchHandler
        **kwargs : dict
            Same keyword args as BatchHandler
        """
        super().__init__(*args, **kwargs)

        self.temporal_weights = np.ones(self.val_data.N_TIME_BINS)
        self.temporal_weights /= np.sum(self.temporal_weights)
        self.training_sample_record = [0] * self.val_data.N_TIME_BINS
        self.normalized_sample_record = [0] * self.val_data.N_TIME_BINS
        self.old_temporal_weights = [0] * self.val_data.N_TIME_BINS
        bin_range = self.data_handlers[0].data.shape[2]
        bin_range -= (self.sample_shape[2] - 1)
        self.temporal_bins = np.array_split(np.arange(0, bin_range),
                                            self.val_data.N_TIME_BINS)
        self.temporal_bins = [b[0] for b in self.temporal_bins]

        logger.info('Using temporal weights: '
                    f'{[round(w, 3) for w in self.temporal_weights]}')

    def update_training_sample_record(self):
        """Keep track of number of observations from each temporal bin"""
        handler = self.data_handlers[self.current_handler_index]
        t_start = handler.current_obs_index[2].start
        bin_number = np.digitize(t_start, self.temporal_bins)
        self.training_sample_record[bin_number - 1] += 1

    def __iter__(self):
        self._i = 0
        self.training_sample_record = [0] * self.val_data.N_TIME_BINS
        return self

    def __next__(self):
        self.current_batch_indices = []
        if self._i < self.n_batches:
            handler_index = np.random.randint(0, len(self.data_handlers))
            self.current_handler_index = handler_index
            handler = self.data_handlers[handler_index]
            high_res = np.zeros((self.batch_size, self.sample_shape[0],
                                 self.sample_shape[1], self.sample_shape[2],
                                 self.shape[-1]), dtype=np.float32)

            for i in range(self.batch_size):
                high_res[i, ...] = handler.get_next(self.temporal_weights)
                self.current_batch_indices.append(handler.current_obs_index)

                self.update_training_sample_record()

            batch = self.BATCH_CLASS.get_coarse_batch(
                high_res, self.s_enhance, t_enhance=self.t_enhance,
                temporal_coarsening_method=self.temporal_coarsening_method,
                output_features_ind=self.output_features_ind,
                output_features=self.output_features)

            self._i += 1
            return batch
        else:
            total_count = self.n_batches * self.batch_size
            self.normalized_sample_record = [c / total_count for c
                                             in self.training_sample_record]
            self.old_temporal_weights = self.temporal_weights.copy()
            raise StopIteration
