import collections

import pandas as pd
import numpy as np
import seaborn as sns
sns.set()

COUNTRY_CODES = ['AE', 'BH', 'EG', 'IR', 'KW', 'LB', 'OM', 'SA', 'SY', 'TR', 'QA', 'YE']
SHEETS_TO_CODES = {
    'Iran': 'IR', 
    'Lebanon': 'LB',
    'Hezbullah': 'HEZ',
    'KSA': 'SA',
    'Syria': 'SY',
    'Bahrain': 'BH',
    'Kuwait': 'KW',
    'Yemen': 'YE',
    'UAE': 'AE',
    'Qatar': 'QA',
    'Egypt': 'EG',
    'Ikhwan': 'MB',
    'Turkey': 'TR',
    'Kurds': 'KURD',
    'Regional websites': 'MENA',
}

pairwise = pd.DataFrame(index=SHEETS_TO_CODES.values(), columns=SHEETS_TO_CODES.values(), dtype=np.float32)

data = pd.read_excel('filtering_data.xlsx', None)
for sheet, code in SHEETS_TO_CODES.items():
    values = data[sheet].loc[:,data[sheet].columns.intersection(COUNTRY_CODES)].sum() / data[sheet].loc[:,data[sheet].columns.intersection(COUNTRY_CODES)].count()
    pairwise[code].update(values)
pairwise

	IR	LB	HEZ	SA	SY	BH	KW	YE	AE	QA	EG	MB	TR	KURD	MENA
IR	NaN	0.044872	0.000000	0.148718	0.027027	0.071429	0.043478	0.045643	0.0	0.000000	0.130435	1.000000	0.000000	0.375	0.320000
LB	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
HEZ	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
SA	0.119565	0.032051	0.727273	NaN	0.054054	0.000000	0.000000	0.124481	0.0	0.685185	0.034783	0.083333	0.052632	0.000	0.360000
SY	0.000000	0.044872	0.000000	0.025641	NaN	0.000000	0.043478	0.000000	0.0	0.037037	0.000000	1.000000	0.052632	0.000	0.346667
BH	0.032609	0.000000	0.181818	0.000000	0.000000	NaN	0.000000	0.000000	0.0	0.148148	0.000000	NaN	0.000000	0.000	NaN
KW	0.000000	0.006410	0.000000	0.000000	NaN	NaN	NaN	0.000000	0.0	NaN	0.000000	NaN	0.000000	0.000	0.053333
YE	0.000000	0.000000	0.000000	0.102564	0.000000	0.000000	0.000000	NaN	0.1	NaN	0.000000	NaN	0.105263	0.000	0.173333
AE	0.032609	0.000000	0.454545	NaN	0.054054	0.000000	0.000000	0.012448	NaN	0.314815	0.069565	0.416667	0.000000	0.000	0.306667
QA	0.000000	0.000000	0.000000	NaN	NaN	NaN	NaN	0.000000	0.0	NaN	0.000000	NaN	0.000000	0.000	NaN
EG	0.010870	NaN	0.000000	NaN	0.000000	0.000000	NaN	0.000000	0.0	0.222222	NaN	NaN	0.210526	0.000	0.133333
MB	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
TR	0.000000	NaN	NaN	NaN	0.000000	0.000000	0.000000	0.000000	0.0	NaN	0.000000	NaN	0.000000	0.375	NaN
KURD	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
MENA	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN

sns.heatmap(pairwise, square=True, cmap='RdYlGn_r', center=0.1)

<matplotlib.axes._subplots.AxesSubplot at 0x7f511ca4ce48>

conflict_actors = collections.defaultdict(list)
for row in data['Conflict Factions'].loc[:,['Conflict', 'Actor', 'Faction']].iterrows():
    if row[1]['Actor'] != 'SD':
        conflict_actors[row[1]['Conflict']].append(row[1]['Actor'])
conflict_actors

defaultdict(list,
            {'SIT': ['SA', 'AE', 'BH', 'IR'],
             'CIS': ['IR', 'HEZ', 'SA', 'QA', 'TR'],
             'H': ['HEZ', 'LB', 'SA', 'KW', 'OM', 'BH', 'QA', 'AE'],
             'WIY': ['YE', 'SA', 'IR'],
             'GC': ['SA', 'AE', 'BH', 'EG', 'QA'],
             'MB': ['SA', 'AE', 'QA'],
             'K': ['TR', 'KURD']})

conflict_names = data['Conflicts'].set_index('Code').to_dict()['Conflict']

def plot_conflict(conflict, conflict_actors, pairwise_filtering):
    actors = conflict_actors[conflict]
    ax = sns.heatmap(pairwise.loc[actors, actors], square=True, cmap='RdYlGn_r', vmin=0, vmax=0.2)
    ax.set(xlabel='Sites Friendly To', ylabel='Accessing From', title=conflict_names[conflict])

conflict = 'SIT'
plot_conflict(conflict, conflict_actors, pairwise)

conflict = 'CIS'
plot_conflict(conflict, conflict_actors, pairwise)

conflict = 'H'
plot_conflict(conflict, conflict_actors, pairwise)

/home/jclark/miniconda3/envs/im/lib/python3.6/site-packages/ipykernel_launcher.py:3: FutureWarning: 
Passing list-likes to .loc or [] with any missing label will raise
KeyError in the future, you can use .reindex() as an alternative.

See the documentation here:
https://pandas.pydata.org/pandas-docs/stable/indexing.html#deprecate-loc-reindex-listlike
  This is separate from the ipykernel package so we can avoid doing imports until

conflict = 'WIY'
plot_conflict(conflict, conflict_actors, pairwise)

conflict = 'GC'
plot_conflict(conflict, conflict_actors, pairwise)

conflict = 'MB'
plot_conflict(conflict, conflict_actors, pairwise)

conflict = 'K'
plot_conflict(conflict, conflict_actors, pairwise)