from petnetsim.petnetsim import *
import matplotlib.pyplot as plt
import statistics
import numpy as np

def make_net(running_probability, crash_probability, in_process_probability, timeout_probability, timeout_retry_probability, timeout_crash_probability, job_execution_min_time, job_execution_max_time, microservice_probabilities, requests_time_min, requests_time_max, in_timeout_min_time, in_timeout_max_time):
    global petri_net

    ## MICROSERVICE NET
    # define places for microservice net
    places = [Place('requests', init_tokens=0),
              Place('ready', init_tokens=1),
              Place('crash_decision'),
              Place('timeout_decision'),
              Place('running_counter'),
              Place('retry_decision'),
              Place('timeout_waiting'),
              Place('timeout_retry_counter'),
              Place('in_process_waiting'),
              Place('timeout_counter'),
              Place('crash_waiting'),
              Place('job_done'),
              Place('all_jobs'),
              Place('job_crash')]
    
    # define transitions; stochastic and timed transitions moved to loop
    transitions = [Transition('run'),
                   #for crash decision
                   TransitionStochastic('running',0),
                   TransitionStochastic('crash',0),
                   #for timeout decision
                   TransitionStochastic('in_process',0),
                   TransitionStochastic('timeout',0),
                   #for simulating process timeout
                   TransitionTimed('in_timeout',0,0,uniform_distribution),
                   #for retry decision
                   TransitionStochastic('timeout_retry',0),
                   TransitionStochastic('timeout_crash',0),
                   #for simulating process duration
                   TransitionTimed('job_execution',0,0,uniform_distribution),
                   TransitionTimed('crash_job_execution',0,0,uniform_distribution)]
    # define transition input arcs
    input_arcs = [('requests','run'),
                  ('ready','run'),
                  ('crash_decision','running'),
                  ('crash_decision','crash'),                  
                  ('timeout_decision','in_process'),
                  ('timeout_decision','timeout'),
                  ('timeout_waiting','in_timeout'),
                  ('retry_decision','timeout_retry'),
                  ('retry_decision','timeout_crash'),
                  ('in_process_waiting','job_execution'),
                  ('crash_waiting','crash_job_execution')]

    #define transition output arcs
    output_arcs = [('run','crash_decision'),
                   ('running','timeout_decision'),
                   ('running','running_counter'),
                   ('in_process','in_process_waiting'),
                   ('job_execution','job_done'),
                   ('job_execution','ready'),
                   ('job_execution','all_jobs'),
                   ('timeout','timeout_waiting'),
                   ('in_timeout','retry_decision'),
                   ('timeout','timeout_counter'),
                   ('timeout_retry','timeout_decision'),
                   ('timeout_retry','timeout_retry_counter'),
                   ('timeout_crash','crash_waiting'),
                   ('crash','crash_waiting'),
                   ('crash_job_execution','job_crash'),
                   ('crash_job_execution','all_jobs'),
                   ('crash_job_execution','ready')]
    
    # join arcs
    arcs = input_arcs + output_arcs
    
    # making microservice net
    microservice_net = PetriNet(places,transitions,arcs)
    
    ## SERVER NET
    # some stuff for cloning microservice net; place 'API_gateway' is a statring place for all cloned nets
    API_places = [Place('API_gateway')]
    API_transitions = []
    API_arcs = []
    
    # cloning microservices
    for microservice_i in range(1, microservices + 1):
        # filling probability and time data to transitions
        for tr in microservice_net.transitions:
            if tr.name == 'running':
                tr.probability = running_probability[microservice_i - 1]
            if tr.name == 'crash':
                tr.probability = crash_probability[microservice_i -1]
            if tr.name == 'in_process':
                tr.probability = in_process_probability[microservice_i - 1]
            if tr.name == 'timeout':
                tr.probability = timeout_probability[microservice_i - 1]
            if tr.name == 'timeout_retry':
                tr.probability = timeout_retry_probability[microservice_i - 1]
            if tr.name == 'timeout_crash':
                tr.probability = timeout_crash_probability[microservice_i -1]
            if tr.name == 'job_execution':
                tr.t_min = job_execution_min_time[microservice_i - 1]
                tr.t_max = job_execution_max_time[microservice_i - 1]
            if tr.name == 'crash_job_execution':
                tr.t_min = job_execution_min_time[microservice_i - 1]
                tr.t_max = job_execution_max_time[microservice_i - 1]
            if tr.name == 'in_timeout':
                tr.t_min = in_timeout_min_time[microservice_i - 1]
                tr.t_max = in_timeout_max_time[microservice_i - 1]


        prefix = f'microservice{microservice_i}_'
        microservice_net.clone(prefix, API_places, API_transitions, API_arcs)
        API_transitions.append(TransitionStochastic('to_microservice' + f'{microservice_i}', microservice_probabilities[microservice_i - 1]))
        API_arcs.append(('to_microservice' + f'{microservice_i}', prefix + 'requests')) 
        API_arcs.append(('API_gateway', 'to_microservice' + f'{microservice_i}'))
    
    # making server net
    server_net = PetriNet(API_places, API_transitions, API_arcs)
    

    ## APP NET
    # some stuff for cloning; place load_balancer is a starting place for each server
    LB_places = [Place('requests', requests), Place('turn1',1), Place('load_balancer')]
    LB_transitions = [TransitionTimed('make_request', requests_time_min, requests_time_max, uniform_distribution)]
    LB_arcs = [('requests','make_request'), ('make_request','load_balancer')]
    
    
    if servers > 1: 
        for i in range(2, servers + 1):
            LB_places.append(Place('turn' + f'{i}'))

    
    # cloning servers
    for server_i in range(1, servers + 1):
        prefix = f'server{server_i}_'
        server_net.clone(prefix, LB_places, LB_transitions, LB_arcs)
        LB_transitions.append(Transition('to_server' + f'{server_i}'))

        LB_arcs.append(('turn' + f'{server_i}', 'to_server' + f'{server_i}'))
        
        if server_i == servers:
            LB_arcs.append(('to_server' + f'{server_i}', 'turn1')) 
        else:
            LB_arcs.append(('to_server' + f'{server_i}', 'turn' + f'{server_i + 1}')) 

        LB_arcs.append(('to_server' + f'{server_i}', prefix + 'API_gateway')) 
        LB_arcs.append(('load_balancer', 'to_server' + f'{server_i}'))
    
    ## uncomment for crash simulation
    #crash_simulation(LB_places,LB_transitions,LB_arcs)

    
    #petri net of whole web application
    petri_net = PetriNet(LB_places, LB_transitions, LB_arcs)

