<reponame>corollary-health/episimmer
import Simulate
import ReadFile
import Time
import random
from vulnerability_detection.Base_Class import Event_VD

events_to_restrict = None

def event_restriction_fn(agent,event_info,current_time_step):
    if event_info["Event Index"] in events_to_restrict:
        return True

    return False


class Event_Vulnerabilities_MC(Event_VD):
    """
    Parameters to pass : states, num_runs, num_events_to_restrict
    """
    def __init__(self, world_obj, parameter_dict):
        super().__init__()
        self.world_obj = world_obj
        self.states = parameter_dict["states"]
        self.num_runs = parameter_dict["num_runs"]
        self.num_events_to_restrict = 1
        self.locations_obj = ReadFile.ReadLocations(self.world_obj.locations_filename, self.world_obj.config_obj)
        self.agents_obj = ReadFile.ReadAgents(self.world_obj.agents_filename, self.world_obj.config_obj)
        self.event_counts = {}
        self.all_events = []
        self.current_run=0
        self.initialise(self.world_obj.eventFiles_list)

    def reset_events_to_restrict(self):
        global events_to_restrict
        events_to_restrict = []
        self.locations_obj = ReadFile.ReadLocations(self.world_obj.locations_filename, self.world_obj.config_obj)

    def initialise(self, eventFiles_listOfList):
        for eventFiles_list in eventFiles_listOfList:
            if eventFiles_list != []:
                for events_filename in eventFiles_list:
                    ReadFile.ReadEvents(events_filename, self.world_obj.config_obj, self.locations_obj, self.agents_obj)
                    for location in self.locations_obj.locations.values():
                        self.all_events += location.events
                        for event_info in location.events:
                            self.event_scores[event_info["Event Index"]] = 0.0
                            self.event_counts[event_info["Event Index"]] = 0
                        location.new_time_step()

    def restrict_events(self, eventFiles_listOfList, num_events_to_restrict):
        global events_to_restrict
        event_infos_to_restrict = [self.all_events[self.current_run%len(self.all_events)]]
        events_to_restrict = [event_info["Event Index"] for event_info in event_infos_to_restrict]
        for event_index in events_to_restrict:
            self.event_counts[event_index]+=1

    def update_event_scores(self, end_state):
        global events_to_restrict
        epidemic_score = 0
        for state in self.states:
            epidemic_score += end_state[state][-1]
        total_agents = 0
        for key in end_state:
            total_agents += end_state[key][-1]
        score = 1.0 - epidemic_score/total_agents
        for event_info in self.all_events:
            if event_info["Event Index"] in events_to_restrict:
                self.event_scores[event_info["Event Index"]] += 1.0 / (self.event_counts[event_info["Event Index"]]) * (score - self.event_scores[event_info["Event Index"]])


    def one_run(self):
        self.reset_events_to_restrict()
        Time.Time.new_world()
        time_steps = self.world_obj.config_obj.time_steps
        if(self.num_events_to_restrict>0):
            self.restrict_events(self.world_obj.eventFiles_list,self.num_events_to_restrict)

        oneTimeEvent_obj = ReadFile.ReadOneTimeEvents(self.world_obj.one_time_event_file)

        sim_obj = Simulate.Simulate(self.world_obj.config_obj, self.world_obj.model, self.world_obj.policy_list, self.world_obj.event_restriction_fn, self.agents_obj, self.locations_obj)
        sim_obj.event_restriction_fn = event_restriction_fn
        sim_obj.onStartSimulation()

        for current_time_step in range(time_steps):
            sim_obj.onStartTimeStep(self.world_obj.interactionFiles_list, self.world_obj.eventFiles_list, self.world_obj.probabilistic_interactionsFiles_list, oneTimeEvent_obj)
            sim_obj.handleTimeStepForAllAgents()
            sim_obj.endTimeStep()
            Time.Time.increment_current_time_step()

        end_state = sim_obj.endSimulation()
        self.update_event_scores(end_state)

    def run_detection(self):
        for i in range(self.num_runs):
            self.current_run=i
            if (i+1)%(self.num_runs/10)==0:
                print("Iteration running : ",i+1)
            self.one_run()

    def get_maximum_event_vulnerability(self,n):
        res_max = dict(sorted(self.event_scores.items(), key = lambda x: (x[1], x[0]))[-n:])
        return res_max

    def get_minimum_event_vulnerability(self,n):
        res_min = dict(sorted(self.event_scores.items(), key = lambda x: (x[1], x[0]))[:n])
        return res_min

    def print_default_output(self,n):
        print(self.get_maximum_event_vulnerability(n))
        print(self.get_minimum_event_vulnerability(n))

