import scipy as sp

from ..modeltest import Task
from .util import add_default_output


class TuningTask(Task):
    """
    Stimulus tuning task.

    Checks if the cortical tuning is sharper than the thalamic input.

    Accepts network models with feature vectors.
    E.g. FEVER and API are accepted.
    """
    ALPHA = .1
    INPUT_SCALING = .2
    NUMBER_SIMULATION_STEPS = 1000

    def __init__(self, return_arrays=False):
        self.return_arrays = return_arrays

    @add_default_output
    def __call__(self, model):

        # only initialize theano when really needed
        import theano
        import theano.tensor as T

        net = model()
        weight_factor = {"E": 1,
                         "I": 1}
        for input_from in ["E", "I"]:
            mean_input_weight = net[input_from, :].sum(axis=1).mean()
            net[input_from, :] *= weight_factor[input_from] / sp.absolute(mean_input_weight)

        selected_neuron_nr = sp.random.randint(model.nr_exc + model.nr_inh)
        feature_vectors = model.get_feature_vectors()
        stimulus = feature_vectors[:, selected_neuron_nr]
        receptive_field_overlap = feature_vectors * stimulus.reshape(-1, 1)
        cosine_similarity_to_stimulus = receptive_field_overlap.sum(axis=0)  # thung according to feature vector
        rectified_receptive_field_overlap = sp.maximum(receptive_field_overlap, 0)
        thalamic_input = rectified_receptive_field_overlap.sum(axis=0)

        n_steps = T.scalar(dtype="uint32")
        adjacency_matrix = T.matrix()
        adj = net.adjacency_matrix
        initial = T.zeros(adj.shape[0])
        external_input = T.vector()
        alpha = T.scalar()

        def step(a, adjacency_matrix, external_input, alpha):
            return a + alpha * (-a + adjacency_matrix.dot(T.nnet.relu(a + external_input)))

        time_trace, updates = theano.scan(step,
                                          outputs_info=[initial],
                                          non_sequences=[adjacency_matrix, external_input, alpha],
                                          n_steps=n_steps)
        simulate_t = theano.function([alpha, adjacency_matrix, external_input, n_steps], time_trace,
                                     allow_input_downcast=True)
        pool_activity = simulate_t(self.ALPHA, adj, thalamic_input * self.INPUT_SCALING, self.NUMBER_SIMULATION_STEPS)
        pool_activity_at_the_end_of_the_trial = pool_activity[-1]
        correlation = sp.corrcoef([cosine_similarity_to_stimulus, pool_activity_at_the_end_of_the_trial])[0, 1]

        neg_cos_neurons = sp.argwhere(cosine_similarity_to_stimulus < 0)[:,0]
        neg_cos_sim_corr = sp.corrcoef([cosine_similarity_to_stimulus[neg_cos_neurons],
                                        pool_activity_at_the_end_of_the_trial[neg_cos_neurons]])[0, 1]

        res = {"tuning_activity_correlation": correlation,
                "pool_activity_min": pool_activity_at_the_end_of_the_trial.min(),
                "pool_activity_max": pool_activity_at_the_end_of_the_trial.max(),
                "neg_cos_sim_corr": neg_cos_sim_corr}

        if self.return_arrays:
            res.update({"pool_activity": pool_activity,
                        "thalamic_input": thalamic_input,
                        "cosine_similarity_to_stimulus": cosine_similarity_to_stimulus,
                        "feature_vectors": feature_vectors,
                        "stimulus": stimulus})

        return res