import networkx as nx
import matplotlib.pyplot as plt
import random as rnd
import pickle
import numpy as np
import os.path


def coupled_euler(b, I_start, S_start, t_start, t_stop, delta_t=1):
    ts = [t_start]
    Is = [I_start]
    Ss = [S_start]
    while ts[-1] < t_stop:
        t = ts[-1]
        I = Is[-1]
        S = Ss[-1]
        ts.append((t + delta_t))
        Is.append((I + ((1 - (1 - b)**I) * S)))
        Ss.append(S - ((1 - (1 - b)**I) * S))
    return ts, Is, Ss


def prepare_infect_graph(G, perc):
    nodes_infected = rnd.sample(range(len(G)), round(len(G)*perc))
    for i in range(len(G)):
        G.nodes[i]['infected_neighbour'] = False
        if i in nodes_infected:
            G.nodes[i]['infected'] = True
        else:
            G.nodes[i]['infected'] = False
    return len(nodes_infected)


def mean_cols(L):
    mean = []
    for col in range(len(L[0])):
        sum = 0
        for row in range(len(L)):
            sum += L[row][col]
        mean.append(sum / len(L))
    return mean


def error_cols(L):
    error = [[], []]
    for col in range(len(L[0])):
        low_error = min([L[row][col] for row in range(len(L))])
        high_error = max([L[row][col] for row in range(len(L))])
        error[0].append(low_error)
        error[1].append(high_error)
    return error[0], error[1]


def run_simulation_ep1(G, i, t_max, n_inf):
    G.graph['frac_infected_t'] = [n_inf / len(G)]
    G.graph['avg_deg_inf'] = []
    skip = False

    for t in range(t_max):
        #TODO
        print(f"N: {len(G)}, t: {t}, inf: ", end="")
        inf_count = 0
        sum_deg = 0

        if skip:
            G.graph['frac_infected_t'].append(1)
        for node in G.nodes():
            if G.nodes[node]['infected']:
                for neighbour in G[node]:
                    if G.nodes[neighbour]['infected']:
                        continue
                    if rnd.random() < i:
                        inf_count += 1
                        G.nodes[node]['infected_neighbour'] = True
                        G.nodes[neighbour]['infected'] = True
                        sum_deg += len(G[neighbour])
        if inf_count:
            G.graph['avg_deg_inf'].append(sum_deg/inf_count)
        else:
            G.graph['avg_deg_inf'].append(0)
        cur_inf = round(G.graph['frac_infected_t'][-1] * len(G)) + inf_count
        print(cur_inf)
        if cur_inf >= len(G):
            skip = True

        G.graph['frac_infected_t'].append(G.graph['frac_infected_t'][-1] + (inf_count / len(G)))


