""" Main module for optimal transport dataset distance.

Throught this module, source and target are often used to refer to the two datasets
being compared. This notation is legacy from NLP, and does not carry other particular
meaning, e.g., the distance is nevertheless symmetric (though not always identical -
due to stochsticity in the computation) to the order of D1/D2. The reason for this
notation is that here X and Y are usually reserved to distinguish between features and labels.

Other important notation:
    X1, X2: feature tensors of the two datasets
    Y1, Y2: label tensors of the two datasets
    N1, N2 (or N,M): number of samples in datasets
    D1, D2: (feature) dimension of the datasets
    C1, C2: number of classes in the datasets
    π: transport coupling

"""
import os
import pdb
from time import time
import itertools
import numpy as np
from tqdm.autonotebook import tqdm
import torch
from functools import partial
import inspect
import logging
import geomloss
import ot
from torch.distributions.multivariate_normal import MultivariateNormal
from matplotlib import cm
import matplotlib as mpl
import matplotlib.pyplot as plt

from sklearn.cluster import k_means, DBSCAN

## Local Imports
from ..plotting import heatmap, gaussian_density_plot, imshow_group_boundaries
from .utils import augmented_dataset, extract_data_targets
from .moments import compute_label_stats
from .wasserstein import efficient_pwdist_gauss, pwdist_exact, pwdist_upperbound, pwdist_means_only
from .utils import register_gradient_hook, process_device_arg, multiclass_hinge_loss, load_full_dataset


import matplotlib
if os.name == 'posix' and "DISPLAY" not in os.environ:
    matplotlib.use('Agg')
    nodisplay = True
else:
    nodisplay = False


logger = logging.getLogger(__name__)

try:
    import ot.gpu
except:
    logger.warning('ot.gpu not found - coupling computation will be in cpu')


cost_routines = {
    1: (lambda x, y: geomloss.utils.distances(x, y)),
    2: (lambda x, y: geomloss.utils.squared_distances(x, y) / 2),
}


