import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from sklearn import cluster, covariance, manifold
from community import community_louvain as louvain
import matplotlib.cm as cm
import networkx as nx
import networkx.algorithms.community as nxcom
from importlib import reload
import csv
import os
import re
%matplotlib inline

df_prices = pd.read_csv("SP500_prices.csv", index_col = 0)
df_indices = pd.read_csv("indices.csv")
df_SP500 = pd.read_csv("SP500.csv",index_col = 0)

df_prices.describe()

df_prices.head()

df_prices_adj = df_prices[['symbol','date', 'adjusted']]
df_prices_adj.columns = ['symbol','date','price']
df_prices_adj.tail()

df_indices.describe()

df_indices.head()

df_indices = df_indices[["symbol","date","price"]]
df_indices.tail()

df_SP500.head()

df = pd.concat([df_prices_adj,df_indices])
df['date'] = pd.to_datetime(df['date'], format='%Y%m%d', errors='ignore')
df.set_index(['date','symbol'],inplace=True)
df=df.unstack()['price']
df.fillna(method='bfill',inplace=True)
df

df = (df-df.mean())/df.std()

df.describe()

%matplotlib inline
fig, ax1 = plt.subplots(figsize=(20, 15))
df.iloc[:,:20].plot(ax=ax1, legend=False)
plt.tight_layout()
plt.show()

df_delta = df.copy()
for column in df_delta.columns.values.tolist():
    df_delta[column] = df_delta[column]- df_delta[column].shift(1)
df_delta.iloc[0]=0
df_delta

%matplotlib inline
fig, ax1 = plt.subplots(figsize=(20, 15))
df_delta.plot(ax=ax1, legend=False)
plt.tight_layout()
plt.show()

def calculate_corr(df_stock_returns, returns_window, corr_window_size, corr_method):
    stocks_cross_corr_dict = {}
    x_days = []
    y_mean_corr = []        
    for i in range(returns_window,len(df_stock_returns),corr_window_size):
        dic_key = i
        stocks_cross_corr_dict[dic_key]=df_stock_returns.iloc[i:(i+W)].corr(method='pearson')
        stocks_cross_corr_dict[dic_key].fillna(0,inplace=True)
        x_days.append(dic_key)
        y_mean_corr.append(np.mean([abs(j) for j in stocks_cross_corr_dict[dic_key].values.flatten().tolist()]))        
    return stocks_cross_corr_dict, x_days,y_mean_corr

%matplotlib inline
start = 21
end = 126
step = 21;
plt.figure(figsize=(20, 10))

for t in range(start, end, step):
    x_days = []
    y_mean_corr = []
    W = t
    _, x_days, y_mean_corr = calculate_corr(df,1,W, 'pearson')
    plt.plot(x_days, y_mean_corr)
    plt.xlabel('Days')
    plt.ylabel('Mean Correlation')
    l = list(range(start, end, step))
    plt.legend(l, loc='upper left')     

plt.show()

# craetes a graph from correlation matrix
cor_matrix = df.corr()

cor_matrix

mat_pos = cor_matrix[cor_matrix>=0]
mat_pos = mat_pos.fillna(0)
symbols = cor_matrix.index.values
mat_pos = np.asmatrix(mat_pos)
G_pos = nx.from_numpy_matrix(mat_pos)
G_pos = nx.relabel_nodes(G_pos,lambda x: symbols[x])
G_pos.remove_edges_from(nx.selfloop_edges(G_pos))

mat_neg = cor_matrix[cor_matrix<0]
mat_neg = mat_neg.fillna(0)
mat_neg = abs(mat_neg)
symbols = cor_matrix.index.values
mat_neg = np.asmatrix(mat_neg)
G_neg = nx.from_numpy_matrix(mat_neg)
G_neg = nx.relabel_nodes(G_neg,lambda x: symbols[x])
G_neg.remove_edges_from(nx.selfloop_edges(G_neg))

list(G_pos.edges(data=True))[:5], list(G_neg.edges(data=True))[:5]

symbol =  df.columns
df_sector = df_SP500[df_SP500['symbol'].isin(symbol)]
df_sector = df_sector[["symbol","sector"]]
tmp = pd.DataFrame({"symbol":df_indices["symbol"].unique(), "sector":"Macroeconomic Indices"})
df_sector = pd.concat([df_sector,tmp],ignore_index=True)
df_sector['sec_idx']= 0
for i,sec in enumerate(df_sector["sector"].unique()):
    df_sector.loc[df_sector['sector']==sec,'sec_idx'] = i+1

df_sector

nx.set_node_attributes(G_pos, df_sector.set_index('symbol')['sec_idx'],'sec_idx')
nx.set_node_attributes(G_neg, df_sector.set_index('symbol')['sec_idx'],'sec_idx')

