# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# Evaluate different ising influence strategies on the Pokec dataset for a range of different linear field gradients.
#
# Created on: 28/09/20

# +
import networkx as nx
import numpy as np
import pandas as pd
import ast
import seaborn as sns
import matplotlib.pyplot as plt
import tqdm
import random
import itertools
import matplotlib
from scipy import stats
import math

from ising_block_level_influence import N_Block_sbm_class as NBlock
from ising_block_level_influence import projection_simplex as proj
from ising_block_level_influence import mean_field_IIM
from spatial_spin_monte_carlo import spatial_spin_monte_carlo as Spins
import Pokec_processing as  PokProc

from pokec_utils import *
# -

# ### Read in processed data
#
# This data has been generated using:
#
# 1. **make_Pokec_reduced_profiles.ipynb** - which then feeds data into:
#
# 2. The pre-processing script **'make_bratislava_graph_and_blocks.py'**.

graph = nx.read_graphml('Data/Bratislava_graph.graphml')
beta_c = Spins.crit_beta_sparse(graph)
bratislava_profiles = pd.read_csv("Data/bratislava_profiles.csv")
coupling_graph = nx.read_graphml('Data/Bratislava_coupling.graphml')
block_data = pd.read_csv('Data/block_info.csv',converters={'Block' : ast.literal_eval})
mean_block_ages = list(block_data['average_age'])
block_sizes = list(block_data['block_size'])
block_names = list(block_data['block_name'])
block_data['age_group'] = [ 'ages_' + k.split('_')[-1] for k in list(block_data['block_name'])]
block_data.head()


def linear_field(x : np.ndarray,gradient :float) :
    return gradient*x
#Scale ages to [-1,1]:
rescaled_ages = [ (k-np.mean(mean_block_ages))/(max(mean_block_ages)-min(mean_block_ages)) for k in mean_block_ages ]

bratislava_profiles_indices = bratislava_profiles.reset_index()
groups = [ bratislava_profiles_indices.loc[bratislava_profiles_indices['block']==block] for block in block_names]
groups_node_ids = [list(k['index']) for k in groups]


# ## Sweep over $\beta$ values

# +
def check_group(input_val,group_label) :
    if input_val == group_label :
        return 1.0
    else :
        return 0.0
    
def mean_and_se(values) :
    return f"{np.mean(values)} +/- {stats.sem(values)} "

#Seed the random number generators:
seed = 1
random.seed(seed)
np.random.seed(seed)

Field_Budget = 2500.0
T=10000
T_Burn = 2*300000
Samples = 15
sample_frac = 1.0 # for snapshot control
figure='4d'

if figure == '3c' :
    # figure 3c params
    init_sign=1.0
    beta_factor_vals = [(10**k) for k in np.linspace(-1,1.8,12)]
    grad_vals = [0.0]
    con_names=['block','unif','full']
    save_path='Data/Pokec_control_eval_as_beta.csv'
elif figure == '4c' :
    # Figure 4c params
    init_sign=-1.0
    beta_factor_vals=[8.0]
    grad_vals =  np.arange(-10.0,12.0,2.0)
    con_names = ['no con','unif','full','block','sv','nc','snapshot']
    save_path='Data/Pokec_control_eval_negative.csv'
    
elif figure == '4d' :
    # figure 4d params
    init_sign=1.0
    beta_factor_vals=[8.0]
    grad_vals =  np.arange(-10.0,12.0,2.0)
    con_names = ['no con','unif','full','block','sv','nc','snapshot']
    save_path='Data/Pokec_control_eval_positive.csv'


eval_data = pd.DataFrame()
control_data = pd.DataFrame()
control_dict={}

