import numpy as np


class DefaultExtension:
    '''Extend a model by replacing removed features with default values.'''
    def __init__(self, values, model):
        self.model = model
        if values.ndim == 1:
            values = values[np.newaxis]
        elif values[0] != 1:
            raise ValueError('values shape must be (dim,) or (1, dim)')
        self.values = values
        self.values_repeat = values

    def __call__(self, x, S):
        # Prepare x.
        if len(x) != len(self.values_repeat):
            self.values_repeat = self.values.repeat(len(x), 0)

        # Replace specified indices.
        x_ = x.copy()
        x_[~S] = self.values_repeat[~S]

        # Make predictions.
        return self.model(x_)


class MarginalExtension:
    '''Extend a model by marginalizing out removed features using their
    marginal distribution.'''
    def __init__(self, data, model):
        self.model = model
        self.data = data
        self.data_repeat = data
        self.samples = len(data)
        # self.x_addr = None
        # self.x_repeat = None

    def __call__(self, x, S):
        # Prepare x and S.
        n = len(x)
        x = x.repeat(self.samples, 0)
        S = S.repeat(self.samples, 0)
        # if self.x_addr != id(x):
        #     self.x_addr = id(x)
        #     self.x_repeat = x.repeat(self.samples, 0)
        # x = self.x_repeat

        # Prepare samples.
        if len(self.data_repeat) != self.samples * n:
            self.data_repeat = np.tile(self.data, (n, 1))

        # Replace specified indices.
        x_ = x.copy()
        x_[~S] = self.data_repeat[~S]

        # Make predictions.
        pred = self.model(x_)
        pred = pred.reshape(-1, self.samples, *pred.shape[1:])
        return np.mean(pred, axis=1)


class UniformExtension:
    '''Extend a model by marginalizing out removed features using a
    uniform distribution.'''
    def __init__(self, values, categorical_inds, samples, model):
        self.model = model
        self.values = values
        self.categorical_inds = categorical_inds
        self.samples = samples

    def __call__(self, x, S):
        # Prepare x and S.
        n = len(x)
        x = x.repeat(self.samples, 0)
        S = S.repeat(self.samples, 0)

        # Prepare samples.
        samples = np.zeros((n * self.samples, x.shape[1]))
        for i in range(x.shape[1]):
            if i in self.categorical_inds:
                inds = np.random.choice(
                    len(self.values[i]), n * self.samples)
                samples[:, i] = self.values[i][inds]
            else:
                samples[:, i] = np.random.uniform(
                    low=self.values[i][0], high=self.values[i][1],
                    size=n * self.samples)

        # Replace specified indices.
        x_ = x.copy()
        x_[~S] = samples[~S]

        # Make predictions.
        pred = self.model(x_)
        pred = pred.reshape(-1, self.samples, *pred.shape[1:])
        return np.mean(pred, axis=1)


class UniformContinuousExtension:
    '''
    Extend a model by marginalizing out removed features using a
    uniform distribution. Specific to sets of continuous features.

    TODO: should we have caching here for repeating x?

    '''
    def __init__(self, min_vals, max_vals, samples, model):
        self.model = model
        self.min = min_vals
        self.max = max_vals
        self.samples = samples

    def __call__(self, x, S):
        # Prepare x and S.
        x = x.repeat(self.samples, 0)
        S = S.repeat(self.samples, 0)

        # Prepare samples.
        u = np.random.uniform(size=x.shape)
        samples = u * self.min + (1 - u) * self.max

        # Replace specified indices.
        x_ = x.copy()
        x_[~S] = samples[~S]

        # Make predictions.
        pred = self.model(x_)
        pred = pred.reshape(-1, self.samples, *pred.shape[1:])
        return np.mean(pred, axis=1)


class ProductMarginalExtension:
    '''Extend a model by marginalizing out removed features the
    product of their marginal distributions.'''
    def __init__(self, data, samples, model):
        self.model = model
        self.data = data
        self.data_repeat = data
        self.samples = samples

    def __call__(self, x, S):
        # Prepare x and S.
        n = len(x)
        x = x.repeat(self.samples, 0)
        S = S.repeat(self.samples, 0)

        # Prepare samples.
        samples = np.zeros((n * self.samples, x.shape[1]))
        for i in range(x.shape[1]):
            inds = np.random.choice(len(self.data), n * self.samples)
            samples[:, i] = self.data[inds, i]

        # Replace specified indices.
        x_ = x.copy()
        x_[~S] = samples[~S]

        # Make predictions.
        pred = self.model(x_)
        pred = pred.reshape(-1, self.samples, *pred.shape[1:])
        return np.mean(pred, axis=1)


class SeparateModelExtension:
    '''Extend a model using separate models for each subset of features.'''
    def __init__(self, model_dict):
        self.model_dict = model_dict

    def __call__(self, x, S):
        output = []
        for i in range(len(S)):
            # Extract model.
            row = S[i]
            model = self.model_dict[str(row)]

            # Make prediction.
            output.append(model(x[i:i+1, row]))

        return np.concatenate(output, axis=0)


class ConditionalExtension:
    '''Extend a model by marginalizing out removed features using a model of
    their conditional distribution.'''
    def __init__(self, conditional_model, samples, model):
        self.model = model
        self.conditional_model = conditional_model
        self.samples = samples
        self.x_addr = None
        self.x_repeat = None

    def __call__(self, x, S):
        # Prepare x.
        if self.x_addr != id(x):
            self.x_addr = id(x)
            self.x_repeat = x.repeat(self.samples, 0)
        x = self.x_repeat

        # Prepare samples.
        S = S.repeat(self.samples, 0)
        samples = self.conditional_model(x, S)

        # Replace specified indices.
        x_ = x.copy()
        x_[~S] = samples[~S]

        # Make predictions.
        pred = self.model(x_)
        pred = pred.reshape(-1, self.samples, *pred.shape[1:])
        return np.mean(pred, axis=1)


class ConditionalSupervisedExtension:
    '''Extend a model using a supervised surrogate model.'''
    def __init__(self, surrogate):
        self.surrogate = surrogate

    def __call__(self, x, S):
        return self.surrogate(x, S)