def set_threshold(G,threshold):
    edges_rm = list(filter(lambda e: abs(e[2]) < threshold, (e for e in G.edges.data('weight'))))    
    ids_rm = list(e[:2] for e in edges_rm)
    H = G.copy()
    H.remove_edges_from(ids_rm)
    return H

H_pos = set_threshold(G_pos,0.5)
H_neg = set_threshold(G_neg,0.5)
list(H_pos.edges(data=True))[:5], list(H_neg.edges(data=True))[:5]

# grid search
# for positive weights
np.random.seed(2021)
t=0
for cor_thresold in np.linspace(0.8,0.85,20):
    H_pos = set_threshold(G_pos,cor_thresold)
    partition = louvain.best_partition(H_pos)
    modularity = louvain.modularity(partition, H_pos)
    values = [partition.get(node) for node in H_pos.nodes()]
    communities = []
    tmp = list(partition.items())
    for i in range(len(set(values))):
        communities.append([n for n,c in tmp if c==i])
    sum_comm_nodes = 0
    k=0
    print("{}th Total number of Communities = {}".format(t ,len(communities)))
    for i, comm_nodes in enumerate(communities):
        if len(comm_nodes)>=10:
            k+=1
            print('community {}th: '.format(i),len(comm_nodes))
        sum_comm_nodes+=len(comm_nodes)
    t+=1
    print('n_big_communities: ',k)
# best partition with 5 big communities at cor_threshold = np.linspace(0.8,0.85,20)[8] 

np.random.seed(2021)
cor_thresold = np.linspace(0.8,0.85,20)[8]
H_pos = set_threshold(G_pos,cor_thresold)
partition = louvain.best_partition(H_pos)
values = [partition.get(node) for node in H_pos.nodes()]
communities = []
tmp = list(partition.items())
for i in range(len(set(values))):
    communities.append([n for n,c in tmp if c==i])
sum_comm_nodes = 0
k=0
print("Total number of Communities = {}".format(len(communities)))
for i, comm_nodes in enumerate(communities):
    if len(comm_nodes)>=10:
        k+=1
        print('community {}th: '.format(i),len(comm_nodes))
    sum_comm_nodes+=len(comm_nodes)
print('n_big_communities: ',k)

nx.set_node_attributes(H_pos,partition,'community')
values = [partition.get(node) for node in H_pos.nodes()]
plt.figure(figsize=(10,10))
nx.draw_spring(H_pos, cmap = plt.get_cmap('jet'), node_color = values, node_size=30, with_labels=False)

np.random.seed(2021)
# for negative weights
t=0
for cor_thresold in np.linspace(0.6,0.65,20):
    H_neg = set_threshold(G_neg,cor_thresold)
    partition = louvain.best_partition(H_neg)
    modularity = louvain.modularity(partition, H_neg)
    values = [partition.get(node) for node in H_neg.nodes()]
    communities = []
    tmp = list(partition.items())
    for i in range(len(set(values))):
        communities.append([n for n,c in tmp if c==i])
    sum_comm_nodes = 0
    k=0
    
    print("{}th Total number of Communities = {}".format(t ,len(communities)))
    for i, comm_nodes in enumerate(communities):
        if len(comm_nodes)>=10:
            k+=1
            print('community {}th: '.format(i),len(comm_nodes))
        sum_comm_nodes+=len(comm_nodes)
    t+=1
    print('n_big_communities: ',k)
# best partition with 5 big communities at cor_threshold = np.linspace(0.6,0.65,20)[11] 

np.random.seed(2021)
cor_thresold = np.linspace(0.6,0.65,20)[11] 
H_neg_otim = set_threshold(G_neg,cor_thresold)
partition = louvain.best_partition(H_neg_otim)
modularity = louvain.modularity(partition, H_neg)
values = [partition.get(node) for node in H_neg.nodes()]
communities = []
tmp = list(partition.items())
for i in range(len(set(values))):
    communities.append([n for n,c in tmp if c==i])
sum_comm_nodes = 0
k=0
print("Total number of Communities = {}".format(len(communities)))
for i, comm_nodes in enumerate(communities):
    if len(comm_nodes)>=10:
        k+=1
        print('community {}th: '.format(i),len(comm_nodes))
    sum_comm_nodes+=len(comm_nodes)
print('n_big_communities: ',k)

nx.set_node_attributes(H_neg, partition,'community')
values = [partition.get(node) for node in H_neg.nodes()]
plt.figure(figsize=(10,10))
nx.draw_spring(H_neg, cmap = plt.get_cmap('jet'), node_color = values, node_size=30, with_labels=False)

nx.write_graphml(H_pos, "model_H_pos.graphml")
nx.write_graphml(H_neg, "model_H_neg.graphml")