for gradient in tqdm.tqdm_notebook(grad_vals) : 
    
    age_field = [linear_field(a,gradient) for a in rescaled_ages ]
    age_field_map = {k:j for k,j in zip(list(block_data['age_group']),age_field)}

    # Block MF setup:
    coupling_graph = nx.convert_node_labels_to_integers(coupling_graph)
    block_graph_ising_system = mean_field_IIM.mean_field_ising_system(coupling_graph,age_field)
    block_graph_ising_system.gamma = 1.0
    block_graph_ising_system.tol = 1E-5
    block_graph_ising_system.max_mf_fp_iterations = 10000
    block_graph_ising_system.mf_fp_init_state = init_sign*np.ones(len(block_sizes))
    block_graph_ising_system.mf_fp_noisy = False

    # Full graph Mf setup
    background_field = np.asarray([age_field_map[k] for k in list(bratislava_profiles['age_group'])])
    relab_graph = nx.relabel.convert_node_labels_to_integers(graph)
    full_graph_ising_system = mean_field_IIM.mean_field_ising_system(relab_graph,background_field)
    full_graph_ising_system.mf_fp_init_state = init_sign*np.ones(len(graph))
    full_graph_ising_system.mf_fp_noisy = False


    for beta_factor in tqdm.tqdm_notebook(beta_factor_vals) : 

        beta = beta_c*beta_factor

        # Block gradient computation: 
        m_block = block_graph_ising_system.mf_magnetization(age_field,beta)
        mag_gradient = block_graph_ising_system.mf_magnetization_gradient(m_block, beta)
        mag_grad_map = { i:j for i,j in zip(block_names,mag_gradient)}
        block_names_list = list(bratislava_profiles['block'])
        block_mf_gradient = np.asarray([mag_grad_map[k] for k in block_names_list])

        # Full gradient computation:
        if 'full' in con_names :
            m_full = full_graph_ising_system.mf_magnetization(background_field,beta)
            mag_grad_full = full_graph_ising_system.mf_sparse_magnetisation_gradient(m_full, beta)
            H_full = np.sum(mag_grad_full)
            mag_grad_full_at_field_budget = (Field_Budget/H_full)*mag_grad_full
            control_dict['full']=mag_grad_full_at_field_budget

        # Compute the controls:
    
        H_block = np.sum(block_mf_gradient)
        uniform_control = (Field_Budget / len(graph)) * np.ones(len(graph))
        
        mag_grad_block_at_field_budget = (Field_Budget/H_block)*block_mf_gradient
        negative_cancelling_field = []
        for field in background_field :
            if field < 0.0:
                negative_cancelling_field.append(-1.0*field)
            else:
                negative_cancelling_field.append(0.0)
        negative_cancelling_field = np.asarray(negative_cancelling_field)
        negative_cancelling_field = (Field_Budget/np.sum(negative_cancelling_field))*np.asarray(negative_cancelling_field) 
        swing_voter_nodes = np.asarray([check_group(k,'ages_22-28') for k in list(bratislava_profiles['age_group']) ])
        H_SV = np.sum(swing_voter_nodes)
        swing_voter_control =  (Field_Budget/H_SV)*swing_voter_nodes
        
        #Snapshot control
        initial_state = init_sign*np.ones(len(graph))
        block_snapshot = Run_MonteCarlo_Snapshot(relab_graph,groups_node_ids ,0, beta_factor,beta_c ,T_Burn=T_Burn,addition_control=None,sampling_method="Metropolis",full_graph_field=background_field,initial_state=initial_state,frac_to_sample=sample_frac)
        snapshot_gradient = [ (1.0-k**2) for k in block_snapshot ] 
        snapshot_grad_map = { i:j for i,j in zip(block_names,snapshot_gradient)}
        snapshot_influence = np.asarray([snapshot_grad_map[k] for k in block_names_list])
        #Normalize:
        H_snapshot = np.sum(snapshot_influence)
        snapshot_influence = (Field_Budget/H_snapshot)*snapshot_influence
        
        control_dict['no con']=np.zeros(len(graph))
        control_dict['unif']=uniform_control
        control_dict['block']=mag_grad_block_at_field_budget
        control_dict['sv']=swing_voter_control
        control_dict['nc']=negative_cancelling_field
        control_dict['snapshot']=snapshot_influence
        
        for s in tqdm.tqdm_notebook( range(Samples) )  : 

            magnetisations_dict={}
            for con_name in con_names : 
                initial_state = init_sign*np.ones(len(graph))
                control_w_background=np.asarray([i+j for i,j in zip(control_dict[con_name],background_field)])
                magnetisations=Run_MonteCarlo_Average(relab_graph, T, beta_factor,beta_c, T_Burn=T_Burn,addition_control=None,sampling_method="Metropolis",full_graph_field=control_w_background,initial_state=initial_state)
                magnetisations_dict[con_name]=magnetisations
            
            
            means=[]
            ses=[]
            #Loop in same order as cont_names
            # control_list = [no_control,unif_control,full_sus_control_mags,block_sus_control_mags,sv_control_mags,nc_control_mags,snapshot_control_mags]
            #for k in control_list :
            for k in list(magnetisations_dict.values()):
                print(mean_and_se(k))
                means.append(np.mean(k))
                ses.append(stats.sem(k))

            eval_data = eval_data.append(pd.DataFrame({'control':list(magnetisations_dict.keys()),'magnetisation':means,'mag_se':ses,'beta_factor':beta_factor*np.ones(len(means)),'gradient':gradient*np.ones(len(means))}))
            eval_data.to_csv(save_path)
# -

eval_data



StarcoderdataJupyterScriptsDedupFiltered