%matplotlib inline

from hypergraph import generators

HG = generators.generic_hypergraph(6, ((3, 4), (4, 3), (5, 1)))

import numpy as np

from functools import partial
from matplotlib import pyplot as plt

from numpy import linalg as LA

from hypergraph import generators
from hypergraph.analytical import prediction
from hypergraph.diffusion_engine import DiffusionEngine
from hypergraph import utils
from hypergraph.markov_diffusion import create_markov_matrix_model_nodes
from hypergraph.markov_diffusion import create_markov_matrix_model_hyper_edges

def hypergraph_analytical_edges(HG):
    phis = []
    number_of_nodes = len(HG.nodes())
    all_phis = 0
    for edge in HG.hyper_edges():
        phis.append(len(edge))
        all_phis += len(edge)
    
    pis = [phi / all_phis for phi in phis]
    return pis



def show_models(HG):
    all_models={"node": {"analytical": partial(prediction, model='hypergraph'),
                    "numerical": create_markov_matrix_model_nodes,
                    "name": "node",
                   },
            "hyperedges": {"analytical": hypergraph_analytical_edges,
                           "numerical": create_markov_matrix_model_hyper_edges,
                           "name": "hyperedges",
                           }
    }
    t_max = 100000
    number_of_walkers = 1
    
    
    model = all_models['node']
    
    def show_model(model):
        markov_matrix = model["numerical"](HG)
        
        
        t_per_walker = int(t_max / number_of_walkers)

        engine = DiffusionEngine(markov_matrix, t_per_walker=t_per_walker)

        freqs_matrix = LA.matrix_power(markov_matrix, 40)[0]
        frequencies, states = engine.simulate(t_max)

        connected_elements = [freq[0] + 1 for freq in frequencies]

        if model['name'] == "node":
            missing_nodes = set(HG.nodes()) - set(connected_elements)
            for missing_node in missing_nodes:
                frequencies.append((missing_node, 0))
                
        frequencies = [(node, frequency) for node, frequency in frequencies]
        frequencies.sort(key=lambda x: x[0])
        xs, ys = zip(*frequencies)

        ys = np.array(ys, dtype='float')
        ys /= sum(ys)
        width = 0.2  

        ys_prediction = model["analytical"](HG)

        plt.figure(figsize=(10, 8))
        plt.bar(xs, ys, width=width, color='crimson', label='Simulated')

        plt.bar(np.array(xs) + width, freqs_matrix, width=width, label='Traversal matrix to N')

        plt.bar(np.array(xs) + 2 * width, ys_prediction, width=width, color='#dcccdd', label='Analytical')
        plt.legend(loc=0)
        plt.title("Comparison after %s iterations with %s walkers" % (t_max, number_of_walkers))
        
    show_model(all_models["node"])
    show_model(all_models["hyperedges"])

show_models(HG)

100000
100000

HG = generators.uniform_hypergraph(15, 14, 4)
show_models(HG)

100000
100000