from ..modeltest import Task
from .util import add_default_output
import time
from collections import namedtuple
from connectome.analysis import standard_analysis
import numpy as np
import numpy.linalg as la
from connectome.model.network import Network

HistoryItem = namedtuple("HistoryItem", "cos_similarity d_relative_magnitude one_norm_a a")


class FEVERDynamics:
    RANDOM_VECTOR_FILLING_FRACTION = .1

    def __init__(self, network: Network, feature_vectors: np.ndarray):
        self.network = network
        self.feature_vectors = feature_vectors
        self.T_max_input_time = 1
        self.dt_over_tau = .01
        self.max_activation = .2

    def __repr__(self):
        return "{}({})".format(self.__class__.__name__, self.network)

    def activation_function(self, a):
        return np.minimum(np.maximum(a, 0), self.max_activation)

    def statistics(self, x, y):
        norm_x = la.norm(x, 2)
        norm_y = la.norm(y, 2)
        if norm_x == 0 or norm_y == 0:
            return 0, norm_x, norm_y
        return (x * y).sum() / norm_x / norm_y, norm_x, norm_y

    def make_random_activity(self):
        n = self.network.nr_neurons
        a = np.random.rand(n,1) - .5
        a[np.random.rand(n,1) < self.RANDOM_VECTOR_FILLING_FRACTION] = 0
        a /= la.norm(a, 2)
        a = self.activation_function(a)
        return a

    def external_input(self, t):
        if self.T_max_input_time > 0:
            return self.external_input_activity_ / self.T_max_input_time if t < self.T_max_input_time else 0
        else:
            return self.external_input_activity_ if t == 0 else 0

    def run(self, time_in_units_of_tau):
        n_time_steps = int(time_in_units_of_tau / self.dt_over_tau)
        nr_neurons = self.network.nr_neurons
        adjacency_matrix = self.network.adjacency_matrix
        history = []

        start_time = time.time()

        self.time_points_ = np.arange(n_time_steps) * self.dt_over_tau

        self.external_input_activity_ = self.make_random_activity()
        self.target_representation_ = self.feature_vectors @ self.external_input_activity_

        activity = np.zeros((nr_neurons,1))
        for step in range(n_time_steps):
            t = step * self.dt_over_tau
            current_representation = self.feature_vectors @ activity
            cos_sim, norm_d, norm_d0 = self.statistics(current_representation, self.target_representation_)
            history.append(HistoryItem(cos_sim,
                                       norm_d / norm_d0,
                                       la.norm(activity, 1) / nr_neurons,
                                       activity.copy().flatten()))
            activity += self.dt_over_tau * (- self.activation_function(activity)
                                            + adjacency_matrix @ self.activation_function(activity)
                                            + self.external_input(t))
        self.end_representation_ = current_representation
        self.duration_ = time.time() - start_time
        self.cosine_similarity_, self.relative_feature_vector_norm_, self.average_population_activity, self.individual_activities_ = map(np.asarray, zip(*history))


class MemoryTask(Task):
    """
    Short term memory task.
    Test the short term memory property of a network.

    Accepts network models with feature vectors.
    E.g., API and FEVER are accepted.

    The task representation is obtained as the sum over the neurons' feature
    vectors weighted by their activity.

    Parameters
    ----------

    simulation_time: float
        Simulation time relative to the neuronal time constant,
    """
    def __init__(self, simulation_time=20):
        self.simulation_time = simulation_time

    @add_default_output
    def __call__(self, model) -> dict:
        network = model()
        net_analysis = standard_analysis(network)
        feature_vectors = model.get_feature_vectors()
        feverdynamics = FEVERDynamics(network, feature_vectors)
        feverdynamics.run(self.simulation_time)
        res = {"cosine_similarity": feverdynamics.cosine_similarity_[-1],
               "target_representation": feverdynamics.target_representation_,
               "final_representation": feverdynamics.end_representation_,
                "activities": feverdynamics.individual_activities_}
        res.update(net_analysis)
        return res
