import random

import numpy as np

from episimmer.read_file import ReadAgents, ReadLocations, ReadOneTimeEvents
from episimmer.simulate import Simulate
from episimmer.utils.time import Time

from .base import AgentVulnerabilities, VulnerableAgents


class VulnerableAgentMC1(VulnerableAgents):
    """
    Parameters to pass : states, num_runs
    """
    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.init_scores()

    def init_scores(self):
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        for agent_index in agents_obj.agents.keys():
            self.agent_scores[agent_index] = 0

    def update_agent_scores(self, agents_obj):
        for agent in agents_obj.agents.values():
            if agent.state in self.states:
                self.agent_scores[agent.index] += 1

    def one_run(self):
        _, agents_obj, _ = self.world_obj.one_world()
        self.update_agent_scores(agents_obj)

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

    def get_maximum_vulnerable_agents(self, n):
        res_max = dict(
            sorted(self.agent_scores.items(),
                   key=lambda x: (x[1], x[0]),
                   reverse=True)[:n])
        return res_max

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

    def print_default_output(self, n):
        print(self.get_maximum_vulnerable_agents(n))
        print(self.get_minimum_vulnerable_agents(n))


class VulnerableAgentMC2(VulnerableAgents):
    """
    Parameters to pass : states, num_runs
    """
    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.agent_start_infection_step = {}
        self.init_scores()

    def init_scores(self):
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        for agent_index in agents_obj.agents.keys():
            self.agent_scores[agent_index] = 0

    def reset_start_infection(self, agents_obj):
        for agent_index in agents_obj.agents.keys():
            self.agent_start_infection_step[agent_index] = None

    def update_agent_scores(self, agents_obj):
        total_time_steps = self.world_obj.config_obj.time_steps
        for agent in agents_obj.agents.values():
            if (self.agent_start_infection_step[agent.index] is not None):
                self.agent_scores[
                    agent.
                    index] += total_time_steps - self.agent_start_infection_step[
                        agent.index]

    def update_agent_start_infection(self, agents_obj, time_step):
        for agent in agents_obj.agents.values():
            if agent.state in self.states and self.agent_start_infection_step[
                    agent.index] is None:
                self.agent_start_infection_step[agent.index] = time_step

    def one_run(self):
        time_steps = self.world_obj.config_obj.time_steps
        Time.new_world()
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        self.reset_start_infection(agents_obj)
        locations_obj = ReadLocations(self.world_obj.locations_filename,
                                      self.world_obj.config_obj)
        oneTimeEvent_obj = ReadOneTimeEvents(
            self.world_obj.one_time_event_file)

        sim_obj = Simulate(self.world_obj.config_obj, self.world_obj.model,
                           self.world_obj.policy_list,
                           self.world_obj.event_restriction_fn, agents_obj,
                           locations_obj)
        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()
            self.update_agent_start_infection(agents_obj, current_time_step)
            Time.increment_current_time_step()

        self.update_agent_scores(agents_obj)

        end_state = sim_obj.endSimulation()
        return end_state, agents_obj, locations_obj

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

    def get_maximum_vulnerable_agents(self, n):
        res_max = dict(
            sorted(self.agent_scores.items(),
                   key=lambda x: (x[1], x[0]),
                   reverse=True)[:n])
        return res_max

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

    def print_default_output(self, n):
        print(self.get_maximum_vulnerable_agents(n))
        print(self.get_minimum_vulnerable_agents(n))


