"""
Taken from https://github.com/dmitrykazhdan/concept-based-xai.

All credit due to Kazhdan et al. (arXiv:2104.06917).
"""

'''
Implements the tabular toy dataset defined by Mahinpei et al. in "Promises and
Pitfalls of Black-Box Concept Learning Models" (found in
https://arxiv.org/abs/2106.13314).

'''

import numpy as np
from .latentFactorData import LatentFactorData


################################################################################
## GLOBAL VARIABLES
################################################################################

CONCEPT_NAMES = [
    'x_pos',
    'y_pos',
    'z_pos',
]

CONCEPT_N_VALUES = [
    1,
    1,
    1,
]


################################################################################
## DATASET LOADER
################################################################################

class TabularToy(LatentFactorData):

    def __init__(
        self,
        num_samples,
        cov=None,
        num_concepts=len(CONCEPT_NAMES),
        train_size=0.85,
        random_state=None,
        task=None,
    ):
        """
        This dataset has 7 features and 3 concepts. It is generated by 3 latent
        variables (x, y, z), randomly sampled from a multivariate normal
        distribution (with the provided covariance matrix), which define all
        seven features of a sample as:

                Feature 0: np.sin(x_vars) + x_vars
                Feature 1: np.cos(x_vars) + x_vars
                Feature 2: np.sin(y_vars) + y_vars
                Feature 3: np.cos(y_vars) + y_vars
                Feature 4 np.sin(z_vars) + z_vars
                Feature 5: np.cos(z_vars) + z_vars
                Feature 6: x_vars**2 + y_vars**2 + z_vars**2

        As concepts, we provide by default a vector of three values
        [x_pos, y_pos,z_pos] indicating whether each concept is positive or not.
        BY default, if a task is not provided, then we assign as a label to each
        sample whether or not two or more of its latent factors are positive.

        :param int num_samples: The total number of samples we will generate
            for this dataset.
        :param Or[np.ndarray, float] cov: A covariance matrix to use for
            sampling the latent variables. If provided as a float, then we will
            assign cov as the covariance between any two distinct latent
            factors. If not given, then we will use the identity matrix.
        :param int num_concepts: A number in {1, 2, 3} indicating how many
            concepts we want to include in our concept representation. If less
            than 3, then we will trim the concept list from the right.
        :param float train_size: A number in [0,1] indicating which fraction of
            the entire dataset will be used for training and which fraction will
            be used for testing.
        :param Function[(ndarray, ndarray), (ndarray, ndarray, ndarray)] task:
            The task to use with the dataset for creating
            labels. We expect a function that takes two np.ndarrays
            (x_data, c_data) corresponding to the dSprites samples and their
            respective concepts respectively, and produces a tuple of three
            np.ndarrays (x_data, c_data, y_data) corresponding to the task's
            samples, ground truth concept values, and labels, respectively.
        """

        if task is None:
            task = lambda x, c: (
                x,
                c,
                (np.sum(c, axis=-1) > 1).astype(np.int32),
            )
        super().__init__(
            dataset_path=None,
            task_name="tabular_toy",
            num_factors=num_concepts,
            sample_shape=[7],
            c_names=CONCEPT_NAMES,
            task_fn=task,
        )
        self.num_samples = num_samples
        self.cov = np.eye(3) if cov is None else cov
        if isinstance(self.cov, (float, int)):
            self.cov = np.array([
                [1, self.cov, self.cov],
                [self.cov, 1, self.cov],
                [self.cov, self.cov, 1],
            ])

        self._get_generators(train_size, random_state)

    def _load_x_c_data(self):
        # Sample the x, y, and z variables
        latent_vars = np.random.multivariate_normal(
            mean=[0, 0, 0],
            cov=self.cov,
            size=(self.num_samples,),
        )
        x_vars = latent_vars[:, 0]
        y_vars = latent_vars[:, 1]
        z_vars = latent_vars[:, 2]

        # The features are just non-linear functions applied to each
        # variable
        features = [
            np.sin(x_vars) + x_vars,
            np.cos(x_vars) + x_vars,
            np.sin(y_vars) + y_vars,
            np.cos(y_vars) + y_vars,
            np.sin(z_vars) + z_vars,
            np.cos(z_vars) + z_vars,
            x_vars**2 + y_vars**2 + z_vars**2,
        ]
        features = np.stack(features, axis=1)

        # The concepts just check if the variables are positive
        x_pos = (x_vars > 0).astype(np.int32)
        y_pos = (y_vars > 0).astype(np.int32)
        z_pos = (z_vars > 0).astype(np.int32)
        concepts = np.squeeze(
            np.stack([x_pos, y_pos, z_pos][:self.num_factors], axis=1)
        )

        # And that's it buds
        return features, concepts
