# -*- coding: utf-8 -*-
"""
Part of the GSM-Degroot package.
Copyright (C) 2020-2025, Ivan Conjeaud and Argyris Kalogeratos. This package is a free software distributed under the GPLv3 license.
"""


import scipy.stats as scs
import numpy as np
import networkx as nx
import math
import random

#In input give the network type: BA (Barabasi Albert), WS (Watts-Strogatz), grid
def initial_graph_non_SBM(N,network_type,network_parameters,mu,sigma,beta,Lambda, gamma):
    if network_type=='BA':
        m=network_parameters[0]
        G=nx.barabasi_albert_graph(N, m)
    elif network_type=='WS':
        k=network_parameters[0]
        p=network_parameters[1]
        G=nx.watts_strogatz_graph(N, k, p)
    elif network_type=='grid':
        dim=network_parameters[0]
        G=nx.grid_graph(dim)
    elif network_type=='ER':
        p=network_parameters[0]
        G=nx.erdos_renyi_graph(N,p)
    #This helps in case the nodes' labels aren't integers (typically for grids)
    G=nx.convert_node_labels_to_integers(G)    
    #Draw opinions
    X=np.random.normal(mu, sigma, N)
    #Draw initial state
    S=scs.bernoulli.rvs([math.exp(Lambda*x)/(1+math.exp(Lambda*x)) for x in X],size=N)
    #Draw reaction to information
    R=scs.bernoulli.rvs(beta, size=N)*gamma
    for k in range(len(R)):
        if R[k]==0:
            R[k]=-gamma
    #Generate influence matrix
    influence_matrix=np.zeros((N,N))
    for i in range(N):
        neighbours=[j for j in range(N) if (i,j) in G.edges()]
        neighbours.append(i)
        influence_on_i=[]
        for j in neighbours:
            influence_on_i.append(1/(len(neighbours)))
        #Introduce randomness in the weights
        K=10 #Number of iterations of the following process
        for _ in range(K):
            k=random.choice(range(len(neighbours)))
            l=random.choice(range(len(neighbours)))
            influence_on_i[k]=influence_on_i[k]/2
            influence_on_i[l]=influence_on_i[l]+influence_on_i[k]
        for j in neighbours:
            influence_matrix[i,j]=influence_on_i[neighbours.index(j)]
    return X, S, R, influence_matrix

#Function below gives opinion and states next period
def graph_next_period(X, S, R, influence_matrix, gamma, Lambda):
    N=len(X)
    #First we compute the GSM based on previous graph's state
    GSM=np.mean(S)*R
    #New opinion
    X_prime=np.matmul(influence_matrix, X)+GSM
    #New state
    S_prime=scs.bernoulli.rvs([math.exp(Lambda*x)/(1+math.exp(Lambda*x)) for x in X_prime],size=N)

    return X_prime, S_prime

def model_non_SBM(N,network_type, network_parameters,mu,sigma,beta,gamma,Lambda, T, M):
    #Initialize lists of opinion metrics
    max_end=[]
    min_end=[]
    pol_end=[]
    max_pol=[]
    max_pol_t=[]
    S_ts=[0 for _ in range(T)]
    #At each trial
    for k in range(M):
        #A new network is generated
        initial=initial_graph_non_SBM(N,network_type,network_parameters,mu,sigma,beta,Lambda, gamma)
        X, S, R, influence_matrix = initial[0], initial[1], initial[2], initial[3]
        maximal_polarization=max(X)-min(X)
        corresponding_time=0
        #The model is simulated
        for t in range(T):
            #Store the number of agents in state 1
            S_ts[t]=(k/(k+1))*S_ts[t]+np.mean(S)/(k+1)
            #Update
            update=graph_next_period(X, S, R, influence_matrix, gamma, Lambda)
            X, S = update[0], update[1]
            #Update max pol and corresponding time if needed
            if maximal_polarization< max(X)-min(X):
                maximal_polarization=max(X)-min(X)
                corresponding_time=t
        #Store final values for current trial
        max_op=max(X)
        min_op=min(X)
        max_end.append(max_op)
        min_end.append(min_op)
        pol_end.append(max_op-min_op)
        max_pol.append(maximal_polarization)
        max_pol_t.append(corresponding_time)
            
    return S_ts, np.mean(max_end), np.mean(min_end), np.mean(pol_end), np.mean(max_pol), np.mean(max_pol_t)    
    

