import random, os, sys, getopt, subprocess, math
import networkx as nx
import numpy as np
import scipy.spatial
import settings

def saveSnapshot(tempG, threshold, outputFile):
        if threshold>1:
            tempNodeLabels = tempG.nodes()
            adjMatr = nx.to_numpy_matrix(tempG)
            newAdjMatr = np.maximum(adjMatr  - (threshold - 1) * np.ones(adjMatr.shape), np.zeros(adjMatr.shape))
            tempG = nx.relabel_nodes(nx.from_numpy_matrix(newAdjMatr), dict(zip(range(0,tempG.order()), tempNodeLabels)))

        nx.write_edgelist(nx.Graph(tempG), outputFile, data=False)

def roulette(options, weights):
    weightsum = float(sum(weights))
    if weightsum == 0.0:
        weightsum = 1.0
    probabilities = [float(x)/weightsum for x in weights]
    r = random.random()
    s = 0.0
    for k in xrange(len(options)):
        s += probabilities[k]
        if r <= s:
            break
    return options[k]

def _random_subset(seq,m):
    """ Return m unique elements from seq.

    This differs from random.sample which can return repeated
    elements if seq holds repeated elements.
    """
    targets=set()
    while len(targets)<m:
        x=random.choice(seq)
        targets.add(x)
    return targets

def rejection_sampling_1D(f_x, g_var, g_mean, M, num_samples ):
    ''' does rejection sampling using a Gaussian distribution as  the
    underlying density, given a target density function f_x and a number of
    samples we want generated.'''
    g_x = lambda x: 1/np.sqrt(2*np.pi*g_var)*np.exp(-(x-g_mean)**2/(2*g_var))
    good_samples = 0;
    keepers = np.zeros(num_samples)
    while good_samples < num_samples:
        # generate a new sample using our underlying distribution
        new_samp = np.random.rand(1)*np.sqrt(g_var)+g_mean
        # generate a new uniform
        new_check = np.random.rand(1)
        # first, let's make sure M is big enough
        while f_x(new_samp)/g_x(new_samp)/M > 1:
            M = 2*M
            return rejection_sampling_1D( f_x, g_var, g_mean, M, num_samples )
        # now let's check if our stuff actually works.
        if new_check >= f_x(new_samp)/g_x(new_samp)/M:
            keepers[good_samples] = new_samp
            good_samples = good_samples + 1
        else:
            continue
    return keepers


def generate_pa(model_name, network_name_suffix, delimiter = ' '):
    #print 'Generating random network using Preferential Attachment'

    model_values = model_name.split('_')
    n = int(model_values[1])
    m = int(model_values[2])


    # Add m initial nodes (m0 in barabasi-speak)
    #G=nx.empty_graph(m)
    # Target nodes for new edges
    targets=list(range(m))
    # List of existing nodes, with nodes repeated once for each adjacent edge
    repeated_nodes=[]
    # Start adding the other n-m nodes. The first node is m.
    source=m

    output_dir = settings.BASE_INPUT_DIR + 'model/pa/%s/%s_%s/' % (model_name,model_name,network_name_suffix)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_file = open(output_dir + '/%s_%s.txt' % (n,m,network_name_suffix),'wb')

    i_time = 1
    while source<n:
        # Add edges to m nodes from the source.
        new_edges = zip([source]*m,targets)
        #G.add_edges_from(new_edges)

        for edge in new_edges:
            output_file.write('%s%s%s%s%s\n' % (i_time,delimiter,edge[0],delimiter,edge[1]))

        i_time += 1

        # Add one node to the list for each new edge just created.
        repeated_nodes.extend(targets)
        # And the new node "source" has m edges to add to the list.
        repeated_nodes.extend([source]*m)
        # Now choose m unique nodes from the existing nodes
        # Pick uniformly from repeated_nodes (preferential attachment)
        targets = _random_subset(repeated_nodes,m)
        source += 1
    output_file.close()

    return i_time