class DatasetDistance():
    """The main class for the Optimal Transport Dataset Distance.

    An object of this class is instantiated with two datasets (the source and
    target), which are stored in it, and various arguments determining how the
    OTDD is to be computed.


    Arguments:
        D1 (Dataset or Dataloader): the first (aka source) dataset.
        D2 (Dataset or Dataloader): the second (aka target) dataset.
        method (str): if set to 'augmentation', the covariance matrix will be
            approximated and appended to each point, if 'precomputed_labeldist',
            the label-to-label distance is computed exactly in advance.
        symmetric_tasks (bool): whether the two underlying datasets are the same.
            If true, will save some computation.
        feature_cost (str or callable): if not 'euclidean', must be a callable
            that implements a cost function between feature vectors.
        src_embedding (callable, optional): if provided, source data will be
            embedded using this function prior to distance computation.
        tgt_embedding (callable, optional): if provided, target data will be
            embedded using this function prior to distance computation.
        ignore_source_labels (bool): for unsupervised computation of distance
        ignore_target_labels (bool): for unsupervised computation of distance
        loss (str): loss type to be passed to samples_loss. only 'sinkhorn' is
            accepted for now.
        debiased_loss (bool): whether to use the debiased version of sinkhorn.
        p (int): the coefficient in the OT cost (i.e., the p in p-Wasserstein).
        entreg (float): the strength of entropy regularization for sinkhorn.
        λ_x (float): weight parameter for feature component of distance.
        λ_y (float): weight parameter for label component of distance.
        inner_ot_method (str): the method to compute the inner (instance-wise)
            OT problem. Must be one of 'gaussian_approx', 'exact', 'jdot', or
            'naive_upperbound'. If set to 'gaussian_approx', the label distributions
            are approximated as Gaussians, and thus their distance is computed as
            the Bures-Wasserstein distance. If set to 'exact', no approximation is
            used, and their distance is computed as an exact Wasserstein problem.
            If 'naive_upperbound', a simple upper bound on the exact distance is
            computed. If 'jdot', the label distance is computed using a classifi-
            cation loss (see JDOT paper for details).
        inner_ot_loss (str): loss type fo inner OT problem.
        inner_ot_debiased (bool): whether to use the debiased version of sinkhorn
            in the inner OT problem.
        inner_ot_p (int): the coefficient in the inner OT cost.
        inner_ot_entreg (float): the strength of entropy regularization for sinkhorn
            in the inner OT problem.
        diagonal_cov (bool): whether to use the diagonal approxiation to covariance.
        min_labelcount (int): classes with less than `min_labelcount` examples will
            be ignored in the computation of the distance.
        online_stats (bool): whether to compute the per-label means and covariance
            matrices online. If false, for every class, all examples are loaded
            into memory.
        coupling_method (str): If 'geomloss', the OT coupling is computed from
            the dual potentials obtained from geomloss (faster, less precise),
            if 'pot', it will recomputed using the POT library.
        sqrt_method (str): If 'spectral' or 'exact', it uses eigendecomposition
            to compute square root matrices (exact, slower). If 'approximate',
            it uses Newton-Schulz iterative algorithm (can be faster, though less exact).
        sqrt_niters (int): Only used if `sqrt_method` is 'approximate'. Determines
            the number of interations used for Newton-Schulz's approach to sqrtm.
        sqrt_pref (int): One of 0 or 1. Preference for cov sqrt used in cross-wass
            distance (only one of the two is needed, see efficient_pairwise_wd_gauss). Useful
            for differentiable settings, two avoid unecessary computational graph.
        nworkers_stats (int): number of parallel workers used in mean and
            covariance estimation.
        coupling_method (str): method to use for computing coupling matrix.
        nworkers_dists(int): number of parallel workers used in distance computation.
        eigen_correction (bool): whether to use eigen-correction on covariance
            matrices for additional numerical stability.
        device (str): Which device to use in pytorch convention (e.g. 'cuda:2')
        precision (str): one of 'single' or 'double'.
        verbose (str): level of verbosity.

    """

    def __init__(self, D1=None, D2=None,
                 ## General Arguments
                 method='precomputed_labeldist',
                 symmetric_tasks=False,
                 feature_cost='euclidean',
                 src_embedding=None,
                 tgt_embedding=None,
                 ignore_source_labels=False,
                 ignore_target_labels=False,
                 ## Outer OT (dataset to dataset) problem arguments
                 loss='sinkhorn', debiased_loss=True, p=2, entreg=0.1,
                 λ_x=1.0, λ_y=1.0,
                 ## Inner OT (label to label) problem arguments
                 inner_ot_method = 'gaussian_approx',
                 inner_ot_loss='sinkhorn',
                 inner_ot_debiased=False,
                 inner_ot_p=2,
                 inner_ot_entreg=0.1,
                 ## Gaussian Approximation Args
                 diagonal_cov=False,
                 min_labelcount=2,
                 online_stats=True,
                 sqrt_method='spectral',
                 sqrt_niters=20,
                 sqrt_pref=0,
                 nworkers_stats=0,
                 ## Misc
                 coupling_method='geomloss',
                 nworkers_dists=0,
                 eigen_correction=False,
                 device='cpu',
                 precision='single',
                 verbose=1, *args, **kwargs):

        self.method = method
        assert self.method in ['precomputed_labeldist', 'augmentation', 'jdot']
        self.symmetric_tasks = symmetric_tasks
        self.diagonal_cov = diagonal_cov
        ## For outer OT problem
        self.p = p
        self.entreg = entreg
        self.loss = loss
        self.debiased_loss = debiased_loss
        self.feature_cost = feature_cost
        self.src_embedding = src_embedding
        self.tgt_embedding = tgt_embedding
        self.ignore_source_labels = ignore_source_labels
        self.ignore_target_labels = ignore_target_labels
        self.λ_x = λ_x
        self.λ_y = λ_y
        ## For inner (label) OT problem - only used if gaussian approx is False
        self.inner_ot_method = inner_ot_method
        self.inner_ot_p = inner_ot_p
        self.inner_ot_entreg = inner_ot_entreg
        self.inner_ot_loss = inner_ot_loss
        self.inner_ot_debiased = inner_ot_debiased
        self.online_stats = online_stats
        self.coupling_method = coupling_method
        self.min_labelcount = min_labelcount
        self.nworkers_stats = nworkers_stats
        self.nworkers_dists = nworkers_dists
        self.sqrt_method = sqrt_method
        if self.sqrt_method == 'exact':
            self.sqrt_method = 'spectral'
        self.sqrt_niters = sqrt_niters
        assert sqrt_pref in [0,1], 'sqrt pref must be 0 or 1'
        self.sqrt_pref   = sqrt_pref
        self.device = device
        self.precision = precision
        self.eigen_correction = eigen_correction
        self.verbose = verbose

        if self.method == 'augmentation' and not self.diagonal_cov:
            logger.error('Method augmentation requires diagonal_cov = True')

        ## Placeholders
        self.Means = [None, None]
        self.Covs = [None, None]
        self.label_distances = None
        self.X1, self.X2 = None, None
        self.Y1, self.Y2 = None, None
        self._pwlabel_stats_1 = None
        self._pwlabel_stats_2 = None

        self.D1 = D1
        if D2 is None:
            self.D2 = self.D1
            self.symmetric_tasks = True
        else:
            self.D2 = D2

        if self.D1 is not None and self.D2 is not None:
            self._init_data(self.D1, self.D2)
        else:
            logger.warning('DatasetDistance initialized with empty data')
        
        if self.src_embedding is not None or self.tgt_embedding is not None:
            self.feature_cost = partial(FeatureCost,
                                   src_emb = self.src_embedding,
                                   src_dim = (3,28,28),
                                   tgt_emb = self.tgt_embedding,
                                   tgt_dim = (3,28,28),
                                   p = self.p, device=self.device)

        self.src_embedding = None
        self.tgt_embedding = None


    def _load_infer_labels(self, D, classes=None, reindex=None, reindex_start=None):

        if classes:
            k = len(classes)
            labeling_fun = lambda X: torch.LongTensor(k_means(X.numpy(), k)[1])
        else:
            labeling_fun = lambda X: torch.LongTensor(DBSCAN(eps=5, min_samples = 4).fit(X).labels_)

        X, Y_infer, Y_true = load_full_dataset(D, targets='infer',
                                 min_labelcount=self.min_labelcount,
                                 labeling_function = labeling_fun,
                                 return_both_targets=True,
                                 force_label_alignment=True,
                                 reindex=reindex, reindex_start=reindex_start)

        return X, Y_infer, Y_true


    def _init_data(self, D1, D2):
        """ Preprocessing of datasets. Extracts value and coding for effective
        (i.e., those actually present in sampled data) class labels.
        """

        targets1, classes1, idxs1 = extract_data_targets(D1)
        targets2, classes2, idxs2 = extract_data_targets(D2)

        ## Get effective dataset number of samples
        self.idxs1, self.idxs2 = idxs1, idxs2
        self.n1 = len(self.idxs1)
        self.n2 = len(self.idxs2)

        if (targets1 is None) or self.ignore_source_labels: # Unsupervised setting
            X, Y_infer, Y_true = self._load_infer_labels(D1, classes1, reindex=True, reindex_start=0)
            self.targets1 = targets1 = Y_infer
            self.X1, self.Y1 = X, Y_infer
            if Y_true is not None: self.Y1_true = Y_true # will not be used by OTDD, stored only for downstream eval
        else:
            self.targets1 = targets1

        ## Effective classes seen in data (idxs here needed to be filtered
        ## in case dataloader has subsampler)
        ## Indices of classes (might be different from class ids!)
        if (targets1 is None) or self.ignore_source_labels:
            vals1, cts1 = torch.unique(targets1, return_counts=True)
        else:
            vals1, cts1 = torch.unique(targets1[idxs1], return_counts=True)

        ## Ignore everything with a label occurring less than k times
        self.V1 = torch.sort(vals1[cts1 >= self.min_labelcount])[0]

        if (targets2 is None) or self.ignore_target_labels:
            reindex_start = len(self.V1) if (self.loss == 'sinkhorn' and self.debiased_loss) else True
            X, Y_infer, Y_true = self._load_infer_labels(D2, classes2, reindex=True, reindex_start=reindex_start)
            self.targets2 = targets2 = Y_infer - reindex_start
            assert self.targets2.min() == 0
            self.X2, self.Y2 = X, Y_infer
            if Y_true is not None: self.Y2_true = Y_true
        else:
            self.targets2 = targets2

        ## Effective classes seen in data (idxs here needed to be filtered
        ## in case dataloader has subsampler)
        ## Indices of classes (might be different from class ids!)
        if (targets2 is None) or self.ignore_target_labels:
            vals2, cts2 = torch.unique(targets2, return_counts=True)
        else:
            vals2, cts2 = torch.unique(targets2[idxs2], return_counts=True)

        ## Ignore everything with a label occurring less than k times
        self.V2 = torch.sort(vals2[cts2 >= self.min_labelcount])[0]


        self.classes1 = [classes1[i] for i in self.V1]
        self.classes2 = [classes2[i] for i in self.V2]


        if self.method == 'jdot': ## JDOT only works if same labels on both datasets
            assert torch.all(self.V1 == self.V2)


        ## Keep track of real classes vs indices (always 0 to n)(useful if e.g., missing classes):
        self.class_to_idx_1 = {i: c for i, c in enumerate(self.V1)}
        self.class_to_idx_2 = {i: c for i, c in enumerate(self.V2)}

    def copy(self, keep=[], drop=[]):
        """ Copy method for Dataset Distance object.

        Copies 'shell' of object only: configs, but no dataset or its derivatives.

        """
        dataattrs = ['D1', 'D2','X1', 'X2','Y1','Y2','V1','V2',
                    'targets1', 'targets2', 'classes1', 'classes2',
                    'idxs1', 'idxs2', 'class_to_idx_1', 'class_to_idx_2',
                    'Covs', 'Means', 'label_distances', '_label_mean_distances']

        initattrs = list(inspect.signature(DatasetDistance).parameters.keys())

        if not keep:
            ## By default, we keep all non-data attribs, drop all data-dependent ones
            keep = set(initattrs) - set(['D1','D2'])
        elif keep == 'all':
            keep = set(self.__dict__.keys()) # postattrs + initattrs
        else:
            keep = set(self.__dict__.keys()).difference(dataattrs).union(keep)

        kept_init_attrs = set(initattrs).intersection(set(keep))
        dobj = DatasetDistance(**{k:self.__dict__[k] for k in kept_init_attrs})

        ## Must also add attribs that are not taken by __init__ method (because they're generated after)
        kept_post_attrs = set(keep).difference(kept_init_attrs)
        dobj.__dict__.update({k:self.__dict__[k] for k in kept_post_attrs})
        return dobj

    def _load_datasets(self, maxsamples=None, device=None):
        """ Dataset loading, wrapper for `load_full_dataset` function.

        Loads full datasets into memory (into gpu if in CUDA mode).

        Arguments:
            maxsamples (int, optional): maximum number of samples to load.
            device (str, optional): if provided, will override class attribute device.
        """
        logger.info('Concatenating feature vectors...')

        ## We probably don't ever want to store the full datasets in GPU
        device = 'cpu'

        dtype = torch.DoubleTensor if self.precision == 'double' else torch.FloatTensor

        if self.loss == 'sinkhorn' and self.debiased_loss:
            ## We will need to relabel targets {0,...,n-1} and {n,...,n-m-1}
            reindex_start_d2 = len(self.V1)
        else:
            ## Suffices to relabel targets {0,...,n-1} and {0,...,m-1}
            reindex_start_d2 = 0

        if self.X1 is None or self.Y1 is None:
            assert not self.ignore_source_labels, 'Should not be here if igoring target labels'
            self.X1, self.Y1 = load_full_dataset(self.D1, targets=True,
                                                 labels_keep=self.V1,
                                                 maxsamples=maxsamples,
                                                 device=device,
                                                 dtype=dtype,
                                                 reindex=True,
                                                 reindex_start = 0)
        if self.X2 is None or self.Y2 is None:
            assert not self.ignore_target_labels, 'Should not be here if igoring target labels'
            if self.symmetric_tasks:
                self.X2, self.Y2 = self.X1, self.Y1
            else:
                self.X2, self.Y2 = load_full_dataset(self.D2, targets=True,
                                                     labels_keep=self.V2,
                                                     maxsamples=maxsamples,
                                                     device=device,
                                                     dtype=dtype,
                                                     reindex=True,
                                                     reindex_start = reindex_start_d2)

        logger.info("Full datasets sizes")
        logger.info(" * D1 = {} x {} ({} unique labels)".format(*
                                                          self.X1.shape, len(self.V1)))
        logger.info(" * D2 = {} x {} ({} unique labels)".format(*
                                                          self.X2.shape, len(self.V2)))

        print(" * D1 = {} x {} ({} unique labels)".format(*
                                                          self.X1.shape, len(self.V1)))
        print(" * D2 = {} x {} ({} unique labels)".format(*
                                                          self.X2.shape, len(self.V2)))


    def _get_label_stats(self, side='both'):
        """ Return per-label means and covariances.

        Computes means and covariances only once, then stores and retrieves in
        subsequent calls.

        """
        ## Check if already computed
        if (not None in self.Means) and (not None in self.Covs):
            return self.Means, self.Covs

        dtype = torch.DoubleTensor if self.precision == 'double' else torch.FloatTensor

        shared_args = {'to_tensor': True, 'nworkers': self.nworkers_stats,
                       'device': self.device, 'online': self.online_stats,
                       'dtype': dtype,
                       'diagonal_cov': self.diagonal_cov}


        if (side=='both' or side == 'src') and (self.Means[0] is None):
            logger.info("Computing per-class means and variances D1....")
            M1, C1 = compute_label_stats(self.D1, self.targets1, self.idxs1,
                                         self.classes1, embedding=self.src_embedding,
                                         **shared_args)
            self.Means[0] = M1.cpu() # No reason to keep this in GPU, convert on the fly
            self.Covs[0]  = C1.cpu()

        if (side == 'both' or side =='tgt') and (self.Means[1] is None):
            if self.symmetric_tasks:
                M2, C2 = self.Means[0], self.Covs[0]
            else:
                logger.info("Computing per-class means and variances D2....")
                M2, C2 = compute_label_stats(self.D2, self.targets2, self.idxs2,
                                     self.classes2, embedding=self.tgt_embedding,
                                     **shared_args)
            self.Means[1] = M2.cpu()
            self.Covs[1]  = C2.cpu()

        return self.Means, self.Covs

    def _get_label_distances(self):
        """ Precompute label-to-label distances.

        Returns tensor of size nclasses_1 x nclasses_2.

        Useful when computing multiple distances on same pair of datasets
        e.g. between subsets of each datasets. Will store them in memory.

        Only useful if method=='precomputed_labeldist', for now.

        Note that _get_label_stats not called for inner_ot_method = `exact`,
        since exact computation does not use Gaussian approximation, so means
        and covariances are not needed.

        Returns:
            label_distances (torch.tensor): tensor of size (C1, C2) with pairwise
                label-to-label distances across the two datasets.

        """
        ## Check if already computed
        if not self.label_distances is None:
            return self.label_distances

        ## If not, compute from scratch
        if self.inner_ot_method == 'gaussian_approx':
            ## Instantiate call to pairwise wasserstein distance
            pwdist = partial(efficient_pwdist_gauss,
                             symmetric=self.symmetric_tasks,
                             diagonal_cov=self.diagonal_cov,
                             sqrt_method=self.sqrt_method,
                             sqrt_niters=self.sqrt_niters,
                             sqrt_pref  =self.sqrt_pref,
                             cost_function = self.feature_cost,
                             device=self.device,
                             return_dmeans=True,
                             return_sqrts=True)

            Means, Covs = self._get_label_stats()

        elif self.inner_ot_method == 'exact':
            ## In this case, need to load data *before* computing label stats.

            if (self.X1 is None) or (self.X2 is None):
                self._load_datasets(maxsamples=None)  # for now, will use *all* data, to be equiv  to Gaussian

            pwdist = partial(pwdist_exact,
                             symmetric=self.symmetric_tasks,
                             p = self.inner_ot_p,
                             loss = self.inner_ot_loss,
                             debias=self.inner_ot_debiased,
                             entreg = self.inner_ot_entreg,
                             cost_function = self.feature_cost,
                             device=self.device)

        elif self.inner_ot_method == 'naive_upperbound':
            pwdist = partial(pwdist_upperbound,
                             symmetric=self.symmetric_tasks,
                             diagonal_cov=self.diagonal_cov,
                             device=self.device,
                             return_dmeans=True)

            Means, Covs = self._get_label_stats()

        elif self.inner_ot_method == 'means_only':
            pwdist = partial(pwdist_means_only,
                             symmetric=self.symmetric_tasks,
                             device=self.device)

            Means, Covs = self._get_label_stats()

        else:
            raise ValueError()


        if self.debiased_loss and not self.symmetric_tasks:
            ## Then we also need within-collection label distances
            if self._pwlabel_stats_1 is None:
                logger.info('Pre-computing pairwise label Wasserstein distances D1 <-> D1...')
                if self.inner_ot_method == 'gaussian_approx':
                    DYY1, DYY1_means, sqrtΣ1 = pwdist(Means[0], Covs[0])
                elif self.inner_ot_method == 'naive_upperbound':
                    DYY1, DYY1_means  = pwdist(Means[0], Covs[0])
                elif self.inner_ot_method == 'means_only':
                    DYY1 = pwdist(Means[0])
                    DYY1_means = DYY1
                else: # Exact
                    DYY1 = pwdist(self.X1, self.Y1)
            else:
                if self.inner_ot_method == 'gaussian_approx':
                    DYY1, DYY1_means, sqrtΣ1 = [self._pwlabel_stats_1[k] for k in ['dlabs','dmeans','sqrtΣ']]
                elif self.inner_ot_method in ['naive_upperbound', 'means_only']:
                    DYY1, DYY1_means = [self._pwlabel_stats_1[k] for k in ['dlabs','dmeans']]
                else:
                    DYY1 = self._pwlabel_stats_1['dlabs']

            if self._pwlabel_stats_2 is None:
                logger.info('Pre-computing pairwise label Wasserstein distances D2 <-> D2...')
                if self.inner_ot_method == 'gaussian_approx':
                    DYY2, DYY2_means, sqrtΣ2 = pwdist(Means[1], Covs[1])
                elif self.inner_ot_method == 'naive_upperbound':
                    DYY2, DYY2_means  = pwdist(Means[1], Covs[1])
                elif self.inner_ot_method == 'means_only':
                    DYY2 = pwdist(Means[1])
                    DYY2_means = DYY2
                else: # Exact
                    DYY2 = pwdist(self.X2, self.Y2)
            else:
                logger.info('Found pre-existing D2 label-label stats, will not recompute')
                if self.inner_ot_method == 'gaussian_approx':
                    DYY2, DYY2_means, sqrtΣ2 = [self._pwlabel_stats_2[k] for k in ['dlabs','dmeans','sqrtΣ']]
                elif self.inner_ot_method in ['naive_upperbound', 'means_only']:
                    DYY1, DYY1_means = [self._pwlabel_stats_2[k] for k in ['dlabs','dmeans']]
                else:
                    DYY2 = self._pwlabel_stats_2['dlabs']
        else:
            sqrtΣ1, sqrtΣ2 = None, None  # Will have to compute during cross
            DYY1 = DYY2 = None
            DYY1_means = DYY2_means = None

        ## Compute Cross-Distances
        logger.info('Pre-computing pairwise label Wasserstein distances D1 <-> D2...')
        if self.inner_ot_method == 'gaussian_approx':
            DYY12, DYY12_means, _ = pwdist(Means[0], Covs[0], Means[1], Covs[1], sqrtΣ1, sqrtΣ2)
        elif self.inner_ot_method == 'naive_upperbound':
            DYY12, DYY12_means    = pwdist(Means[0], Covs[0], Means[1], Covs[1])
        elif self.inner_ot_method == 'means_only':
            DYY12    = pwdist(Means[0], Means[1])
            DYY12_means = DYY12
        else:
            DYY12 = pwdist(self.X1,self.Y1,self.X2, self.Y2)
            DYY12_means = None


        if self.debiased_loss and self.symmetric_tasks:
            ## In this case we can reuse DXY to get DYY1 and DYY
            DYY1, DYY2 = DYY12, DYY12
            if self.inner_ot_method in ['gaussian_approx', 'naive_upperbound', 'means_only']:
                DYY1_means, DYY2_means = DXY_means, DXY_means

        if self.debiased_loss:
            D = torch.cat([torch.cat([DYY1, DYY12], 1),
                           torch.cat([DYY12.t(), DYY2], 1)], 0)
            if self.inner_ot_method in ['gaussian_approx', 'naive_upperbound', 'means_only']:
                D_means = torch.cat([torch.cat([DYY1_means, DYY12_means], 1),
                                 torch.cat([DYY12_means.t(), DYY2_means], 1)], 0)
        else:
            D = DYY12
            if self.inner_ot_method in ['gaussian_approx', 'naive_upperbound', 'means_only']:
                D_means = DYY12_means

        ## Collect and save
        self.label_distances  = D
        if self.inner_ot_method == 'gaussian_approx':
            self._label_mean_distances = D_means
            self._pwlabel_stats_1 = {'dlabs':DYY1, 'dmeans':DYY1_means, 'sqrtΣ':sqrtΣ1}
            self._pwlabel_stats_2 = {'dlabs':DYY2, 'dmeans':DYY2_means, 'sqrtΣ':sqrtΣ2}
        elif self.inner_ot_method  in ['naive_upperbound', 'means_only']:
            self._label_mean_distances = D_means
            self._pwlabel_stats_1 = {'dlabs':DYY1, 'dmeans':DYY1_means}#, 'sqrtΣ':sqrtΣ1}
            self._pwlabel_stats_2 = {'dlabs':DYY2, 'dmeans':DYY2_means}#, 'sqrtΣ':sqrtΣ2}
        else:
            self._pwlabel_stats_1 = {'dlabs':DYY1}#
            self._pwlabel_stats_2 = {'dlabs':DYY2}#, 'dmeans':DYY2_means, 'sqrtΣ':sqrtΣ2}


        return self.label_distances

    def distance(self, maxsamples=10000, return_coupling=False):
        """ Compute dataset distance.

            Note:
                Currently both methods require fully loading dataset into memory,
                this can probably be avoided, e.g., via subsampling.

            Arguments:
                maxsamples (int): maximum number of samples used in outer-level
                    OT problem. Note that this is different (and usually smaller)
                    than the number of samples used when computing means and covs.
                return_coupling (bool): whether to return the optimal coupling.

            Returns:
                dist (float): the optimal transport dataset distance value.
                π (tensor, optional): the optimal transport coupling.

        """
        device_dists = self.device
        GPU_LIMIT = 10000
        if (self.n1 > GPU_LIMIT or self.n2 > GPU_LIMIT) and maxsamples > GPU_LIMIT and self.device != 'cpu':
            logger.warning('Warning: maxsamples = {} > 5000, and device = {}. Loaded data' \
                   ' might not fit in GPU. Computing distances on' \
                   ' CPU.'.format(maxsamples, self.device))
            device_dists = 'cpu'
        if self.X1 is None or self.X2 is None:
            if (not self.method == 'jdot') and (self.λ_y is not None) and (self.λ_y > 0):
                s = time()
                _ = self._get_label_distances()
                logger.info('/* Time to precompute label distances: {} */'.format(time() - s))

            self._load_datasets(maxsamples, device=device_dists)
        if self.method == 'augmentation':
            DA = (self.X1, self.Y1)
            DB = (self.X2, self.Y2)

            if self.λ_x != 1.0 or self.λ_y != 1.0:
                raise NotImplementedError('Unevenly weighted feature/label' \
                    'not available for method=augmentation yet')

            if not hasattr(self, 'XμΣ1') or self.XμΣ1 is None:
                XA = augmented_dataset(DA, self.Means[0], self.Covs[0], maxn=maxsamples)#, diagonal_cov=self.diagonal_cov)
                del DA
                XB = augmented_dataset(DB, self.Means[1], self.Covs[1], maxn=maxsamples)#, diagonal_cov=self.diagonal_cov)
                del DB
                self.XμΣ1 = XA
                self.XμΣ2 = XB
            else:
                XA, XB = self.XμΣ1, self.XμΣ2

            loss = geomloss.SamplesLoss(
                loss=self.loss, p=self.p,
                debias=self.debiased_loss,
                blur=self.entreg**(1 / self.p), # "blur" of geomloss is eps^(1/p).
                backend='tensorized',
                )
            ## By default, use constant weights = 1/number of samples
            dist = loss(XA, XB)
            del XA, XB
        elif self.method == 'jdot':
            loss = geomloss.SamplesLoss(
                loss=self.loss, p=self.p,
                cost=partial(batch_jdot_cost, alpha = self.λ_x),
                debias=self.debiased_loss,
                blur=self.entreg**(1 / self.p),
                backend='tensorized'
            )
            if maxsamples and self.X1.shape[0] > maxsamples:
                idxs_1 = sorted(np.random.choice(
                    self.X1.shape[0], maxsamples, replace=False))
            else:
                idxs_1 = np.s_[:]  # hack to get a full slice

            if maxsamples and self.X2.shape[0] > maxsamples:
                idxs_2 = sorted(np.random.choice(
                    self.X2.shape[0], maxsamples, replace=False))
            else:
                idxs_2 = np.s_[:]  # hack to get a full slice


            Z1 = torch.cat((self.X1[idxs_1],
                            self.Y1[idxs_1].type(self.X1.dtype).unsqueeze(1)), -1)
            Z2 = torch.cat((self.X2[idxs_2],
                            self.Y2[idxs_2].type(self.X2.dtype).unsqueeze(1)), -1)
            Z1 = Z1.to(device_dists)
            Z2 = Z2.to(device_dists)
            dist = loss(Z1,Z2)
        elif self.method == 'precomputed_labeldist':
            if self.λ_y is None or self.λ_y == 0:
                W = None
            else:
                W = self._get_label_distances().to(torch.device(device_dists))

            ## This one leverages precomputed pairwise label distances
            cost_geomloss = partial(
                batch_augmented_cost,
                W=W,
                λ_x=self.λ_x,
                λ_y=self.λ_y,
                feature_cost=self.feature_cost
            )

            loss = geomloss.SamplesLoss(
                loss=self.loss, p=self.p,
                cost=cost_geomloss,
                debias=self.debiased_loss,
                blur=self.entreg**(1 / self.p),
                backend='tensorized'
            )

            if maxsamples and self.X1.shape[0] > maxsamples:
                idxs_1 = sorted(np.random.choice(
                    self.X1.shape[0], maxsamples, replace=False))
            else:
                idxs_1 = np.s_[:]  # hack to get a full slice

            if maxsamples and self.X2.shape[0] > maxsamples:
                idxs_2 = sorted(np.random.choice(
                    self.X2.shape[0], maxsamples, replace=False))
            else:
                idxs_2 = np.s_[:]  # hack to get a full slice


            Z1 = torch.cat((self.X1[idxs_1],
                            self.Y1[idxs_1].type(self.X1.dtype).unsqueeze(1)), -1)
            Z2 = torch.cat((self.X2[idxs_2],
                            self.Y2[idxs_2].type(self.X2.dtype).unsqueeze(1)), -1)
            Z1 = Z1.to(device_dists)
            Z2 = Z2.to(device_dists)
            try:
                dist = loss(Z1,Z2)
            except:
                dist = 0


            if return_coupling:
                with torch.no_grad():
                    C = cost_geomloss(Z1.unsqueeze(0), Z2.unsqueeze(0)).squeeze()
                    if self.coupling_method == 'geomloss':
                        loss.potentials = True
                        u, v = loss(Z1, Z2)
                        π = torch.exp(1 / self.entreg * (u.t() + v - C))  # * (pq)
                    elif self.coupling_method == 'pot':
                        C = C.cpu()
                        π = ot.sinkhorn(ot.unif(Z1.shape[0]), ot.unif(Z2.shape[0]),
                                        C / C.max(), self.entreg, numItermax=50,
                                        method='sinkhorn_epsilon_scaling', verbose=True)
                    else:
                        pass  # nonadimisslbe args already caught in argparse
                    del C
            del Z1, Z2


        ## Admittedly ugly but might be necessary to avoid memory clogs
        torch.cuda.empty_cache()

        if return_coupling:
            return dist, π
        else:
            return dist

    def compute_coupling(self, entreg=None, gpu=None, **kwargs):
        """ Compute the optimal transport coupling.

        Arguments:
            entreg (float): strength of entropy regularization.
            gpu (bool): whether to use gpu for coupling computation.
            **kwargs: arbitrary keyword args passed to ot.sinkhorn

        Returns:
            π (tensor): tensor of size (N1, N2) with optimal transport coupling.

        """
        if self.X1 is None or self.X2 is None:
            self._load_datasets()
        entreg = entreg if entreg else self.entreg
        Z1 = torch.cat((self.X1, self.Y1.type(self.X1.dtype).unsqueeze(1)), -1)
        Z2 = torch.cat((self.X2, self.Y2.type(self.X2.dtype).unsqueeze(1)), -1)
        ## Compute on device of Z1, Z1. If cuda is available but Z1,Z2 are in cpu
        ## it was decided in .distance() that they're too large for GPU.
        device = Z1.device
        a = ot.unif(Z1.shape[0])
        b = ot.unif(Z2.shape[0])
        W = self._get_label_distances().to(device)

        C = batch_augmented_cost(Z1.unsqueeze(0), Z2.unsqueeze(0),W=W).squeeze()
        C = C.cpu()
        if gpu is None:
            gpu = self.device != 'cpu'
        if 'method' in kwargs and kwargs['method'] == 'emd':
            π = ot.emd(a, b, C / C.max())
        elif not gpu:
            π = ot.sinkhorn(a, b, C / C.max(), entreg, **kwargs)
        else:
            kwargs['verbose'] = False
            π = ot.gpu.sinkhorn(a, b, C / C.max(), entreg, **kwargs)
        self.π = π
        return π

    def final_distance(self):
        """ Computes the outer-level OT distance between the datasets.
        FIXME: This function seems to be deprecated.
        """
        α = ot.utils.unif(self.n1)
        β = ot.utils.unif(self.n2)
        if normalize_dists == 'max':
            D_norm = D / D.max()
        else:
            D_norm = D
        π, log = ot.bregman.sinkhorn(
            α, β, D_norm, reg=reg, method='sinkhorn', log=True)
        d = (π * D).sum()
        return d, π

    def subgroup_distance(self, labels_a, labels_b, maxsamples=500):
        """
        Compute dataset distance between subsets of the two datasets, where
        the subsets are defined through subroups of labels.

        """
        mask_a = np.isin(self.Y1.cpu(), labels_a)
        idxs_a = mask_a.nonzero()[0].squeeze()
        mask_b = np.isin(self.Y2.cpu(), labels_b)
        idxs_b = mask_b.nonzero()[0].squeeze()

        if self.method == 'augmentation':
            DA = (self.X1[mask_a], self.Y1[mask_a])
            DB = (self.X2[mask_b], self.Y2[mask_b])
            XA = augmented_dataset(
                DA, self.Means[0], self.Covs[0], maxn=maxsamples)
            del DA
            XB = augmented_dataset(
                DB, self.Means[1], self.Covs[1], maxn=maxsamples)
            del DB
            pdb.set_trace()
            loss = geomloss.SamplesLoss(
                loss="sinkhorn", p=2, blur=self.entreg**(1 / self.p))
            ## By default, use constant weights = 1/number of samples
            dist = loss(XA, XB)
            del XA, XB
        elif self.method == 'precomputed_labeldist':
            ## This one leverages precomputed pairwise label distances
            cost_geomloss = partial(batch_augmented_cost,
                                    W=self._get_label_distances(),
                                    V1=len(self.V1),
                                    V2=len(self.V2))

            loss = geomloss.SamplesLoss(loss='sinkhorn', p=2, cost=cost_geomloss,
                                        blur=self.entreg**(1 / self.p), backend='tensorized')

            if maxsamples and mask_a.sum() > maxsamples:
                idxs_a = idxs_a[sorted(np.random.choice(
                    len(idxs_a), maxsamples, replace=False))]
            if maxsamples and mask_b.sum() > maxsamples:
                idxs_b = idxs_b[sorted(np.random.choice(
                    len(idxs_b), maxsamples, replace=False))]
            ZA = torch.cat(
                (self.X1[idxs_a], self.Y1[idxs_a].float().unsqueeze(1)), -1)
            ZB = torch.cat(
                (self.X2[idxs_b], self.Y2[idxs_b].float().unsqueeze(1)), -1)
            dist = loss(ZA, ZB)
            del ZA, ZB

        torch.cuda.empty_cache()

        return dist

    def plot_label_distances(self, plot_means=False, ax=None, show=True, cbar=True,
                             cmap="YlGn", cbarlabel="default", save_path=None, xlabel=None, ylabel=None,
                             fontsize=10,
                             **kwargs):
        LD = self._get_label_distances().sqrt()
        LMD = self._label_mean_distances.sqrt()

        if LD.shape[0] > len(self.V1):
            ## Means we also have self, distance, don't want to plot those usually
            LD = LD[:len(self.V1), len(self.V1):]
            LMD = LMD[:len(self.V1), len(self.V1):]

        if not ax:
            ncol = 2 if plot_means else 1
            fig, ax = plt.subplots(1, ncol, figsize=(ncol * 5, 5))
        elif type(ax) is np.ndarray:
            assert len(ax) == 2
            ncol = 2
        else:
            assert not plot_means
            ncol = 1

        ax0 = ax[0] if plot_means else ax
        if cbarlabel == 'default': cbarlabel = r"Wasserstein Distance $d(y,y')$"
        ax0.set_title('Label-to-Label Distance',fontsize=fontsize)
        im, cbar = heatmap(LD.cpu(), self.classes1, self.classes2, ax=ax0,
                           cmap=cmap, cbar=cbar, cbarlabel=cbarlabel, **kwargs)
        if plot_means:
            ax[1].set_title('Label-to-label Distance')
            im, cbar = heatmap(LMD.cpu(), self.V1.tolist(), self.V2.tolist(), ax=ax[1],
                               cmap=cmap, cbar=cbar, cbarlabel=cbarlabel, **kwargs)

        if xlabel: ax.set_xlabel(xlabel, fontsize=fontsize)
        if ylabel: ax.set_ylabel(ylabel, fontsize=fontsize)
        if save_path:
            fig = plt.gcf()
            fig.tight_layout()
            plt.savefig(save_path, format='pdf', dpi=300)
        if show:
            plt.show()

    def plot_label_stats(self, same_plot=False, show_target = True,
                         label_groups=True, label_means=True,
                         color_by='domain',
                         pad=0.1, ax=None, show=True, shift=(1,1)):
        ## Assert that both datasets are 2 dim
        Means, Covs = self._get_label_stats()
        if self.X1 is None or self.X2 is None:
            self._load_datasets()

        k1,k2 = len(self.classes1), len(self.classes2)

        if ax is None:
            ncol = 1 if same_plot else 2
            fig, ax = plt.subplots(1, ncol, figsize=(14, 7))
        elif type(ax) is np.ndarray and len(ax) == 2:
            assert not same_plot
            ncol = 2
        else:
            assert same_plot
            ncol = 1

        color_by = 'label'
        if color_by == 'domain':
            ## Only two colours, don't distinguish between label groups
            ## These are for the density plots
            colors = ['red', 'blue']
            cmaps  = [cm.Reds, cm.Blues]

            ## These are for the scatter plots
            scatter_colors = [self.Y1-self.Y1.min(),self.Y2-self.Y2.min()]
            scatter_cmaps = [
                mpl.colors.ListedColormap(cm.get_cmap('Reds', 2*k1)(np.linspace(0.4,0.8,k1))),
                mpl.colors.ListedColormap(cm.get_cmap('Blues', 2*k2)(np.linspace(0.4,0.8,k2)))
            ]
            markers = ['o', 'o']

        elif color_by == 'label':
            ## These are for the density plots
            colors = ['red', 'blue']
            cmaps = [cm.Reds, cm.Blues]

            ## These are for the scatter plots
            scatter_colors = [self.Y1-self.Y1.min(),self.Y2-self.Y2.min()]
            scatter_cmaps = [cm.get_cmap('tab10', k1),cm.get_cmap('tab10', k2)]
            markers = ['*', 'o']

        else:
            raise ValueError('Unrecognized value')

        ## Set plot limits
        lims = {'x': [None, None], 'y': [None, None]}
        padx1 = (self.X1[:, 0].max() - self.X1[:, 0].min()) * pad
        pady1 = (self.X1[:, 1].max() - self.X1[:, 1].min()) * pad
        padx2 = (self.X2[:, 0].max() - self.X2[:, 0].min()) * pad
        pady2 = (self.X2[:, 1].max() - self.X2[:, 1].min()) * pad

        lims['x'][0] = (self.X1[:, 0].min() - padx1,
                        self.X1[:, 0].max() + padx1)
        lims['y'][0] = (self.X1[:, 1].min() - pady1,
                        self.X1[:, 1].max() + pady1)
        lims['x'][1] = (self.X2[:, 0].min() - padx2,
                        self.X2[:, 0].max() + padx2)
        lims['y'][1] = (self.X2[:, 1].min() - pady2,
                        self.X2[:, 1].max() + pady2)

        ## Maybe repeated calls to single gaussian_distrib_plot, change colors
        for i in range(2):
            if i == 1 and not show_target: continue
            X, Y, c = (self.X1, self.Y1, self.classes1) if i == 0 else (
                self.X2, self.Y2, self.classes2)
            X = X.clone().detach() # In case we're in dynamic setting
            Y = Y.clone().detach()
            axi = ax[i] if ncol == 2 else ax
            if same_plot:
                axi.set_xlim(min(lims['x'][0][0], lims['x'][1][0]),
                             max(lims['x'][0][1], lims['x'][1][1]))
                axi.set_ylim(min(lims['y'][0][0], lims['y'][1][0]),
                             max(lims['y'][0][1], lims['y'][1][1]))
            else:
                axi.set_xlim(lims['x'][i])
                axi.set_ylim(lims['y'][i])
            for j, (μ, Σ) in enumerate(zip(Means[i], Covs[i])):
                μi, Σi = μ.clone().detach(), Σ.clone().detach()
                if Σ.ndim == 1: Σi = torch.diag(Σi)
                try:
                    P = MultivariateNormal(μi, Σi)
                    gaussian_density_plot(P, X=X[Y == j], method='exact',
                                      nsamples=100, label_means=label_means,
                                      color=colors[i], cmap=cmaps[i], ax=axi)
                except:
                    logger.warning('Gaussian density plot failed - probably singular covariance')
                axi.scatter(X[:, 0], X[:, 1], marker=markers[i], s=8,
                            c=scatter_colors[i], cmap=scatter_cmaps[i])#, normalize = )

                if label_groups:
                    axi.text(μ[0] + shift[0], μ[1] + shift[1], r"$y={}$".format(c[j]))
            axi.set_xlabel('')
            axi.set_ylabel('')

        if show:
            plt.show()

        return ax

    @staticmethod
    def plot_coupling(pi, Y1, Y2, ax=None, boundaries=None, sorting=False, title=None,
                      ticks=False, xlabel=None, ylabel=None, axlabel_fontsize=10,
                      show=True, save_path=None):
        Y1 -= Y1.min()
        Y2 -= Y2.min()
        if not ax:
            maxn = max(len(Y1), len(Y2))
            fig, ax = plt.subplots(
                figsize=(5 * len(Y1) / maxn, 5 * len(Y2) / maxn))

        if sorting:
            if type(sorting) is bool:
                _,idxs1 = zip(*sorted(zip(Y1.numpy(),range(len(Y1)))))
                _,idxs2 = zip(*sorted(zip(Y2.numpy(),range(len(Y2)))))
                pi = pi[idxs1, :][:, idxs2]
            elif type(sorting) is tuple:
                pi = pi[sorting[0], :][:, sorting[1]]

        ax.imshow(pi,  cmap='Reds', aspect="auto")
        ax.set_xticks([])
        ax.set_yticks([])

        if boundaries:
            src_group_sizes = torch.bincount(Y1)
            trg_group_sizes = torch.bincount(Y2)
            gnames = [range(10), range(10)]
            imshow_group_boundaries(ax, src_group_sizes, trg_group_sizes,
                                    group_names=gnames)

        if not ticks:
            ax.set_xticks([])
            ax.set_yticks([])
        else:
            ytick_spacing = max(round(pi.shape[0] / 20), 1)
            xtick_spacing = max(round(pi.shape[1] / 20), 1)
            rownames = Y1
            colnames = Y2
            ax.set_xticks(range(len(colnames))[::xtick_spacing])
            ax.set_xticklabels([str(a) for a in colnames[::xtick_spacing]])
            ax.set_yticks(range(len(rownames))[::ytick_spacing])
            ax.set_yticklabels([str(a) for a in rownames[::ytick_spacing]])

        if xlabel:
            ax.set_xlabel(r'Dataset: {}'.format(xlabel), fontsize=axlabel_fontsize)
        if ylabel:
            ax.set_ylabel(r'Dataset: {}'.format(ylabel), fontsize=axlabel_fontsize)
        if title:
            ax.set_title(title, pad=5 + 20 * (boundaries))
        if save_path:
            plt.savefig(save_path, format='pdf', dpi=300)
        if show:
            plt.show()