import Simulate
import ReadFile
import Time
import random
from vulnerability_detection.Base_Class import Event_VD

events_to_restrict = None

def event_restriction_fn(agent,event_info,current_time_step):
    if event_info["Event Index"] in events_to_restrict:
        return True

    return False


class Event_Vulnerabilities_MC(Event_VD):
    """
    Parameters to pass : states, num_runs, num_events_to_restrict
    """
    def __init__(self, world_obj, parameter_dict):
        super().__init__()
        self.world_obj = world_obj
        self.states = parameter_dict["states"]
        self.num_runs = parameter_dict["num_runs"]
        self.num_events_to_restrict = 1
        self.locations_obj = ReadFile.ReadLocations(self.world_obj.locations_filename, self.world_obj.config_obj)
        self.agents_obj = ReadFile.ReadAgents(self.world_obj.agents_filename, self.world_obj.config_obj)
        self.event_counts = {}
        self.all_events = []
        self.current_run=0
        self.initialise(self.world_obj.eventFiles_list)

    def reset_events_to_restrict(self):
        global events_to_restrict
        events_to_restrict = []
        self.locations_obj = ReadFile.ReadLocations(self.world_obj.locations_filename, self.world_obj.config_obj)

    def initialise(self, eventFiles_listOfList):
        for eventFiles_list in eventFiles_listOfList:
            if eventFiles_list != []:
                for events_filename in eventFiles_list:
                    ReadFile.ReadEvents(events_filename, self.world_obj.config_obj, self.locations_obj, self.agents_obj)
                    for location in self.locations_obj.locations.values():
                        self.all_events += location.events
                        for event_info in location.events:
                            self.event_scores[event_info["Event Index"]] = 0.0
                            self.event_counts[event_info["Event Index"]] = 0
                        location.new_time_step()

    def restrict_events(self, eventFiles_listOfList, num_events_to_restrict):
        global events_to_restrict
        event_infos_to_restrict = [self.all_events[self.current_run%len(self.all_events)]]
        events_to_restrict = [event_info["Event Index"] for event_info in event_infos_to_restrict]
        for event_index in events_to_restrict:
            self.event_counts[event_index]+=1

    def update_event_scores(self, end_state):
        global events_to_restrict
        epidemic_score = 0
        for state in self.states:
            epidemic_score += end_state[state][-1]
        total_agents = 0
        for key in end_state:
            total_agents += end_state[key][-1]
        score = 1.0 - epidemic_score/total_agents
        for event_info in self.all_events:
            if event_info["Event Index"] in events_to_restrict:
                self.event_scores[event_info["Event Index"]] += 1.0 / (self.event_counts[event_info["Event Index"]]) * (score - self.event_scores[event_info["Event Index"]])


    def one_run(self):
        self.reset_events_to_restrict()
        Time.Time.new_world()
        time_steps = self.world_obj.config_obj.time_steps
        if(self.num_events_to_restrict>0):
            self.restrict_events(self.world_obj.eventFiles_list,self.num_events_to_restrict)

        oneTimeEvent_obj = ReadFile.ReadOneTimeEvents(self.world_obj.one_time_event_file)

        sim_obj = Simulate.Simulate(self.world_obj.config_obj, self.world_obj.model, self.world_obj.policy_list, self.world_obj.event_restriction_fn, self.agents_obj, self.locations_obj)
        sim_obj.event_restriction_fn = event_restriction_fn
        sim_obj.onStartSimulation()

        for current_time_step in range(time_steps):
            sim_obj.onStartTimeStep(self.world_obj.interactionFiles_list, self.world_obj.eventFiles_list, self.world_obj.probabilistic_interactionsFiles_list, oneTimeEvent_obj)
            sim_obj.handleTimeStepForAllAgents()
            sim_obj.endTimeStep()
            Time.Time.increment_current_time_step()

        end_state = sim_obj.endSimulation()
        self.update_event_scores(end_state)

    def run_detection(self):
        for i in range(self.num_runs):
            self.current_run=i
            if (i+1)%(self.num_runs/10)==0:
                print("Iteration running : ",i+1)
            self.one_run()

    def get_maximum_event_vulnerability(self,n):
        res_max = dict(sorted(self.event_scores.items(), key = lambda x: (x[1], x[0]))[-n:])
        return res_max

    def get_minimum_event_vulnerability(self,n):
        res_min = dict(sorted(self.event_scores.items(), key = lambda x: (x[1], x[0]))[:n])
        return res_min

    def print_default_output(self,n):
        print(self.get_maximum_event_vulnerability(n))
        print(self.get_minimum_event_vulnerability(n))
Parameters to pass : states, num_runs, num_events_to_restrict