import pandas as pd
import numpy as np
import networkx as nx
from scipy.linalg import block_diag
import matplotlib.pyplot as plt
%matplotlib inline

np.random.seed(seed=42)

def compute_stationary_distribution(P):
    """Returns the stationary distribution of the markov chain.

       Solves x = xP, where x is the stationary distribution.
       
       x - xP = 0 <-> x(I - P) = 0, such that sum(x) = 1
       
       :param P: The transition matrix of the markov chain.
    """
    n = P.shape[0]
    # formulates as a x = b
    a = np.eye(n) - P
    a = np.vstack((a.T, np.ones(n)))
    b = np.matrix([0] * n + [1]).T
    
    return np.linalg.lstsq(a, b)[0]

def build_graph(rankings):
    rank_graph = nx.DiGraph()
    for ranking, ic in rankings:
        for node in ranking.values():
            rank_graph.add_node(node)
    for ranking, ic in rankings:
        for rank_up in range(1, len(ranking)):
            for rank_dn in range(rank_up + 1, len(ranking) + 1):
                u = ranking[rank_up]
                v = ranking[rank_dn]

                if u in rank_graph and v in rank_graph[u]:
                    rank_graph.edges[u, v]['weight'] += ic
                else:
                    rank_graph.add_edge(u, v, weight=ic)
                    
    return rank_graph

def build_markov_chain(rank_graph):
    for node in rank_graph.nodes():
        sum_weights = 0

        for adj_node in rank_graph[node]:
            sum_weights += rank_graph[node][adj_node]['weight']

        for adj_node in rank_graph[node]:
            rank_graph.edges[node, adj_node]['weight'] /= sum_weights
            
    return rank_graph

def compute_transition_matrix(rank_graph):
    nb_nodes = len(rank_graph.nodes())
    transitions = pd.DataFrame(np.zeros((nb_nodes, nb_nodes)),
                                     index=rank_graph.nodes(),
                                     columns=rank_graph.nodes())
    for node in rank_graph.nodes():
        for adj_node in rank_graph[node]:
            transitions.set_value(node, adj_node, rank_graph[node][adj_node]['weight'])
            
    return transitions.div(transitions.sum(axis=1), axis=0)

def random_walk(markov_chain, counts, max_random_steps=1000):
    cur_step = 0
    current_node = graph_nodes[np.random.randint(len(graph_nodes))]
    while cur_step < max_random_steps:
        counts[current_node] += 1
        neighbours = [(adj, markov_chain[current_node][adj]['weight'])
                      for adj in markov_chain[current_node]]

        adj_nodes = [adj for adj, w in neighbours]
        adj_weights = [w for adj, w in neighbours]

        if not adj_nodes:
            return counts
        
        current_node = np.random.choice(adj_nodes, p=adj_weights)

        cur_step += 1
        
    return counts

nb_assets = 250
nb_observations = 1 * 252

residual_returns = np.random.multivariate_normal(
    np.zeros(nb_assets), cov=np.eye(nb_assets), size=nb_observations)

residual_returns = pd.DataFrame(residual_returns)

ranked_cumulative_residual_returns = (residual_returns
                                      .sum()
                                      .rank(ascending=False))

nb_experts_available = [
    5, 10, 15, 20, 25, 30, 40, 50,
    60, 70, 80, 90, 100, 125, 150,
    175, 200, 250, 300, 350, 400,
    450, 500,
]

results = {}