# function for simulating server crash, used in make_net function to append some places, transitions and arcs to lists
def crash_simulation(LB_places, LB_transitions, LB_arcs):
    LB_places.append(Place('crash'))
    LB_places.append(Place('crash_request',1))
    LB_transitions.append(TransitionTimed('make_crash',crash_time_max))
    LB_arcs.append(('crash_request','make_crash'))
    LB_arcs.append(('make_crash','crash'))
    LB_arcs.append(Inhibitor('crash', 'to_server2'))
    for i in range(1, microservices + 1):
        LB_arcs.append(Inhibitor('crash', f'server2_microservice{i}_run'))
    LB_arcs.append(('to_server1', 'turn1'))

# just one simulation
def run(max_steps):
    #init variables
    global place_tokens
    global place_time  

    #making empty dictionaries of tokens and time in every places
    place_tokens = dict()
    place_time = dict()
    for p in petri_net.places:
            place_tokens[p.name] = []
            place_time[p.name] = []
    #main loop
    while not petri_net.ended and petri_net.step_num < max_steps:
        petri_net.step()    
        #filling dictionaries
        for p in petri_net.places:
            place_tokens[p.name].append(p.tokens)
            place_time[p.name].append(petri_net.time)

# more simulations (by number in loops) 
def loop_run():
    # defining variables
    global stats
    global probability_of_job_done
    global probability_of_job_crash
    global probability_of_timeout
    global probability_of_running
    global probability_of_microservice_requests
    stats = dict()
    time = [] 
    #
    for loop in range(1,loops + 1):
        petri_net.reset()
        stats[f'loop_{loop}'] = dict()
        while not petri_net.ended and petri_net.step_num < max_steps: 
            petri_net.step()
        
        time.append(petri_net.time)
        
        for p in petri_net.places:
            stats[(f'loop_{loop}')][f'{p.name}'] = []            
            stats[(f'loop_{loop}')][f'{p.name}'].append(p.tokens)
        
    # probability of crash and job done   
    probability_of_job_done = dict()
    probability_of_job_crash = dict()

    for i in range(1, servers + 1):
        probability_of_job_done[f'server_{i}'] = {}
        probability_of_job_crash[f'server_{i}'] = {}
        for k in range(1, microservices + 1):
            probability_of_job_done[f'server_{i}'][f'microservice_{k}'] = []
            probability_of_job_crash[f"server_{i}"][f'microservice_{k}'] = []
            for j in range(1, loops + 1): 
                job_done = stats[(f'loop_{j}')].get(f'server{i}_microservice{k}_job_done')
                all_jobs = stats[(f'loop_{j}')].get(f'server{i}_microservice{k}_all_jobs')
                if all_jobs[0] == 0:
                    probability_done = 0
                else:
                    probability_done = job_done[0] / all_jobs[0]
                probability_crash = 1 - probability_done
                probability_of_job_done[f'server_{i}'][f'microservice_{k}'].append(probability_done)
                probability_of_job_crash[f'server_{i}'][f'microservice_{k}'].append(probability_crash)


    # probability of timeout
    probability_of_timeout = dict()
    probability_of_running = dict()

    for i in range(1, servers + 1):        
        probability_of_timeout[f"server_{i}"] = {}
        probability_of_running[f"server_{i}"] = {}
        for k in range(1, microservices + 1):
            probability_of_timeout[f'server_{i}'][f'microservice_{k}'] = []
            probability_of_running[f"server_{i}"][f'microservice_{k}'] = []
            for j in range(1, loops + 1): 
                timeout = stats[(f'loop_{j}')].get(f'server{i}_microservice{k}_timeout_counter')
                running = stats[(f'loop_{j}')].get(f'server{i}_microservice{k}_running_counter')
                if running[0] == 0:
                    probability_timeout = 0
                else:
                    probability_timeout = timeout[0] / running[0]
                probability_running = 1 - probability_timeout
                probability_of_timeout[f'server_{i}'][f'microservice_{k}'].append(probability_timeout)
                probability_of_running[f'server_{i}'][f'microservice_{k}'].append(probability_running)
    
    # probability, if from timeout retry or crash
    probability_of_microservice_requests = dict()

    for i in range(1, servers + 1):
        probability_of_microservice_requests[f"server_{i}"] = {}
        for k in range(1, microservices + 1):
            probability_of_microservice_requests[f'server_{i}'][f'microservice_{k}'] = []
            for j in range(1, loops + 1): 
                request_no = stats[(f'loop_{j}')].get(f'server{i}_microservice{k}_all_jobs')
                probability_of_microservice_requests[f'server_{i}'][f'microservice_{k}'].append(request_no[0])