class AgentVulnerabilitiesMC(AgentVulnerabilities):
    """
    Parameters to pass : states, num_runs, num_agents_to_remove
    """
    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_agents_to_remove = parameter_dict['num_agents_to_remove']
        self.agents_to_remove = None
        self.agent_counts = {}
        self.init_scores()

    def init_scores(self):
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        for agent_index in agents_obj.agents.keys():
            self.agent_scores[agent_index] = 0.0
            self.agent_counts[agent_index] = 0

    def remove_agents(self, agents_obj, num_agents_to_remove):
        agents = list(agents_obj.agents)
        self.agents_to_remove = random.sample(agents, num_agents_to_remove)
        for agent in self.agents_to_remove:
            agents_obj.agents.pop(agent)
            self.agent_counts[agent] += 1

    def update_agent_scores(self, end_state):
        agent_score = 0
        for state in self.states:
            agent_score += end_state[state][-1]
        total_agents = 0
        for key in end_state:
            total_agents += end_state[key][-1]
        score = 1.0 - agent_score / total_agents
        for agent in self.agents_to_remove:
            # self.agent_scores[agent] += score
            self.agent_scores[agent] += 1.0 / (self.agent_counts[agent]) * (
                score - self.agent_scores[agent])

    def one_run(self):
        Time.new_world()
        time_steps = self.world_obj.config_obj.time_steps
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        if (self.num_agents_to_remove > 0):
            self.remove_agents(agents_obj, self.num_agents_to_remove)
        locations_obj = ReadLocations(self.world_obj.locations_filename,
                                      self.world_obj.config_obj)
        oneTimeEvent_obj = ReadOneTimeEvents(
            self.world_obj.one_time_event_file)

        sim_obj = Simulate(self.world_obj.config_obj, self.world_obj.model,
                           self.world_obj.policy_list,
                           self.world_obj.event_restriction_fn, agents_obj,
                           locations_obj)
        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.increment_current_time_step()

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

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

    def get_maximum_agent_vulnerability(self, n):
        res_max = dict(
            sorted(self.agent_scores.items(),
                   key=lambda x: (x[1], x[0]),
                   reverse=True)[:n])
        return res_max

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

    def print_default_output(self, n):
        print(self.get_maximum_agent_vulnerability(n))
        print(self.get_minimum_agent_vulnerability(n))


