import math, random, os, time, sys, io
import json
import numpy as np
import networkx as nx
import scipy
import time
from networkx.drawing.nx_pydot import write_dot

def clean_shuffle_graph(G):
  random_seed_state =  int(random.random()*100000) # quick hack to go back to random afterwards
  random.seed(42)
  node_mapping = dict(zip(sorted(G.nodes()), sorted(G.nodes(), key=lambda _: random.random()))) # maybe sorted not really deterministic
  G = nx.relabel_nodes(G, node_mapping)
  G = nx.convert_node_labels_to_integers(G)
  if not nx.is_connected(G):
    print('Graph is not connected, try a differnt one.')
    assert(nx.is_connected(G))
  random.seed(random_seed_state)
  return G

def gen_sis(G, steps = 1000, inf_rate=1.0, rec_rate=2.0, noise=0.1, statesList = None):
    S = [1., 0.]
    I = [0., 1.]
    states = [random.choice([S, I]) for i in range(G.number_of_nodes())]
    for _ in range(steps):
        rates = np.zeros(G.number_of_nodes())
        for n in range(G.number_of_nodes()):
            rates[n] = noise
            if states[n] == I:
                rates[n] += rec_rate
            if states[n] == S:
                rates[n] += inf_rate * len([n_j for n_j in G.neighbors(n) if states[n_j] == I])
            rates[n] = 1.0/rates[n] # numpy uses mean as rate param
        jump_time = np.random.exponential(rates)
        change_n = np.argmin(jump_time)
        states[change_n] = S if states[change_n] == I else I     
        statesList.append(states.copy())
    return states

statesList = []
steps = 1000 + random.choice(range(1000))
# G_grid10x10 = nx.grid_2d_graph(10,10)
G_grid10x10 = nx.newman_watts_strogatz_graph(120, 4, 0.15, seed=42)
G = clean_shuffle_graph(G_grid10x10)
TS_data = gen_sis(G, steps=steps,statesList=statesList)
# print(G.number_of_nodes())
# f = open("graph.dot", "rw")
# nx.drawing.nx_pydot.write_dot(G,f)
# f.write("\n")
# f.close()
# Writing the json file to stdout
# f = StringIO("")
dotGraph = ""

#Write dot file to string
with io.StringIO() as f:
    write_dot(G, f)
    f.seek(0)
    dotGraph = f.read()

output = {
        "name" : "testGraph",
        "dotGraph" : dotGraph,
        "steps" : steps,
        "states" : statesList,
        }

jsonOutput = json.dumps(output)
print(jsonOutput)