class IncomparableDatasetDistance(DatasetDistance):
    """ Dataset Distance subclass for datasets that have different feature dimension.

    Note:
        Proceed with caution, this class is still experimental and in active
        development

    """
    def __init__(self, *args, **kwargs):
        super(IncomparableDatasetDistance, self).__init__(*args, **kwargs)
        if self.debiased_loss:
            raise ValueError('Debiased GWOTDD not implemented yet')

    def _get_label_distances(self):
        """
            TODO: We could instead modify method in parent class to allow for only
            within-domain label distance computation.
        """
        Means, Covs = self._get_label_stats()

        ## Instantiate call to pairwise wasserstein distance
        pwdist = partial(efficient_pwdist_gauss,
                         symmetric=self.symmetric_tasks,
                         diagonal_cov=self.diagonal_cov,
                         sqrt_method=self.sqrt_method,
                         sqrt_niters=self.sqrt_niters,
                         sqrt_pref  =self.sqrt_pref,
                         device=self.device,
                         return_dmeans=True,
                         return_sqrts=True)


        if not self._pwlabel_stats_1 is None:
            logger.info('Found pre-existing D1 label-label stats, will not recompute')
            DYY1, DYY1_means, sqrtΣ1 = [self._pwlabel_stats_1[k] for k in ['dlabs','dmeans','sqrtΣ']]
        else:
            logger.info('Pre-computing pairwise label Wasserstein distances D1 <-> D1...')
            DYY1, DYY1_means, sqrtΣ1 = pwdist(Means[0], Covs[0])

        if not self._pwlabel_stats_2 is None:
            logger.info('Found pre-existing D2 label-label stats, will not recompute')
            DYY2, DYY2_means, sqrtΣ2 = [self._pwlabel_stats_2[k] for k in ['dlabs','dmeans','sqrtΣ']]
        else:
            logger.info('Pre-computing pairwise label Wasserstein distances D1 <-> D1...')
            DYY2, DYY2_means, sqrtΣ2 = pwdist(Means[1], Covs[1])


        self._pwlabel_stats_1 = {'dlabs':DYY1, 'dmeans':DYY1_means, 'sqrtΣ':sqrtΣ1}
        self._pwlabel_stats_2 = {'dlabs':DYY2, 'dmeans':DYY2_means, 'sqrtΣ':sqrtΣ2}

        return DYY1, DYY2


    def _compute_intraspace_distances(self):
        if self.X1 is None or self.X2 is None:
            self._load_datasets()
        DYY1, DYY2 = self._get_label_distances()

        Z1 = torch.cat((self.X1, self.Y1.type(self.X1.dtype).unsqueeze(1)), -1)
        Z1 = Z1.to(self.device)
        C1 = batch_augmented_cost(Z1.unsqueeze(0), Z1.unsqueeze(0), W=DYY1).squeeze()

        Z2 = torch.cat((self.X2, self.Y2.type(self.X2.dtype).unsqueeze(1)), -1)
        Z2 = Z2.to(self.device)
        C2 = batch_augmented_cost(Z2.unsqueeze(0), Z2.unsqueeze(0), W=DYY2).squeeze()

        return C1, C2

    def distance(self, maxsamples=10000, return_coupling=False):
        C1, C2 = self._compute_intraspace_distances()
        a = torch.ones(self.X1.shape[0]).to(self.device)/self.X1.shape[0]
        b = torch.ones(self.X2.shape[0]).to(self.device)/self.X2.shape[0]

        ## Normalize distances
        C1 = (C1 - C1.min())/C1.max()
        C2 = (C2 - C2.min())/C2.max()

        π, log = ot.gromov.entropic_gromov_wasserstein(C1, C2, a, b,
                                loss_fun = 'square_loss', epsilon=self.entreg,
                                log=True, verbose=True)
        dist = log['gw_dist']

        if return_coupling:
            return dist, π
        else:
            return dist