def main():
    # Ep1.b
    n = 5
    max_t = 500

    if not os.path.isfile("ep1.b.i") or not os.path.isfile("ep1.b.ii"):
        Iti = []
        Itii = []

        for _ in range(n):
            G1bi = nx.fast_gnp_random_graph(10**5, 5.0 / 10**5)
            G1bii = nx.fast_gnp_random_graph(10**5, 0.8 / 10**5)
            infi = prepare_infect_graph(G1bi, 0.001)
            infii = prepare_infect_graph(G1bii, 0.001)
            run_simulation_ep1(G1bi, 0.01, max_t, infi)
            run_simulation_ep1(G1bii, 0.1, max_t, infii)
            Iti.append(G1bi.graph['frac_infected_t'])
            Itii.append(G1bii.graph['frac_infected_t'])

        try:
            f = open("ep1.b.i", "wb")
            pickle.dump(Iti, f)
            f = open("ep1.b.ii", "wb")
            pickle.dump(Itii, f)
        except:
            print("Failed to save")
    else:
        try:
            f = open("ep1.b.i", "rb")
            Iti = pickle.load(f)
            f = open("ep1.b.ii", "rb")
            Itii = pickle.load(f)
        except:
            print("Failed to load")
            exit()

    mean_Iti = mean_cols(Iti)
    mean_Itii = mean_cols(Itii)


    plt.plot(range(max_t + 1), mean_Iti, color="g")
    plt.plot(range(max_t + 1), mean_Itii, color="b")
    plt.fill_between(range(max_t + 1), *error_cols(Iti), alpha=0.3, color="g")
    plt.fill_between(range(max_t + 1), *error_cols(Itii), alpha=0.3, color="b")
    plt.xlabel("t")
    plt.ylabel("Fraction infected")
    plt.legend(["(i) i = 0.01, <k> = 5.0", "(ii) i = 0.1, <k> = 0.8"])
    plt.savefig("ep1_b.png")
    plt.clf()


    # Ep1.h
    R0i = []
    R0ii = []
    for i in range(max_t):
        R0i.append((mean_Iti[i + 1] - mean_Iti[i]) / mean_Iti[i])
        R0ii.append((mean_Itii[i + 1] - mean_Itii[i]) / mean_Itii[i])

    plt.plot(range(max_t), R0i, color="g")
    plt.plot(range(max_t), R0ii, color="b")
    plt.xlabel("t")
    plt.ylabel("I'(t)/I(y)")
    plt.legend(["case (i)", "case (ii)"])
    plt.savefig("ep1_h.png")
    plt.clf()


    # Ep3.a
    t, Ii, Si = coupled_euler(5.03 * 10**-7, 100, 100000 - 100, 0, 500)
    _, Iii, Sii = coupled_euler(8.43 * 10**-7, 100, 100000 - 100, 0, 500)
    for i in range(len(Ii)):
        Ii[i] = Ii[i]/100000
        Iii[i] = Iii[i]/100000
    plt.plot(range(max_t + 1), mean_Iti, color="g")
    plt.plot(range(max_t + 1), mean_Itii, color="b")
    plt.plot(range(max_t + 1), Ii, color="r")
    plt.plot(range(max_t + 1), Iii, color="y")
    plt.legend(["case (i) num", "case (ii) num", "case (i) ODE", "case (ii) ODE"])
    plt.savefig("ep3_a.png")
    plt.clf()

    # Ep4.a
    N = 10**5
    γ = 2.5
    k = 5
    s = []

    if not os.path.isfile("ep4.a") or not os.path.isfile("ep4.a.deg"):
        It4 = []

        for _ in range(2):
            while len(s) < N:
                n_val = int(nx.utils.powerlaw_sequence(k, γ)[0])
                if n_val:
                    s.append(n_val)
            if sum(s) % 2:
                s.append(1)
            G4 = nx.configuration_model(s)
            G4 = nx.Graph(G4)
            G4.remove_edges_from(G4.selfloop_edges())

            inf4 = prepare_infect_graph(G4, 0.001)
            run_simulation_ep1(G4, 0.01, 500, inf4)
            It4.append(G4.graph['frac_infected_t'])

        avg_deg = G4.graph['avg_deg_inf']

        try:
            f = open("ep4.a", "wb")
            pickle.dump(It4, f)
            f = open("ep4.a.deg", "wb")
            pickle.dump(avg_deg, f)
        except:
            print("Failed to save")
    else:
        try:
            f = open("ep4.a", "rb")
            It4 = pickle.load(f)
            f = open("ep4.a.deg", "rb")
            avg_deg = pickle.load(f)
        except:
            print("Failed to load")
            exit()

    mean_It4 = mean_cols(It4)
    plt.plot(range(max_t + 1), mean_Iti, color="g")
    plt.plot(range(max_t + 1), mean_It4, color="m")
    plt.xlabel("t")
    plt.ylabel("Fraction infected")
    plt.legend(["Case (i)", "Power law network"])
    plt.savefig("ep4_a.png")
    plt.clf()

    plt.plot(range(max_t), avg_deg, color="g")
    plt.xlabel("t")
    plt.ylabel("Degree")
    plt.savefig("ep4_c.png")
    plt.clf()

if __name__ == "__main__":
    main()