def generate_network_evolution(model_name, network_name_suffix, delimiter = ' '):

    model_values = model_name.split('_')
    node_arrival_type = model_values[1]
    sleep_type = model_values[2]

    #parandomrandom edge addition type - first edge PA, then randomly choose neighbor and randomly choose neighbor's neighbor (default)
    #random edge addition type - completely random (random)
    #simple preferential attachment - pa20
    edge_addition_type = model_values[3]

    if not edge_addition_type in ['default','random','pa20']:
        raise Exception('Unknown edge addition type \'%s\'' % edge_addition_type)

    N_max = int(model_values[4])

    if node_arrival_type == 'exp':
        node_arrival_function = lambda x: int(math.exp(0.2*x))
    elif node_arrival_type == 'quad':
        node_arrival_function = lambda x: int(0.2 * x * x + 4 *  x + 1)
    elif node_arrival_type == 'lin':
        node_arrival_function = lambda x: 20 * x
    else:
        raise Exception('Unknown node arrival type \'%s\'' % node_arrival_type)
    #larger l -> larger lifetime
    if edge_addition_type == 'pa20':
        sleep_type = 'type0' # there is no wakeup for preferential attachment

    if sleep_type == 'type1':
        l = 0.008 * 4
        alpha = 0.8
        beta = 0.002
    elif sleep_type == 'type2':
        l = 0.005 * 4
        alpha = 0.9
        beta = 0.004


    output_dir = settings.BASE_INPUT_DIR + 'model/ne/%s/%s_%s/' % (model_name,model_name,network_name_suffix)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_file = open(output_dir + '/%s_%s.txt' % (model_name,network_name_suffix),'wb')

    G = nx.Graph()
    nodes_lifetimes = np.zeros(N_max)
    nodes_wakeup_times =  np.zeros(N_max)
    N_previous = 1
    t = 1

    # List of existing nodes, with nodes repeated once for each adjacent edge
    repeated_nodes=[0]

    while True:
        N = node_arrival_function(t)
        i_iterations = 0
        # while N < N_previous:
        #     t+=1
        #     N = node_arrival_function(t)
        #     if i_iterations>100:
        #         raise Exception('Network does not grow for more than 100 steps')
        #     i_iterations+=1
        N_added = N - N_previous


        for i in range(0,N_added):
            #print '%s: %s < %s' % (t, N_previous + i, N_max)
            if N_previous + i >= N_max:
                break;

            u =  N_previous+i
            #add node u
            G.add_node(u)

            #sample lifetime
            if sleep_type != 'type0':
                nodes_lifetimes[u] = np.random.exponential(1.0/l)
            elif edge_addition_type == 'pa20':
                if G.order()<=20:
                    continue

            #add first edge using preferential attachment
            if edge_addition_type == 'pa20':
                for i in range(0,20):
                    v = list(_random_subset(list(set(repeated_nodes)),1))[0]
                    repeated_nodes.extend([u,v])
                    G.add_edge(u,v)
                    output_file.write('%s%s%s%s%s\n' % (t,delimiter,u,delimiter,v))
            else:
                if edge_addition_type == 'default':
                    v = list(_random_subset(repeated_nodes,1))[0]
                elif edge_addition_type == 'random':
                    v = list(_random_subset(list(set(repeated_nodes)),1))[0]
                repeated_nodes.extend([u,v])
                G.add_edge(u,v)
                output_file.write('%s%s%s%s%s\n' % (t,delimiter,u,delimiter,v))


            if sleep_type != 'type0':
                d = G.degree(u)
                Z = math.gamma(1 - alpha) / math.pow(beta * d,1-alpha)
                p_sleep = lambda x: (1.0/Z) * math.pow(x,-1.0 * alpha) * math.exp(-1.0 * beta * d *x)
                sigma = rejection_sampling_1D(p_sleep,5,2,0.3,1)[0]

                nodes_wakeup_times[u] = t + sigma

        N_previous = G.order()
        if N_previous >= N_max:
            break


        if sleep_type != 'type0':
            wakup_nodes_sorted = np.argsort(nodes_wakeup_times)
            for u in wakup_nodes_sorted:
                if t<nodes_wakeup_times[u]<=t+1:
                    if nodes_lifetimes[u]<nodes_wakeup_times[u]:
                        nodes_wakeup_times[u] = 0
                        continue

                    if edge_addition_type == 'default':
                        v = list(_random_subset(G.neighbors(u),1))[0]
                        v_neighbors = G.neighbors(v)
                        if len(v_neighbors)==1:
                            w=v
                        else:
                            v_neighbors.remove(u)
                            w = list(_random_subset(v_neighbors,1))[0]
                    elif edge_addition_type == 'random':
                        w = list(_random_subset(list(set(repeated_nodes)-{u}),1))[0]


                    if not G.has_edge(u,w):
                        repeated_nodes.extend([u,w])
                    G.add_edge(u,w)
                    output_file.write('%s%s%s%s%s\n' % (t,delimiter,u,delimiter,v))

                    d = G.degree(u)

                    Z = math.gamma(1 - alpha) / math.pow(beta * d,1-alpha)
                    p_sleep = lambda x: (1.0/Z) * math.pow(x,-1.0 * alpha) * math.exp(-1.0 * beta * d *x)
                    sigma = rejection_sampling_1D(p_sleep,5,2,0.3,1)[0]

                    nodes_wakeup_times[u] = t + sigma

                elif nodes_wakeup_times[u]>t+1:
                    break

        t+=1

    output_file.close()
    return t