# just for model design
def basic_info():
    print('CONFLICT GROUPS')
    print(petri_net.conflict_groups_str)
    print('------------------------------------')
    petri_net.reset()
    print('INITIAL TOKENS:')
    petri_net.print_places()
    print('------------------------------------')
    print('TRANSITIONS STATS:')
    for t in petri_net.transitions:
        print(t.name, t.fired_times, sep=': ')
    print('------------------------------------')
    print('PLACES STATS:')
    petri_net.print_places()
    
# for ploting waiting requests for each microservice
def plot_requests():    
    run(max_steps)
    # plot requests for each microservice
    fig, axs = plt.subplots(1, microservices, facecolor='w', edgecolor='k', sharey=True)
    fig.subplots_adjust(hspace = .5, wspace=.1)
    axs = axs.ravel()
    fig.suptitle('Waiting requests for each microservice', fontsize=16)
    for i in range(1, servers + 1):     
        for j in range(1, microservices + 1):
            axs[j-1].plot(place_time[('server%s_microservice%s_requests' % (i,j))],place_tokens[('server%s_microservice%s_requests' % (i,j))], label=('Server %s' % i))
            axs[j-1].set_title(f'Microservice {j}', fontsize=16)
               
    for ax in axs.flat:
        ax.set_xlabel('time [s]',fontsize=16)
        ax.set_ylabel('requests',fontsize=16)
        #ax.axvline(x = crash_time_max, color='k', linestyle='-.', linewidth=0.5, label = 'Server 2 crash')
        handles, labels = ax.get_legend_handles_labels()

    fig.legend(handles, labels, loc=(0.85,0.75), fontsize=16)

    for ax in fig.get_axes():
        ax.label_outer()

    