class FeatureCost():
    """ Class implementing a cost (or distance) between feature vectors.

    Arguments:
        p (int): the coefficient in the OT cost (i.e., the p in p-Wasserstein).
        src_embedding (callable, optional): if provided, source data will be
            embedded using this function prior to distance computation.
        tgt_embedding (callable, optional): if provided, target data will be
            embedded using this function prior to distance computation.

    """
    def __init__(self, src_embedding=None, tgt_embedding=None, src_dim=None,
                 tgt_dim=None, p=2, device='cpu'):
        assert (src_embedding is None) or (src_dim is not None)
        assert (tgt_embedding is None) or (tgt_dim is not None)
        self.src_emb = src_embedding
        self.tgt_emb = tgt_embedding
        self.src_dim = src_dim
        self.tgt_dim = tgt_dim
        self.p = p
        self.device = device

    def _get_batch_shape(self, b):
        if b.ndim == 3: return b.shape
        elif b.ndim == 2: return (1,*b.shape)
        elif b.ndim == 1: return (1,1,b.shape[0])

    def _batchify_computation(self, X, side='x', slices=20):
        if side == 'x':
            out = torch.cat([self.src_emb(b).to('cpu') for b in torch.chunk(X, slices, dim=0)])
        else:
            out = torch.cat([self.tgt_emb(b).to('cpu') for b in torch.chunk(X, slices, dim=0)])
        return out.to(X.device)

    def __call__(self, X1, X2):
        _orig_device = X1.device
        device = process_device_arg(self.device)
        if self.src_emb is not None:
            B1, N1, D1 = self._get_batch_shape(X1)
            try:
                X1 = self.src_emb(X1.view(-1,*self.src_dim).to(self.device)).reshape(B1, N1, -1)
            except: # Memory error?
                print('Batchifying feature distance computation')
                X1 = self._batchify_computation(X1.view(-1,*self.src_dim).to(self.device), 'x').reshape(B1, N1, -1)
        if self.tgt_emb is not None:
            B2, N2, D2 = self._get_batch_shape(X2)
            try:
                X2 = self.tgt_emb(X2.view(-1,*self.tgt_dim).to(self.device)).reshape(B2, N2, -1)
            except:
                print('Batchifying feature distance computation')
                X2 = self._batchify_computation(X2.view(-1,*self.tgt_dim).to(self.device), 'y').reshape(B2, N2, -1)
        if self.p == 1:
            c = geomloss.utils.distances(X1, X2)
        elif self.p == 2:
            c = geomloss.utils.squared_distances(X1, X2) / 2
        else:
            raise ValueError()
        return c.to(_orig_device)