class AgentVulnerabilitiesMC2(AgentVulnerabilities):
    """
    Parameters - states, num_runs, num_agents_to_remove, chunk_len_range (default  = [1,50])
    """
    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_agents_to_remove = parameter_dict['num_agents_to_remove']
        self.agents_to_remove = None

        self.time_steps = self.world_obj.config_obj.time_steps
        self.start_agent_score = None
        self.end_agent_score = None
        chunk_len_range = []
        if ('chunk_len_range' not in parameter_dict.keys()
                or not parameter_dict['chunk_len_range']):
            chunk_len_range = [1, 50]
        else:
            chunk_len_range = parameter_dict['chunk_len_range']
        self.range_list = list(
            range(int(chunk_len_range[0]),
                  int(chunk_len_range[1]) + 1))
        self.init_scores()

    def init_scores(self):
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        for agent_index in agents_obj.agents.keys():
            self.agent_scores[agent_index] = 0

    def remove_agents(self, agents_obj, num_agents_to_remove):
        agents = list(agents_obj.agents)
        self.agents_to_remove = random.sample(agents, num_agents_to_remove)
        for agent in self.agents_to_remove:
            agents_obj.agents.pop(agent)

    def reset_world(self, agents_obj):
        self.start_agent_score = None
        self.end_agent_score = None

    def get_scores(self, agents_obj):
        score = 0
        for agent in agents_obj.agents.values():
            for state in self.states:
                if agent.state == state:
                    score += 1
                    break
        return score

    def update_agent_scores(self, agents_obj, time_step, mode):
        if (mode == 'start'):
            self.start_agent_score = self.get_scores(agents_obj)
        elif (mode == 'end'):
            self.end_agent_score = self.get_scores(agents_obj)
            for agent in self.agents_to_remove:
                self.agent_scores[
                    agent] += self.end_agent_score - self.start_agent_score

    def one_run(self, start, end):

        Time.new_world()
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        self.reset_world(agents_obj)
        if (self.num_agents_to_remove > 0):
            self.remove_agents(agents_obj, self.num_agents_to_remove)
        locations_obj = ReadLocations(self.world_obj.locations_filename,
                                      self.world_obj.config_obj)
        oneTimeEvent_obj = ReadOneTimeEvents(
            self.world_obj.one_time_event_file)

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

        flag = 0
        for current_time_step in range(self.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()
            if (current_time_step == start):
                self.update_agent_scores(agents_obj, current_time_step,
                                         'start')
                flag += 1
            if (current_time_step == end):
                self.update_agent_scores(agents_obj, current_time_step, 'end')
                flag += 1
                break
            Time.increment_current_time_step()

        end_state = sim_obj.endSimulation()
        assert flag == 2

    def generate_bounds(self):
        chunk_len = random.choice(self.range_list)

        if (chunk_len <= 0):
            raise ValueError(
                'Inappropriate chunk len range provided. Please provide positive ranges.'
            )
        if (chunk_len >= self.time_steps):
            return 0, self.time_steps - 1

        start_list = list(range(self.time_steps - chunk_len + 1))
        start = random.choice(start_list)
        end = start + chunk_len - 1
        return start, end

    def run_detection(self):
        for i in range(self.num_runs):
            if (i + 1) % (self.num_runs / 10) == 0:
                print('Iteration running : ', i + 1)
            start, end = self.generate_bounds()
            if start is None:
                continue
            self.one_run(start, end)

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

    def get_minimum_agent_vulnerability(self, n):
        res_min = dict(
            sorted(self.agent_scores.items(),
                   key=lambda x: (x[1], x[0]),
                   reverse=True)[:n])
        return res_min

    def print_default_output(self, n):
        print(self.get_maximum_agent_vulnerability(n))
        print(self.get_minimum_agent_vulnerability(n))


class BanditAlgos(AgentVulnerabilities):
    """
    Parameters - states, num_runs, mode (default = "EPS")
    """
    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']
        if ('mode' not in parameter_dict.keys() or not parameter_dict['mode']):
            self.mode = 'EPS'
        else:
            self.mode = parameter_dict['mode']
        self.agent_counts = {}
        self.eps = 0.01
        self.rm_agent = None
        self.init_scores()

    def init_scores(self):
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        for agent_index in agents_obj.agents.keys():
            # self.agent_scores[int(agent_index)] = 1.0
            self.agent_scores[int(agent_index)] = self.get_init_score(
                agent_index)
            self.agent_counts[int(agent_index)] = 0

    def get_init_score(self, agent_index):
        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        agents_obj.agents.pop(agent_index)
        end_state, _, _ = self.one_run_helper(agents_obj)
        return self.get_score(end_state)

    def remove_agents(self, agents_obj):

        agents = list(agents_obj.agents)
        if (self.mode == 'EPS'):
            if random.random() < self.eps:
                self.rm_agent = random.randint(0, len(agents) - 1)
            else:
                self.rm_agent = max(range(len(agents)),
                                    key=lambda x: self.agent_scores[x])

        elif (self.mode == 'UCB'):
            self.rm_agent = max(range(len(agents)),
                                key=lambda x: self.agent_scores[x] + np.
                                sqrt(2 * np.log(self.t) /
                                     (1 + self.agent_counts[x])))

        else:
            raise Exception('Enter Valid mode.')

        agents_obj.agents.pop(str(self.rm_agent))
        self.agent_counts[self.rm_agent] += 1

    def get_score(self, end_state):
        agent_score = 0
        for state in self.states:
            agent_score += end_state[state][-1]
        total_agents = 0
        for key in end_state:
            total_agents += end_state[key][-1]
        score = 1.0 - agent_score / total_agents
        return score

    def update_agent_scores(self, end_state):
        score = self.get_score(end_state)
        self.agent_scores[self.rm_agent] += 1.0 / (self.agent_counts[
            self.rm_agent]) * (score - self.agent_scores[self.rm_agent])

    def one_run_helper(self, agents_obj):
        Time.new_world()
        time_steps = self.world_obj.config_obj.time_steps
        locations_obj = ReadLocations(self.world_obj.locations_filename,
                                      self.world_obj.config_obj)
        oneTimeEvent_obj = ReadOneTimeEvents(
            self.world_obj.one_time_event_file)

        sim_obj = Simulate(self.world_obj.config_obj, self.world_obj.model,
                           self.world_obj.policy_list,
                           self.world_obj.event_restriction_fn, agents_obj,
                           locations_obj)
        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.increment_current_time_step()

        end_state = sim_obj.endSimulation()
        return end_state, agents_obj, locations_obj

    def one_run(self):

        agents_obj = ReadAgents(self.world_obj.agents_filename,
                                self.world_obj.config_obj)
        self.remove_agents(agents_obj)
        end_state, _, _ = self.one_run_helper(agents_obj)
        self.update_agent_scores(end_state)

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

    def get_maximum_agent_vulnerability(self, n):
        res_max = dict(
            sorted(self.agent_scores.items(),
                   key=lambda x: (x[1], x[0]),
                   reverse=True)[:n])
        return res_max

    def print_default_output(self, n):
        print(self.get_maximum_agent_vulnerability(n))