# for plotting probability of successful request execution
def plot_running_stats():
    loop_run()
    #ind = np.arange(loops)
    width = 0.1
    # plot requests for each microservice
    fig, axs = plt.subplots(1, 1, facecolor='w', edgecolor='k', sharex=True)
    fig.subplots_adjust(hspace = .5, wspace=.1)
    #axs = axs.ravel()
    fig.suptitle('Probability of successful request execution', fontsize=16)
    probability = {}
    std = {}
    
    for i in range(1, servers + 1):
        probability[i] = []
        std[i] = []
        for j in range(1,microservices + 1):
            probability[i].append(statistics.mean(probability_of_job_done[('server_%s' % i)][('microservice_%s' % j)]))  
            std[i].append(statistics.stdev(probability_of_job_done[('server_%s' % i)][('microservice_%s' % j)]))  
    
    ind = np.arange(microservices)
    for i in range(1, servers +  1):   
        #for j in range(1,microservices + 1):
        axs.bar(ind, np.array(probability[i]) * 100, width, 
        yerr=np.array(std[i]) * 100, label=('Server %s' % i))
        ind = ind + width
        
        axs.set_title(f'Server {i}', fontsize=16)
        axs.legend(fontsize=16)
    
    ind = np.arange(microservices) + ((width / 2) * (servers - 1))
    labels = [('microservice %s' %i) for i in range(1, microservices + 1)]
 
    axs.set_xticks(ind)
    axs.set_xticklabels(labels,fontsize=16)
    axs.set_ylabel('probability',fontsize=16)
    axs.grid(axis='y')

# for plotting number of requests for each microservice
def plot_microservice_stats():
    loop_run()
    #ind = np.arange(loops)
    width = 0.1
    # plot requests for each microservice
    fig, axs = plt.subplots(1, 1, facecolor='w', edgecolor='k', sharex=True)
    fig.subplots_adjust(hspace = .5, wspace=.1)
    #axs = axs.ravel()
    
    fig.suptitle('Number of requests for each microservice', fontsize=16)
    MS_requests = {}
    std = {}
    for i in range(1, servers + 1):
        MS_requests[i] = []
        std[i] = []
        for j in range(1,microservices + 1):
            MS_requests[i].append(statistics.mean(probability_of_microservice_requests[('server_%s' % i)][('microservice_%s' % j)]))  
            std[i].append(statistics.stdev(probability_of_microservice_requests[('server_%s' % i)][('microservice_%s' % j)]))  
    
    ind = np.arange(microservices)
    for i in range(1, servers +  1):   
        #for j in range(1,microservices + 1):
        axs.bar(ind, np.array(MS_requests[i]), width, yerr=np.array(std[i]), label=('Server %s' % i))
        ind = ind + width
        
        axs.set_title(f'Server {i}', fontsize=16)
        axs.legend(fontsize=16)
    
    ind = np.arange(microservices) + ((width / 2) * (servers - 1))
    labels = [('microservice %s' %i) for i in range(1, microservices + 1)]
 
    axs.set_xticks(ind)
    axs.set_xticklabels(labels,fontsize=16)
    axs.set_ylabel('requests',fontsize=16)
    axs.grid(axis='y')

# for plotting load balancer stats (waiting requests)
def plot_loadbalancer_stats():
    width = 0.1
    run(max_steps)
    loop_run()
    fig1 = plt.figure(figsize=(20,8))
    axs1 = fig1.add_subplot(111)
    fig1.suptitle('Waiting requests for microservice', fontsize=16)
    for i in range(1, servers + 1):     
        axs1.plot(place_time[('server%s_microservice%s_requests' % (i,1))],place_tokens[('server%s_microservice%s_requests' % (i,1))], label=('Server %s' % i))
        axs1.set_title('Microservice %s' % 1,fontsize=16)
        axs1.legend(fontsize=16)

    axs1.set_xlabel('time [s]',fontsize=16)
    axs1.set_ylabel('requests',fontsize=16)