def batch_jdot_cost(Z1, Z2, p=2, alpha=1.0, feature_cost=None):
    " https://papers.nips.cc/paper/6963-joint-distribution-optimal-transportation-for-domain-adaptation.pdf"
    B, N, D1 = Z1.shape
    B, M, D2 = Z2.shape
    assert (D1 == D2) or (feature_cost is not None)
    Y1 = Z1[:, :, -1].long()
    Y2 = Z2[:, :, -1].long()
    if feature_cost is None or feature_cost == 'euclidean': # default is euclidean
        C1 = cost_routines[p](Z1[:, :, :-1], Z2[:, :, :-1])
    else:
        C1 = feature_cost(Z1[:, :, :-1], Z2[:, :, :-1])
    ## hinge loss assumes classes and indices are same for both - shift back to [0,K]
    C2 = multiclass_hinge_loss(Y1.squeeze()-Y1.min(), Y2.squeeze()-Y2.min()).reshape(B, N, M)
    return alpha*C1 + C2


def batch_augmented_cost(Z1, Z2, W=None, Means=None, Covs=None, feature_cost=None,
                         p=2, λ_x=1.0, λ_y=1.0):
    """ Batch ground cost computation on augmented datasets.

    Defines a cost function on augmented feature-label samples to be passed to
    geomloss' samples_loss. Geomloss' expected inputs determine the requirtements
    below.

    Args:
        Z1 (torch.tensor): torch Tensor of size (B,N,D1), where last position in
            last dim corresponds to label Y.
        Z2 (torch.tensor): torch Tensor of size (B,M,D2), where last position in
            last dim corresponds to label Y.
        W (torch.tensor): torch Tensor of size (V1,V2) of precomputed pairwise
            label distances for all labels V1,V2 and returns a batched cost
            matrix as a (B,N,M) Tensor. W is expected to be congruent with p.
            I.e, if p=2, W[i,j] should be squared Wasserstein distance.
        Means (torch.tensor, optional): torch Tensor of size (C1, D1) with per-
            class mean vectors.
        Covs (torch.tensor, optional): torch Tehsor of size (C2, D2, D2) with
            per-class covariance matrices
        feature_cost (string or callable, optional): if None or 'euclidean', uses
            euclidean distances as feature metric, otherwise uses this function
            as metric.
        p (int): order of Wasserstein distance.
        λ_x (float): weight parameter for feature component of distance
        λ_y (float): weight parameter for label component of distance

    Returns:
        D (torch.tensor): torch Tensor of size (B,N,M)

    Raises:
        ValueError: If neither W nor (Means, Covs) are provided.

    """

    B, N, D1 = Z1.shape
    B, M, D2 = Z2.shape
    assert (D1 == D2) or (feature_cost is not None)

    Y1 = Z1[:, :, -1].long()
    Y2 = Z2[:, :, -1].long()

    ## Compute distances between features X in the usual way
    if λ_x is None or λ_x == 0:
        ## Features ignored in d(z,z'), C1 is dummy
        logger.info('no d_x')
        C1 = torch.zeros(B,N,M)
    elif feature_cost is None or feature_cost == 'euclidean': # default is euclidean
        C1 = cost_routines[p](Z1[:, :, :-1], Z2[:, :, :-1])  # Get from GeomLoss
    else:
        C1 = feature_cost(Z1[:, :, :-1], Z2[:, :, :-1])

    if λ_y is None or λ_y == 0:
        ## Labels ignored in d(z,z'), C2 is dummy
        logger.info('no d_y')
        C2 = torch.zeros_like(C1)
        λ_y = 0.0
    elif W is not None:
        ## Label-to-label distances have been precomputed and passed
        ## Stores flattened index corresponoding to label pairs
        M = W.shape[1] * Y1[:, :, None] + Y2[:, None, :]
        C2 = W.flatten()[M.flatten(start_dim=1)].reshape(-1,Y1.shape[1], Y2.shape[1])
    elif Means is not None and Covs is not None:
        ## We need to compate label distances too
        dmeans = cost_routines[p](Means[0][Y1.squeeze()], Means[1][Y2.squeeze()])
        dcovs  = torch.zeros_like(dmeans)
        pdb.set_trace("TODO: need to finish this. But will we ever use it?")
    else:
        raise ValueError("Must provide either label distances or Means+Covs")

    assert C1.shape == C2.shape

    ## NOTE: geomloss's cost_routines as defined above already divide by p. We do
    ## so here too for consistency. But as a consequence, need to divide C2 by p too.
    D = λ_x * C1  +  λ_y * (C2/p)

    return D
