###
from conmorph.connections import load_overlap
import numpy as np
import matplotlib.pyplot as plt

file_path = "/home/pfeiffer/Projects/subiculum_interneuron_polarity/data/"
file_name = "circular_inhibitory_neurons_overlap_matrix_0"
path_to_store = file_path + "/" + file_name + ".pickle"

overlap_matrix = load_overlap(path_to_store)

plt.imshow(overlap_matrix.overlaps)
plt.show()

### Use a transfer function to turn the overlaps into connection probabilities

# The transfer function determines how much an overlap is weighted in comparison to others, the identity weighs all overlaps equally, a higher power will prefer high overlaps
def transfer_function(overlap):
    return overlap**2

# From the transferred overlap, one normalizes the results such that if each entry is treated as a probability the expected number of connections equals the number of (experimentally observed ones)

# what is the expected number of connections, given a number N of non-zero overlaps, it is just the sum of probabilites, at each overlap the mean number is the probability, so that the sum is the total number of expected connections, note that the number of overlaps has to be bigger than the expected connectivity, that is a first sanity check of the number of ovelaps

p_ex_ex = 0
p_ex_in = 0.02
p_in_in = 0.01
p_in_ex = 0.02


ex_ex, ex_in, in_in, in_ex = overlap_matrix.get_submatrices()


# Problem to normalize a matrix with only few more entries than necessary, make all one that are above one leads to too few connections
def normalize(transferred_overlaps, connection_probabiity):
    number_of_available_connections = np.sum(transferred_overlaps>0)
    number_of_expected_connections = connection_probabiity * np.prod(transferred_overlaps.shape)
    if number_of_available_connections < number_of_expected_connections:
        print("Warning: for an expected number of {:d} connections only {:d} overlaps available. No valid connection probabilities available")

    normalization = number_of_expected_connections / np.sum(transferred_overlaps)
    normalized = transferred_overlaps*normalization
    normalized[np.where(normalized>1.0)] = 1.0
    return normalized

conn_probs = [normalize(transfer_function(overlap), p) for  overlap, p in zip(overlap_matrix.get_submatrices(), [p_ex_ex, p_ex_in, p_in_in, p_in_ex])]



total_connection_probability_matrix = np.vstack((np.concatenate((conn_probs[0], conn_probs[1]), axis=1), np.concatenate((conn_probs[2], conn_probs[3]), axis=1)))

def realize(connection_probability_matrix):
    return np.random.binomial(1, connection_probability_matrix)


connectivity = realize(total_connection_probability_matrix)

n_ex = overlap_matrix.excitatory_neurons.number
n_in = overlap_matrix.inhibitory_neurons.number

# Get realized connectivities

c_ex_ex = np.sum(connectivity[:n_ex, :n_ex])
c_ex_in = np.sum(connectivity[:n_ex, n_ex+1:-1])
c_in_ex = np.sum(connectivity[n_ex+1:-1, :n_ex])
c_in_in = np.sum(connectivity[n_ex+1:-1, n_ex+1:-1])

first_order_connectivity = [float(c_ex_ex) / n_ex ** 2, float(c_ex_in) / n_ex / n_in, float(c_in_ex) / n_ex / n_ex,
       float(c_in_in) / n_in ** 2]

print(first_order_connectivity)
plt.bar(range(4), first_order_connectivity)
plt.xticks(range(4), ["E-E", "E-I", "I-E", "I-I"])
plt.ylabel("Abundance")
plt.show()
exit()

print([np.sum(realize(p_conn)) for p_conn in conn_probs])

plt.imshow(realize(conn_probs[-1]))
plt.show()


print([np.prod(overlap.shape)* p for  overlap, p in zip(overlap_matrix.get_submatrices(), [p_ex_ex, p_ex_in, p_in_in, p_in_ex])])

print([np.sum(p_conn) for p_conn in conn_probs])