# for plotting API gateway stats
def plot_gateway_stats():
    width = 0.1
    loop_run()
    fig2 = plt.figure(figsize=(20,8))
    axs2 = fig2.add_subplot(111)
    fig2.suptitle('Requests for microservices with different probability', fontsize=16)
    data = {}
    for j in range(1,microservices + 1):
        data[j] = []
        for i in range(1, loops + 1):
            a = stats[(f'loop_{i}')].get(f'server{1}_microservice{j}_all_jobs')
            data[j].append(a[0])     
    
    ind = np.arange(loops)
    for i in range(1, microservices + 1):
        axs2.bar(ind, np.array(data[i]), width, label=(f'Microservice {i}'))  
        ind = ind + width

    axs2.set_title('Server 1',fontsize=16)
    axs2.legend(fontsize=16)

    ind = np.arange(loops) + (width / 2)
    labels = [('loop %s' %i) for i in range(1, loops + 1)]

    axs2.set_xticks(ind )
    axs2.set_xticklabels(labels,fontsize=16)
    axs2.set_ylabel('requests',fontsize=16)

# for plotting requests of microservices with differrent probability
def plot_jobdone_stats():
    width = 0.1

    loop_run()
   
    fig2 = plt.figure(figsize=(20,8))
    axs2 = fig2.add_subplot(111)
    fig2.suptitle('Requests for microservices with different probability', fontsize=16)
    data = {}
    for j in range(1,microservices + 1):
        data[j] = []
        for i in range(1, loops + 1):
            a = stats[(f'loop_{i}')].get(f'server{1}_microservice{j}_all_jobs')
            data[j].append(a[0])     
    
    ind = np.arange(loops)
    for i in range(1, microservices + 1):
        axs2.bar(ind, np.array(data[i]), width, label=(f'Microservice {i}'))  
        ind = ind + width

    axs2.set_title('Server 1',fontsize=16)
    axs2.legend(fontsize=16)

    ind = np.arange(loops) + (width / 2)
    labels = [('loop %s' %i) for i in range(1, loops + 1)]

    axs2.set_xticks(ind )
    axs2.set_xticklabels(labels,fontsize=16)
    axs2.set_ylabel('requests',fontsize=16)


def main():
    # init variables
    globals()['requests'] = 500
    globals()['servers'] = 1
    globals()['microservices'] = 5
    globals()['loops'] = 10
    globals()['max_steps'] = 50000
    globals()['crash_time_max'] = 1.5
    global microservice_probabilities
   
    # probability and times of transitions for individual microservices
    running_probability = [0.9, 0.9, 0.9, 0.9, 0.9]
    crash_probability = [0.1, 0.1, 0.1, 0.1, 0.1]
    in_process_probability = [0.9, 0.9, 0.9, 0.9, 0.9]
    timeout_probability = [0.1, 0.1, 0.1, 0.1, 0.1]
    timeout_retry_probability = [0.8, 0.8, 0.8, 0.8, 0.8]
    timeout_crash_probability = [0.2, 0.2, 0.2, 0.2, 0.2]
    job_execution_min_time = [0.005, 0.005, 0.005, 0.005, 0.005]
    job_execution_max_time = [0.05, 0.05, 0.05, 0.05, 0.05]
    in_timeout_min_time = [0.1, 0.1, 0.1, 0.1, 0.1]
    in_timeout_max_time = [0.5, 0.5, 0.5, 0.5, 0.5]

    # probability for API gateway for choosing microservices
    microservice_probabilities = [0.5, 0.4, 0.3, 0.2, 0.1]

    # times for incoming requests
    requests_time_min = 0.001
    requests_time_max = 0.01
    
    # this make the whole net 
    make_net(running_probability, crash_probability, in_process_probability, timeout_probability, timeout_retry_probability, timeout_crash_probability, job_execution_min_time, job_execution_max_time, microservice_probabilities, requests_time_min, requests_time_max, in_timeout_min_time, in_timeout_max_time)
    
    # used for ploting loadbalancer graphs in thesis
    #plot_loadbalancer_stats()

    # used for plotting gateway graphs in thesis
    #plot_gateway_stats()
    
    #loop_run()
    #plot_requests()
    #plot_running_stats()
    plot_microservice_stats()

    plt.show()

if __name__ == '__main__':
    main()