for nb_experts in nb_experts_available:

    prediction_rankings = []

    for expert in range(nb_experts):
        confidence = np.random.randint(1, 61) / 100
        pct_coverage = np.random.randint(2, 51) / 100
        expert_coverage = np.random.choice(nb_assets, replace=False,
                                           size=int(nb_assets * pct_coverage))

        predictions = []
        for node in expert_coverage:
            prediction = (confidence * residual_returns[node] + 
                          np.sqrt(1 - confidence**2) * np.random.normal(
                              0, np.std(residual_returns[node]), nb_observations))
            predictions.append(prediction)

        expert_predictions = pd.DataFrame(np.array(predictions).T,
                                          columns=expert_coverage)
        ranked_cumulated_predictions = (
            expert_predictions
            .sum()
            .rank(ascending=False))

        prediction_rankings.append(
            ({value: key for key, value in dict(ranked_cumulated_predictions).items()},
             confidence))
        
    rank_graph = build_graph(prediction_rankings)
    markov_chain = build_markov_chain(rank_graph)

    transition_matrix = compute_transition_matrix(rank_graph)
    stationary_distribution = compute_stationary_distribution(
        transition_matrix.values)

    eq_ranking = (pd.DataFrame(stationary_distribution,
                       index=transition_matrix.columns)[0]
                  .rank(ascending=True))

    graph_nodes = list(markov_chain.nodes())

    max_random_steps = 10000
    counts = {node: 0 for node in graph_nodes}

    nb_jumps = 100
    for jump in range(nb_jumps):
        counts = random_walk(markov_chain, counts, max_random_steps)

    # true ranking
    true_ranking = sorted([(key, value)
                           for key, value in dict(ranked_cumulative_residual_returns).items()],
                          key=lambda x: x[1])
    true_ranking = pd.Series([count for node, count in true_ranking],
                         index=[node for node, count in true_ranking])

    # rank graph method
    final_ranking = sorted([(node, count) for node, count in counts.items()],
                           key=lambda x: x[0])

    graph_ranking = (pd.Series([count for node, count in final_ranking],
                              index=[node for node, count in final_ranking])
                     .rank(ascending=True))

    # average ranking
    rankings = pd.DataFrame(np.nan,
                            columns=sorted(graph_nodes),
                            index=range(len(prediction_rankings)))
    for idrank in range(len(prediction_rankings)):
        for rank, asset in prediction_rankings[idrank][0].items():
            rankings[asset].loc[idrank] = rank / len(prediction_rankings[idrank][0])

    sum_weights = sum([prediction_rankings[idrank][1]
                       for idrank in range(len(prediction_rankings))])
    average_rankings = (rankings
                        .multiply([prediction_rankings[idrank][1] / sum_weights
                                   for idrank in range(len(prediction_rankings))],
                                  axis='rows')
                        .sum(axis=0))

    average_ranking = average_rankings.rank(ascending=True)

    results[nb_experts] = {
        'avg': average_ranking.corr(true_ranking, method='spearman'),
        'graph': graph_ranking.corr(true_ranking, method='spearman'),
        'eq': eq_ranking.corr(true_ranking, method='spearman'),
        'corr_graph_eq': eq_ranking.corr(graph_ranking, method='spearman'),
        'corr_graph_avg': graph_ranking.corr(average_ranking, method='spearman')}

    print('number of rankings available:', nb_experts, '\n')
    print(results[nb_experts])
    print('--------------------------------------\n\n')

plt.figure(figsize=(20, 10))
plt.plot(nb_experts_available, 
         [results[nb_experts]['avg'] for nb_experts in nb_experts_available],
         '--o', label='average')
plt.plot(nb_experts_available, 
         [results[nb_experts]['graph'] for nb_experts in nb_experts_available],
         '--o', label='network-based')
plt.plot(nb_experts_available, 
         [results[nb_experts]['eq'] for nb_experts in nb_experts_available],
         '--o', label='stat. distr.')
plt.title('Correlation to the true ranking', fontsize=24)
plt.legend(fontsize=24)
plt.xlabel('Number of rankings available', fontsize=16)
plt.ylabel('Correlation', fontsize=16)
plt.show()

plt.figure(figsize=(20, 10))
plt.plot(nb_experts_available, 
         [results[nb_experts]['corr_graph_avg'] for nb_experts in nb_experts_available],
         '--o', label='corr(graph, average)')
plt.plot(nb_experts_available, 
         [results[nb_experts]['corr_graph_eq'] for nb_experts in nb_experts_available],
         '--o', label='corr(graph, equation)')
plt.title('Correlation between methods', fontsize=24)
plt.legend(fontsize=24)
plt.xlabel('Number of rankings available', fontsize=16)
plt.ylabel('Correlation', fontsize=16)
plt.show()