def generate_geogd(model_name, network_name_suffix, delimiter = ' '):
    '''
    Geometric model with Gene Duplication (based on paper "GEOMETRIC EVOLUTIONARY DYNAMICS OF PROTEIN INTERACTION NETWORKS")
    :param model_name:
    :param network_name_suffix:
    :param delimiter:
    :return:
    '''

    model_values = model_name.split('_')
    node_arrival_type = model_values[1]
    cutoff = float('0.' + model_values[2][1:]) #'05' -> 0.5
    N_max = int(model_values[3])
    dim = 3
    eps = 1

    if node_arrival_type == 'exp':
        node_arrival_function = lambda x: int(math.exp(0.2*x))
    elif node_arrival_type == 'quad':
        node_arrival_function = lambda x: int(0.2 * x * x + 4 *  x + 1)
    elif node_arrival_type == 'lin':
        node_arrival_function = lambda x: 20 * x
    else:
        raise Exception('Unknown node arrival type \'%s\'' % node_arrival_type)

    output_dir = settings.BASE_INPUT_DIR + 'model/geogd/%s/%s_%s/' % (model_name,model_name,network_name_suffix)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_file = open(output_dir + '/%s_%s.txt' % (model_name,network_name_suffix),'wb')

    G = nx.Graph()
    N_previous = 5
    t = 1

    points = []

    #seed network - K5
    for i in xrange(5):
        points.append(generate_random_point([0]*dim,eps/2.0,dim))
    for i1 in xrange(len(points)):
        for i2 in xrange(i1+1,len(points)):
            if scipy.spatial.distance.euclidean(points[i1],points[i2])<eps:
                G.add_edge(i1,i2)
                output_file.write('%s%s%s%s%s\n' % (t,delimiter,i1,delimiter,i2))

    while True:
        N = node_arrival_function(t)
        N_added = N - N_previous
        print '%s = %s - %s' % (N_added, N, N_previous)
        for i in range(0,N_added):
            if N_previous + i >= N_max:
                break;

            i_p = random.choice(range(len(points)))

            p_child = generate_random_point(points[i_p],eps * (1 if random.uniform(0,1)<cutoff else 10),dim)
            i_child = len(points)
            points.append(p_child)
            G.add_node(i_child)

            for i1 in xrange(len(points)-1):
                if scipy.spatial.distance.euclidean(points[i1],points[i_child])<eps:
                    G.add_edge(i1,i_child)
                    output_file.write('%s%s%s%s%s\n' % (t,delimiter,i1,delimiter,i_child))

        N_previous = G.order()
        if N_previous >= N_max:
            break
        t+=1


    output_file.close()
    return t

def generate_random_point(center, radius, dim=3):
    p = []
    for d in range(0,dim):
        p.append(center[d] + random.uniform(-radius,+radius))
    return p


