import numpy as np
from ..algorithms import pair_vector as pv
import json
import matplotlib.pyplot as plt

'''
    For a graph G, this computes the resolution parameter value for which the CM-modularity vector has latitude 
    approximately equal to lat.
'''
def lat2CM_res(G,lat):
    mG = len(G.edges)
    N = len(G) * (len(G) - 1) / 2
    dG = pv.degree_product(G, normalized=True)
    d1 = dG * pv.ones
    dd = dG * dG
    if lat<np.arccos(d1/(N*dd)**0.5):
        return float('inf')
    if lat>np.arccos(-mG/(N*mG)**0.5):
        return -float('inf')
    lat_prev = np.arccos(-mG/(N*mG)**0.5)
    res_prev = 0
    for res_c in (list(np.linspace(0,0.1,11))
                  +list(np.linspace(0.2,1,9))
                  +list(np.linspace(1.5,5,8))
                  +list(np.linspace(6,20,15))
                  +list(np.linspace(30,100,8))
                  +list(np.linspace(200,1000,9))):
        lat_c = pv.query_CM(G,res=res_c).latitude()
        if lat_c<lat:
            return res_c + (res_prev-res_c)*(lat-lat_c)/(lat_prev-lat_c)
        lat_prev = lat_c
        res_prev = res_c
    return res_prev


default_lats = np.linspace(0,np.pi,51)


'''
    Computes all the query vectors, their Louvain projections, and their correlation distance to the ground truth, for
    the network G with ground truth T and corresponding name.
    The parameters compute_CM_mod and compute_meridians specify which query vectors are included in the experiment.
    Because this experiment involves many runs of the Louvain projection, it may take a while to finish. 
    The option save_intermediate saves the json file after each computed candidate, so that nothing is lost when 
    interrupting the experiment.
'''
def perform_experiment(name,G,T,
                       desired_lats=default_lats,
                       load=True,
                       save_intermediate=True,
                       compute_CM_mod=True,
                       compute_meridians=['CM meridian', 'ER modularity', 'Wedges']):
    meridian2map = {
        'CM meridian': pv.query_CM,
        'ER modularity': pv.connectivity,
        'Wedges': pv.wedges,
    }
    scores = {}
    if load:
        with open(name + '_scores.json') as f:
            scores = json.load(f)
    if compute_CM_mod and 'CM modularity' not in scores:
        scores['CM modularity'] = []
    for x_name in compute_meridians:
        if x_name not in scores:
            scores[x_name] = []
    bT = pv.clustering_binary(T)
    for ilat,lat in enumerate(desired_lats):
        for x_name in compute_meridians:
            if len(scores[x_name]) > ilat:
                continue
            x_map = meridian2map[x_name]
            x = x_map(G)
            q = x.latitude_on_meridian(lat)
            bC = pv.clustering_binary(pv.louvain_projection(q))
            scores[x_name].append(bT.meridian_angle(bC))
        if compute_CM_mod and len(scores['CM modularity']) <= ilat:
            res = lat2CM_res(G, lat)
            mod_score = np.pi / 2
            if res not in [float('inf'), -float('inf')]:
                q_mod = pv.query_CM(G, res=res)
                bC_mod = pv.clustering_binary(pv.louvain_projection(q_mod))
                mod_score = bT.meridian_angle(bC_mod)
            scores['CM modularity'].append(mod_score)
        if save_intermediate:
            with open(name + '_scores.json', 'w') as fp:
                json.dump(scores, fp)
    with open(name + '_scores.json', 'w') as fp:
        json.dump(scores, fp)
    return scores


'''
    Plots the experiment for the dataset specified by the name. This requires the file [name]_scores.json (which is
     generated by perform_experiment) to be present in the current directory. lats need to be the same latitude values 
     as were used in the generation of the experiment.
     To prevent certain meridians from being plotted, list them in omit_vectors.
     By specifying gt_lat, the ground-truth latitude will be annotated on the horizontal axis.
'''
def plot_experiment(name,lats=default_lats,omit_vectors=[],gt_lat=None):
    fig, ax = plt.subplots()
    with open(name + '_scores.json') as f:
        scores = json.load(f)
    for v_name, dCCs in scores.items():
        if v_name not in omit_vectors:
            ax.plot(lats, dCCs, label=v_name)
    ax.legend()
    ax.set_xlabel('Query latitude')
    ax.set_ylabel(r'$d_{CC}(b(T),b(C))$')
    xticks = [0, np.pi / 4, np.pi / 2, 3 * np.pi / 4, np.pi]
    xticklabels = ['$0$', r'$\frac{1}{4}\pi$', r'$\frac{1}{2}\pi$', r'$\frac{3}{4}\pi$', r'$\pi$']
    if gt_lat is not None:
        xticklabels = [r'$\ell(b(T))$']+[label for (tick,label) in zip(xticks,xticklabels) if abs(gt_lat-tick)>np.pi/10]
        xticks = [gt_lat]+[tick for tick in xticks if abs(gt_lat-tick)>np.pi/10]
        if len(xticks)!=len(xticklabels):
            print(xticks)
            print(xticklabels)
    plt.setp(ax,
             xticks=xticks,
             xticklabels=xticklabels,
             yticks=[0, np.pi / 6, np.pi / 3, np.pi / 2],
             yticklabels=[r'$0$', r'$\frac{1}{6}\pi$', r'$\frac{1}{3}\pi$', r'$\frac{1}{2}\pi$'])
    ax.get_figure().savefig(name + '_benchmark.svg',bbox_inches='tight')


def generate_figures():
    from .benchmarknetworks import load_dataset
    for name in ['karate', 'dolphins', 'polbooks', 'football', 'eu-core', 'polblogs']:
        G, T = load_dataset(name)
        perform_experiment(name, G, T, load=False)
        plot_experiment(name, gt_lat=pv.clustering_binary(T).latitude())

def generate_figures_for_large_graphs():
    from .benchmarknetworks import load_dataset
    for name in ['cora', 'as']:
        G, T = load_dataset(name)
        perform_experiment(name, G, T, load=False)
        plot_experiment(name, gt_lat=pv.clustering_binary(T).latitude())