def initial_graph_SBM(N,mus,sigmas,betas,Lambda,gamma, clusters, D,r):
    #Clusters is simply a list with the proportions of the population in each cluster: [x, y, z] with x+y+z=1
    #The probability matrix used for the generation of SBM has D on the diagonal, r outside
    
    #Transforms frequencies into numbers for the clusters' sizes
    sizes=[]
    for i in range(len(clusters)):
        sizes.append(int(clusters[i]*N))
    #Adjust the sizes if necessary
    if sum(sizes)!=N:
        for i in range(N-sum(sizes)):
            choice=random.choice([i for i in range(len(clusters))])
            sizes[choice]+=1
    #Define probs matrix
    probs=np.zeros((len(clusters),len(clusters)))
    #Fill the matrix
    for i in range(len(clusters)):
        probs[i][i]=D
        for j in range(len(clusters)):
            if j!=i:
                probs[i][j]=r
    #Generate the network
    G=nx.stochastic_block_model(sizes, probs)
    #In the graph generated, nodes in a cluster have consecutive labels
    X=np.array([])
    S=np.array([])
    R=np.array([])
    for k in range(len(sizes)):
        #Draw initial opinion
        X_cluster=np.random.normal(mus[k], sigmas[k], sizes[k])
        X=np.concatenate((X,X_cluster))
       #Draw initial state
        S_cluster=scs.bernoulli.rvs([math.exp(Lambda*x)/(1+math.exp(Lambda*x)) for x in X_cluster],size=sizes[k])
        S=np.concatenate((S,S_cluster))
        #Draw reaction to information
        R_cluster=scs.bernoulli.rvs(betas[k], size=sizes[k])
        for l in range(len(R_cluster)):
            if R_cluster[l]==0:
                R_cluster[l]=-gamma
            else:
                R_cluster[l]=gamma
        R=np.concatenate((R,R_cluster))
    #Generate influence matrix
    influence_matrix=np.zeros((N,N))
    for i in range(N):
        neighbours=[j for j in range(N) if (i,j) in G.edges()]
        neighbours.append(i)
        influence_on_i=[]
        for j in neighbours:
            influence_on_i.append(1/(len(neighbours)))
        #Introduce randomness in the weights
        K=10 #Number of iterations of the following process
        for _ in range(K):
            k=random.choice(range(len(neighbours)))
            l=random.choice(range(len(neighbours)))
            influence_on_i[k]=influence_on_i[k]/2
            influence_on_i[l]=influence_on_i[l]+influence_on_i[k]
        for j in neighbours:
            influence_matrix[i,j]=influence_on_i[neighbours.index(j)]
    return X, S, R, influence_matrix

def model_SBM(N,mus,sigmas,betas,gamma,Lambda, T, M, clusters, D, r):
    #Initialize lists of opinion metrics
    max_end=[]
    min_end=[]
    pol_end=[]
    max_pol=[]
    max_pol_t=[]
    S_ts=[0 for _ in range(T)]
    #At each trial
    for k in range(M):
        #A new network is generated
        initial=initial_graph_SBM(N,mus,sigmas,betas,Lambda,gamma, clusters, D,r)
        X, S, R, influence_matrix = initial[0], initial[1], initial[2], initial[3]
        maximal_polarization=max(X)-min(X)
        corresponding_time=0
        #The model is simulated
        for t in range(T):
            #Store the number of agents in state 1
            S_ts[t]=(k/(k+1))*S_ts[t]+np.mean(S)/(k+1)
            #Update
            update=graph_next_period(X, S, R, influence_matrix, gamma, Lambda)
            X, S = update[0], update[1]
            #Update max pol and corresponding time if needed
            if maximal_polarization< max(X)-min(X):
                maximal_polarization=max(X)-min(X)
                corresponding_time=t
        #Store final values for current trial
        max_op=max(X)
        min_op=min(X)
        max_end.append(max_op)
        min_end.append(min_op)
        pol_end.append(max_op-min_op)
        max_pol.append(maximal_polarization)
        max_pol_t.append(corresponding_time)
            
    return S_ts, np.mean(max_end), np.mean(min_end), np.mean(pol_end), np.mean(max_pol), np.mean(max_pol_t)

#The following allows to cope with potential numerical errors
def model_SBM_robust(N,mus,sigmas,betas,gamma,Lambda, T, M, clusters, D, r):
    #Initialize lists of opinion metrics
    max_end=[]
    min_end=[]
    pol_end=[]
    max_pol=[]
    max_pol_t=[]
    S_ts=[0 for _ in range(T)]
    #At each trial: try to compute the model until all metrics are generated, try again otherwise
    for k in range(M):
        success=False
        while success==False:
            try:
        #A new network is generated
                initial=initial_graph_SBM(N,mus,sigmas,betas,Lambda,gamma, clusters, D,r)
                X, S, R, influence_matrix = initial[0], initial[1], initial[2], initial[3]
                maximal_polarization=max(X)-min(X)
                corresponding_time=0
        #The model is simulated
                for t in range(T):
            #Store the number of agents in state 1
                    S_ts[t]=(k/(k+1))*S_ts[t]+np.mean(S)/(k+1)
            #Update
                    update=graph_next_period(X, S, R, influence_matrix, gamma, Lambda)
                    X, S = update[0], update[1]
            #Update max pol and corresponding time if needed
                    if maximal_polarization< max(X)-min(X):
                        maximal_polarization=max(X)-min(X)
                        corresponding_time=t
        #Store final values for current trial
                max_op=max(X)
                min_op=min(X)
                max_end.append(max_op)
                min_end.append(min_op)
                pol_end.append(max_op-min_op)
                max_pol.append(maximal_polarization)
                max_pol_t.append(corresponding_time)
                success=True
            except:
                success=False
            
    return S_ts, np.mean(max_end), np.mean(min_end), np.mean(pol_end), np.mean(max_pol), np.mean(max_pol_t)
    