def generate_sfgd(model_name, network_name_suffix, delimiter = ' '):
    '''
    Scale-Free model with Gene Duplication (based on paper "Modeling of Protein Interaction Networks")
    :param model_name:
    :param network_name_suffix:
    :param delimiter:
    :return:
    '''

    model_values = model_name.split('_')
    node_arrival_type = model_values[1]
    p = float('0.' + model_values[2][1:]) #'05' -> 0.5
    q = float('0.' + model_values[3][1:]) #'05' -> 0.5
    N_max = int(model_values[4])

    if node_arrival_type == 'exp':
        node_arrival_function = lambda x: int(math.exp(0.2*x))
    elif node_arrival_type == 'quad':
        node_arrival_function = lambda x: int(0.2 * x * x + 4 *  x + 1)
    elif node_arrival_type == 'lin':
        node_arrival_function = lambda x: 20 * x
    else:
        raise Exception('Unknown node arrival type \'%s\'' % node_arrival_type)

    output_dir = settings.BASE_INPUT_DIR + 'model/sfgd/%s/%s_%s/' % (model_name,model_name,network_name_suffix)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_file = open(output_dir + '/%s_%s.txt' % (model_name,network_name_suffix),'wb')

    G = nx.Graph()
    N_previous = 2
    t = 1



    #seed network - K2
    G.add_edge(0,1)
    output_file.write('%s%s%s%s%s\n' % (t,delimiter,0,delimiter,1))


    while True:
        N = node_arrival_function(t)
        N_added = N - N_previous

        for i in range(0,N_added):
            if N_previous + i >= N_max:
                break;

            u =  N_previous+i
            #add node u
            G.add_node(u)

            node_parent = random.choice(range(G.order()))

            if random.uniform(0,1)<p:
                    G.add_edge(u,node_parent)
                    output_file.write('%s%s%s%s%s\n' % (t,delimiter,node_parent,delimiter,u))

            for v in G.neighbors(node_parent):
                is_added = True
                if random.uniform(0,1)<q:
                    if random.uniform(0,1)<0.5:
                        G.remove_edge(v,node_parent)#note: this is not reflected in the output
                    else:
                        is_added = False
                if is_added:
                    G.add_edge(u,v)
                    output_file.write('%s%s%s%s%s\n' % (t,delimiter,v,delimiter,u))

        N_previous = G.order()
        if N_previous >= N_max:
            break
        t+=1

    output_file.close()
    return t



snapshots_delta_t_list = ['10']
number_of_samples = 1
delimiter = ' '
use_existing_networks = False
accumulating_snapshots = False



opts, args = getopt.getopt(sys.argv[1:], "c:s:ah", ["help","use-existing-networks"])
for o, a in opts:
    if o == '-s':
        snapshots_delta_t_list = a.split(',')
    elif o == '-c':
        number_of_samples = int(a)
    elif o == '--use-existing-networks':
        use_existing_networks = True
    elif o in ('-h','--help'):
        print """Usage:
%s model_name
-s                      Snapshot delta t
-c                      Number of different samples to create
-a                      Accumulating snapshots (i.e., edges from earlier snapshots appear in later snapshots)
--use_existing_networks If samples were already generated, do not re-generate them
-h,--help               Print this""" % sys.argv[0]
        sys.exit(0)

model_name_list = args[0].split(',')

i_model = 1
for model_name in model_name_list:
    print '%s/%s. Generating model %s' % (i_model, len(model_name_list),model_name)
    model_params = model_name.split('_')
    for i in range(1,number_of_samples+1):


        network_name = model_name +'_' + str(i)
        if use_existing_networks and os.path.exists(settings.BASE_INPUT_DIR + 'model/%s/%s/%s/%s.txt' % (model_params[0], model_name, network_name,network_name)):
            print 'Using existing network #%s' % i
        else:
            print 'Generating network #%s' % i
            if model_params[0] == 'pa':
                generate_pa(model_name, str(i),delimiter)
            elif model_params[0] == 'ne':
                generate_network_evolution(model_name, str(i),delimiter)
            elif model_params[0] == 'geogd':
                generate_geogd(model_name,str(i),delimiter)
            elif model_params[0] == 'sfgd':
                generate_sfgd(model_name,str(i),delimiter)


        print 'Splitting into snapshots...'
        for snapshots_delta_t in snapshots_delta_t_list:
            accumulating_snapshots = 'a' in snapshots_delta_t
            print 'delta t = %s' % snapshots_delta_t.replace('a','') + (' (accumulating)' if accumulating_snapshots else '')
            
            run_parameters =  ['python', 'split_into_snapshots.py', '-s','model/%s/%s' % (model_params[0], model_name), network_name, str(snapshots_delta_t.replace('a','')), '1']
            if accumulating_snapshots:
                run_parameters.insert(2,'-a')
            p = subprocess.Popen(run_parameters, stdout=subprocess.PIPE)
            p_output = p.stdout.read()
            if p_output != '':
                print(p_output)

        print 'Aggregating...'
        p = subprocess.Popen(['python', 'aggregate_network.py',  '-s','model/%s/%s' % (model_params[0], model_name), network_name, '1'], stdout=subprocess.PIPE)
        p_output = p.stdout.read()
        if p_output != '':
            print(p_output)

        print 'Converting to leda format...'
        p = subprocess.Popen(['python', 'convert_network_format.py',  '-s','model/%s/%s' % (model_params[0], model_name), network_name, 'list','leda'], stdout=subprocess.PIPE)
        p_output = p.stdout.read()
        if p_output != '':
            print(p_output)