import argparse
import numpy as np
import torch
import time
import pickle
import os
import matplotlib.pyplot as plt

from attacks.action import softmax_attack, random_attack, guassian_attack
from trainer.maddpg.maddpg import MADDPGAgentTrainer
from trainer.maddpg.maddpg_infer import MADDPGInferAgentTrainer
from trainer.maddpg.maddpg_merge import MADDPGMergeAgentTrainer
from utils.common import save_state, restore_state, latest_checkpoint, generate_checkpoint_path, get_device

# torch.manual_seed(0)

def mkdir(dirs):
    for directory in dirs:
        dirs = directory.split('/')
        _dirs = '.'
        for _dir in dirs:
            _dirs = os.path.join(_dirs, _dir)
            if not os.path.exists(_dirs):
                os.mkdir(_dirs)

def parse_args():
    parser = argparse.ArgumentParser(
        "Reinforce Learning for Multi-Agent Environments")

    # Environment
    parser.add_argument("--scenario", type=str, default="simple", help="name of the scenario script")
    parser.add_argument("--max-episode-len", type=int, default=25, help="maximum episode length")
    parser.add_argument("--num-episodes", type=int, default=60000, help="number of episodes")
    # parser.add_argument("--num-adversaries", type=int, default=0, help="number of adversaries")
    parser.add_argument("--policy", type=str, default="maddpg", help="policy for agents")
    parser.add_argument("--num-agents", type=int, default=3, help="number of agents")

    # Core training parameters
    parser.add_argument("--lr", type=float, default=1e-2, help="learning rate for Adam optimizer")
    parser.add_argument("--gamma", type=float, default=0.95, help="discount factor")
    parser.add_argument("--batch-size", type=int, default=1024, help="number of episodes to optimize at the same time")
    parser.add_argument("--num-units", type=int, default=64, help="number of units in the mlp")
    parser.add_argument("--grad-norm-clip", type=float, default=0.5, help="max norm for clip gradients")
    parser.add_argument('--use-gpu', action='store_true', default=False)
    parser.add_argument('--gpu-id', type=int, default=0, help='id of gpu')
    parser.add_argument('--buffer-size', type=int, default=1e6, help='size of replay buffer')
    parser.add_argument('--depth', type=int, default=3, help='depth of MLP')
    parser.add_argument('--update-gap', type=int, default=100, help='amount of samples for online fashion')
    parser.add_argument('--beta', type=float, default=1e-2, help='factor for online fashion')

    # Checkpointing
    parser.add_argument("--exp-name", type=str, default="test", help="name of the experiment")
    parser.add_argument("--save-dir", type=str, default="./save_dir/", help="directory in which training state and model should be saved")
    parser.add_argument("--save-rate", type=int, default=1000, help="save model once every time this many episodes are completed")
    parser.add_argument("--load-dir", type=str, default="", help="directory in which training state and model are loaded")
    parser.add_argument("--summary-dir", type=str, default="", help="directory in which the summary of tranining process should be saved")
    parser.add_argument("--plot", action="store_true", default=False)
    parser.add_argument('--run_index', type=int, default=0, help="index of which run")

    # Evaluation
    parser.add_argument("--restore", action="store_true", default=False)
    parser.add_argument("--display", action="store_true", default=False)
    parser.add_argument("--benchmark", action="store_true", default=False)
    parser.add_argument("--benchmark-iters", type=int, default=100000, help="number of iterations run for benchmarking")
    parser.add_argument("--benchmark-dir", type=str, default="./benchmark_files/", help="directory where benchmark data is saved")
    parser.add_argument("--plots-dir", type=str, default="./learning_curves/", help="directory where plot data is saved")
    return parser.parse_args()

def make_env(scenario_name, arglist, benchmark=False):
    from multiagent.environment import MultiAgentEnv
    import multiagent.scenarios as scenarios

    # load scenario from script
    # scenario = scenarios.load(scenario_name + ".py").Scenario()
    from scenarios.spread import Scenario
    # from scenarios.spread_policy_reg import Scenario
    scenario = Scenario()

    # create world
    world = scenario.make_world()

    # create multiagent environment
    if benchmark:
        env = MultiAgentEnv(world, scenario.reset_world, scenario.reward,
                            scenario.observation, scenario.benchmark_data)
    else:
        env = MultiAgentEnv(world, scenario.reset_world,
                            scenario.reward, scenario.observation)

    return env

def get_trainers(env, num_adversaries, obs_shape_n, arglist):
    trainers = []
    if arglist.policy == 'maddpg':
        trainer = MADDPGAgentTrainer
    elif arglist.policy == 'maddpg_pi':
        trainer = MADDPGInferAgentTrainer
    elif arglist.policy == 'maddpg_si':
        trainer = MADDPGMergeAgentTrainer
    else:
        trainer = MADDPGAgentTrainer
    for i in range(num_adversaries):
        trainer_instance = train()
        trainers.append(trainer_instance)
    for i in range(num_adversaries, env.n):
        trainer_instance = trainer("agent_%d" % i, obs_shape_n, env.action_space,
                                   i, arglist, local_q=(arglist.policy == 'ddpg'))
        trainers.append(trainer_instance)
    return trainers

def run(arglist):
    # prepare save dir
    # print(arglist.save_dir)

    arglist.device = get_device(arglist.use_gpu, arglist.gpu_id)

    # save params of arglist
    arglist_dict = vars(arglist)
    with open(os.path.join(arglist.save_dir, 'params.txt'), 'w', encoding='utf-8') as f:
        keys = list(arglist_dict.keys())
        keys.sort()
        for k in keys:
            f.write(k + ':\n\t' + str(arglist_dict[k]) + '\n')

    # Create Environment
    env = make_env(arglist.scenario, arglist, arglist.benchmark)

    # create agent trainers
    obs_shape_n = [env.observation_space[i].shape for i in range(env.n)]
    # num_adversaries = min(env.n, arglist.num_adversaries)
    num_adversaries = 0
    trainers = get_trainers(env, num_adversaries, obs_shape_n, arglist)
    print('Policy:{}'.format(arglist.policy))

    # Prepare necessary data recorders
    mean_episode_rewards = []
    episode_rewards = [0.0]                         # sum of rewards for all agents
    final_ep_rewards = []                           # sum of rewards for training curve
    final_ep_ag_rewards = []                        # agent rewards for training curve
    agent_rewards = [[0.0] for _ in range(env.n)]   # individual agent reward
    agent_info = [[[]]]                             # placeholder for benchmarking info

    # Run experiments
    print('Start iterations...')
    obs_n = env.reset()
    episode_step = 0    # amount of steps inside one episode
    num_episodes = 0    # amount of episodes
    train_step = 0      # amount of training steps
    t_start = time.time()

    # Load previous results if necessary
    if arglist.load_dir == '':
        arglist.load_dir = arglist.save_dir
    if arglist.display or arglist.restore or arglist.benchmark:
        print('Loading previous states...')
        path, latest_index = latest_checkpoint(arglist.load_dir)
        restore_state(trainers, path, map_location=arglist.device)
        num_episodes = latest_index
    if arglist.display or arglist.benchmark:
        for agent in trainers:
            agent.eval()
    else:
        for agent in trainers:
            agent.train()
            # agent.sync_target_nets()

    while True:
        # Get action
        action_n = [agent.action(obs) for agent, obs in zip(trainers, obs_n)]

        # Evolution of the environment
        new_obs_n, reward_n, done_n, info_n = env.step(action_n)

        episode_step += 1
        done = all(done_n)
        terminal = (episode_step >= arglist.max_episode_len)

        # Collect experience for replay buffer
        for i, agent in enumerate(trainers):
            agent.experience(obs_n[i], action_n[i], reward_n[i],
                             new_obs_n[i], done_n[i], terminal)
        obs_n = new_obs_n

        for i, reward in enumerate(reward_n):
            episode_rewards[-1] += reward
            agent_rewards[i][-1] += reward

        if done or terminal:
            num_episodes += 1
            obs_n = env.reset()
            episode_step = 0
            episode_rewards.append(0.0)
            for a in agent_rewards:
                a.append(0.0)
            agent_info.append([[]])

        # Increment global train step counter
        train_step += 1

        # # for benchmarking learned policies
        # if arglist.benchmark:
        #     for i, info in enumerate(info_n):
        #         agent_info[-1][i].append(info_n['n'])
        #     if train_step > arglist.benchmark_iters and (done or terminal):
        #         file_name = arglist.benchmark_dir + arglist.exp_name + '.pkl'
        #         print('Finished benchmarking, now saving...')
        #         with open(file_name, 'wb') as f:
        #             pickle.dump(agent_info[:-1],f)
        #         break3, q-value: -115.43821716308594, p-loss:115.20866394042969, time: 134.907
        #     continue

        # For displaying learned policies
        if arglist.display:
            time.sleep(0.1)
            env.render()
            continue

        # Update parameters of all trainers if not in display or benchmark mode
        info = None
        for agent in trainers:
            agent.preupdate()   # reset sample index of replay buffer
        for agent in trainers:
            info = agent.update(trainers, train_step)

        # Save mode, display training output
        if terminal and ( num_episodes % arglist.save_rate == 0):
            path = generate_checkpoint_path(arglist.save_dir, num_episodes)
            save_state(trainers, path)

            _steps = train_step
            _episodes = num_episodes
            _mean_episode_reward = np.mean(episode_rewards[-arglist.save_rate:])
            mean_episode_rewards.append(_mean_episode_reward)
            q_value = info['q-value']
            p_loss = info['p_loss']
            _time = round(time.time() - t_start, 3)
            if num_adversaries == 0:
                print("steps: {}, epsiodes: {}, mean episode reward: {:.4f}, q-value: {:.4f}, p-loss:{}, time: {}".format(
                    _steps, _episodes, _mean_episode_reward, q_value, p_loss, _time))
            else:
                _agent_episode_reward = [np.mean(reward[-arglist.save_rate:]) for reward in agent_rewards]
                print("steps: {}, epsiodes: {}, mean episode reward: {}, agent episode reward: {}, q-value: {}, p-loss:{}, time: {}".format(
                    _steps, _episodes, _mean_episode_reward,_agent_episode_reward, q_value, p_loss, _time
                ))
            t_start = time.time()
            # Keep track of final episode reward
            final_ep_rewards.append(_mean_episode_reward)
            for reward in agent_rewards:
                final_ep_ag_rewards.append(np.mean(reward[-arglist.save_rate:]))

            # Save final episode reward for plotting training curve later
            if num_episodes >= arglist.num_episodes:
                np.savetxt(os.path.join(arglist.save_dir, 'episode_reward.txt'), episode_rewards[arglist.save_rate:-1], fmt='%.1f')
                plt.plot(list(range(1, len(episode_rewards) - arglist.save_rate)), episode_rewards[arglist.save_rate:-1])
                plt.savefig(os.path.join(arglist.save_dir, 'episode_reward.png'))
                np.savetxt(os.path.join(arglist.save_dir, 'mean_episode_rewards.txt'), mean_episode_rewards, fmt='%.1f')
                plt.clf()
                plt.plot(list(range(1, len(mean_episode_rewards) + 1)), mean_episode_rewards)
                plt.savefig(os.path.join(arglist.save_dir, 'mean_episode_rewards.png'))
                # reward_file_name = arglist.plots_dir + arglist.exp_name + '_rewards.pkl'
                # with open(reward_file_name, 'wb') as f:
                #     pickle.dump(final_ep_rewards, f)

                # ag_rewards_file_name = arglist.plots_dir + arglist.exp_name + '_agrewards.pkl'
                # with open(ag_rewards_file_name, 'wb') as f:
                #     pickle.dump(final_ep_ag_rewards, f)
                print('...Finished tottal of {} episodes'.format(len(episode_rewards)))
                break

if __name__ == "__main__":
    for i in range(5):
        np.random.seed(seed=i)
        arglist = parse_args()
        arglist.run_index = i
        save_dir = './save_dir/{}/agent_{}/run_{}/'.format(arglist.policy, arglist.num_agents, arglist.run_index)
        mkdir([save_dir,])
        arglist.save_dir = save_dir
        run(arglist)


# Richie Pajak independent study

import csv
import math
import os
import pandas as pd
import pygame
import random
import tkinter as tk

from pygame.locals import (
    K_ESCAPE,
    KEYDOWN,
    QUIT,
)

SCREEN_WIDTH = 1500
SCREEN_HEIGHT = 900

def find_empty_home(): # This function returns the first found empty home
    for x in homes:
        if not x.owned:
            x.owned = True
            return x
    return None

def find_home_location(): # This funciton calculates the proper location of a new house using the number of houses that already exist
    row = len(homes) // 5
    col = len(homes) % 5
    rowcoord = (row * 100) + 100
    colcoord = (col * 100) + 525
    return colcoord, rowcoord

def find_business_location(): # This funciton calculates the proper location of a new business using the number of businesses that already exist
    row = len(businesses) // 2
    col = len(businesses) % 2
    rowcoord = (row * 100) + 100
    colcoord = (col * 100) + 75
    return colcoord, rowcoord

def create_home():
    new_home = Home()
    homes.add(new_home)
    all_sprites.add(new_home)

def create_agent():
    new_agent = Agent()
    agents.add(new_agent)
    all_sprites.add(new_agent)

def create_business(type):
    new_business = Business(type)
    businesses.add(new_business)
    all_sprites.add(new_business)

class Agent(pygame.sprite.Sprite):
    def __init__(self):
        super(Agent, self).__init__()
        self.size = 15
        self.surf = pygame.Surface((self.size, self.size))
        self.surf.fill((random.randint(0, 200), random.randint(0, 200), random.randint(0, 200)))
        self.rect = self.surf.get_rect()
        self.food = random.randint(5, 10)
        self.water = random.randint(5, 10)
        self.money = 100
        self.home = find_empty_home()
        self.destination = self.home
        self.at_home = True
        self.consumer = True
        self.business_options = []

        if self.home is None:
            self.rect.center = (0, 0)
        else:
            self.rect.center = self.home.rect.center

        self.curr_prio = None

        self.work_prio = -1 * self.money

        self.prio_list = {
            'food': -1 * self.food,
            'water': -1 * self.water
        }

    def update(self, businesses): # The update function is called each tick of the simulation. It handles everything an agent needs to do each time the simulation progresses

        if self.curr_prio == 'home':
            self.destination = self.home
        elif self.curr_prio == 'work':
            for x in businesses:
                if x.being_worked() is False:
                    self.destination = x
                    x.set_worked(True)
                    break
        else:
            if len(self.business_options) != 0:
                min_price = float('inf')
                best_option = None
                for x in self.business_options:
                    if x.get_sell_price() < min_price:
                        if x.get_product_amount() > 0:
                            min_price = x.get_sell_price()
                            best_option = x
                if best_option is None:
                    best_option = self.home
                    self.curr_prio = 'home'
                self.destination = best_option
                if self.rect.colliderect(self.destination.rect) and businesses.has(self.destination):
                    if self.buy(self.destination):
                        self.curr_prio = 'home'
                        self.business_options.clear()
                    else:
                        self.business_options.pop(0)
            else:
                self.curr_prio = 'home'

        self.move()

        if self.rect.center == self.home.rect.center:
            self.at_home = True
        else:
            self.at_home = False

    def determine_business(self, businesses): # This function creates a list of all businesses that sell the product the agent is interested in
        for x in businesses:
            if x.product_type == self.curr_prio:
                self.business_options.append(x)

    def buy(self, business): # This funciton causes agents to attempt to purchase from a business
        num_purchases = 4
        success = False
        for x in range(0, num_purchases):
            if self.money >= business.get_sell_price():
                if business.sell():
                    success = True
                    self.spend_money(business.get_sell_price())
                    self.gain_product(business.product_type, 1)
        return success

    def spend_money(self, amount):
        self.money = self.money - amount

    def gain_money(self, amount):
        self.money = self.money + amount

    def gain_product(self, type, amount): # This function causes an agent to gain food or water
        if type == 'food':
            self.food += amount
        if type == 'water':
            self.water += amount

    def calc_prios(self): # This function causes an agent's prios to be an inverse of their food, water, and money
        #print(self.food, self.water, self.money)
        self.prio_list['food'] = 20 - self.food
        self.prio_list['water'] = 20 - self.water
        self.work_prio = 100 - self.money

    def lose_products(self): # This function is called at the end of each day cycle and causes agents to expend 1 food and 1 water
        if self.food > 0:
            self.food = self.food - 1
        if self.water > 0:
            self.water = self.water - 1

    def get_highest_prio(self): # This function determines an agent's highest priority and sets it to curr_prio
        prio_options = ['food']
        for x in self.prio_list:
            if self.prio_list.get(x) > self.prio_list[prio_options[0]]:
                prio_options.clear()
                prio_options.append(x)
            elif self.prio_list.get(x) == self.prio_list[prio_options[0]]:
                prio_options.append(x)
        index = random.randint(0, len(prio_options) - 1)
        self.curr_prio = prio_options[index]

    def set_curr_prio(self, prio):
        self.curr_prio = prio

    def get_curr_prio(self):
        return self.curr_prio

    def get_work_prio(self):
        return self.work_prio

    def set_worker(self):
        self.consumer = False

    def set_consumer(self):
        self.consumer = True

    def is_at_home(self):
        if self.at_home is True:
            return True
        else:
            return False

    def move(self):
        direction = self.get_direction(self.destination.rect.centerx, self.destination.rect.centery)
        self.rect.move_ip(direction)

    def get_direction(self, x, y):  # This function returns the direction that an agent should move based on the
        # coordinates of its destination that are passed in
        direction_x = 0
        direction_y = 0
        if self.rect.centerx < x:
            direction_x = 1
        if self.rect.centerx > x:
            direction_x = -1
        if self.rect.centery < y:
            direction_y = 1
        if self.rect.centery > y:
            direction_y = -1

        direction = (direction_x, direction_y)
        return direction

class Business(pygame.sprite.Sprite):
    def __init__(self, type):
        super(Business, self).__init__()
        self.size = 50
        self.surf = pygame.Surface((self.size, self.size))
        self.surf.fill((random.randint(0, 200), random.randint(0, 200), random.randint(0, 200)))
        self.rect = self.surf.get_rect()
        self.rect.center = find_business_location()
        self.product_amount = 15
        self.product_type = type
        self.sell_price = 5
        self.money = 0
        self.production_amount = 15
        self.is_worked = False
        self.worker = None
        self.can_be_worked = True

    def update(self):
        if self.product_type == 'food':
            self.surf.fill((255, 0, 0))
        elif self.product_type == 'water':
            self.surf.fill((0, 255, 0))
        else:
            self.surf.fill((0, 0, 255))
        return

    def get_sell_price(self):
        return self.sell_price

    def get_product_amount(self):
        return self.product_amount

    def sell(self): # This function is called when an agent attempts to buy a product, it determines if this business is able to sell and causes the business to gain money and lose a product
        if self.product_amount > 0 and self.can_be_worked:
            self.money = self.money + self.sell_price
            self.product_amount = self.product_amount - 1
            return True
        return False

    def give_profits(self): # This function gives all profits made in the day to the worker who is working the business
        if self.can_be_worked:
            self.worker.gain_money(self.money)
            self.money = 0

    def produce(self): # This funciton causes the business to produce new products
        if self.can_be_worked:
            self.product_amount = self.production_amount
        print(self.product_amount)

    def find_worker(self, agents): # This function finds the agent with the highest work_prio that is not already working and sets it to be working at this business
        if self.can_be_worked:
            current_worker = None
            prio = float('-inf')
            for x in agents:
                if x.get_work_prio() > prio and x.get_curr_prio() != 'work':
                    current_worker = x
                    prio = x.get_work_prio()
            current_worker.set_worker()
            current_worker.set_curr_prio('work')
            return current_worker
        return None

    def being_worked(self):
        return self.is_worked

    def set_worked(self, status):
        self.is_worked = status

    def clear_worker(self):
        self.worker = None
        self.is_worked = False

    def price_change(self): # This function changes the price of goods at the end of each day
        if self.can_be_worked:
            if self.product_amount > 0:
                if self.sell_price > 1:
                    self.sell_price = self.sell_price - 1
                    print("decrease", self.sell_price)
            else:
                self.sell_price = self.sell_price + 1
                print("increase", self.sell_price)

    def set_production_amount(self, amount): # This function changes the production amount of a business when a disaster occurs
        amount = float(amount) / 100
        amount = 1 - amount
        self.production_amount = int(math.ceil(self.production_amount * amount))

    def reset_production_amount(self):
        self.production_amount = 15

class Home(pygame.sprite.Sprite):
    def __init__(self):
        super(Home, self).__init__()
        self.size = 35
        self.surf = pygame.Surface((self.size, self.size))
        self.surf.fill((random.randint(0, 200), random.randint(0, 200), random.randint(0, 200)))
        self.rect = self.surf.get_rect()
        self.owned = False
        self.rect.center = find_home_location()

disaster_stats = {
    'disaster_type': '',
    'disaster_start_day': 0,
    'disaster_end_day': 0,
    'disaster_severity': 0
}

print(disaster_stats)

relief_stats = {
    'relief_type': '',
    'relief_start_day': 0,
    'relief_severity': 0
}

# This section of code handles the gui that is used at the beginning of the program to allow users to set the parameters of the simulation
def disaster_apply():
    disaster_stats['disaster_type'] = dropdown_result.get()
    disaster_stats['disaster_start_day'] = int(e1.get())
    disaster_stats['disaster_end_day'] = int(e2.get())
    disaster_stats['disaster_severity'] = int(e3.get())
    print(disaster_stats)

def relief_apply():
    relief_stats['relief_type'] = dropdown_result.get()
    relief_stats['relief_start_day'] = int(e4.get())
    relief_stats['relief_severity'] = int(e5.get())

disaster_menu = tk.Tk()
tk.Label(disaster_menu, text="Disaster Type").grid(row=0)
tk.Label(disaster_menu, text="Disaster Start Day").grid(row=1)
tk.Label(disaster_menu, text="Disaster End Day").grid(row=2)
tk.Label(disaster_menu, text="Disaster Severity").grid(row=3)

disaster_list = [
    'None',
    'Pandemic',
    'Famine',
    'Drought',
    'Tornado',
]

dropdown_result = tk.StringVar(disaster_menu)
dropdown_result.set(disaster_list[0])  # default value

disaster_dropdown = tk.OptionMenu(disaster_menu, dropdown_result, *disaster_list)

e1 = tk.Entry(disaster_menu)
e2 = tk.Entry(disaster_menu)
e3 = tk.Entry(disaster_menu)
e1.insert(10, '0')
e2.insert(10, '0')
e3.insert(10, '0')

disaster_dropdown.grid(row=0, column=1)
e1.grid(row=1, column=1)
e2.grid(row=2, column=1)
e3.grid(row=3, column=1)

tk.Button(disaster_menu, text='Apply', command=disaster_apply).grid(row=6,
                                                                    column=0,
                                                                    sticky=tk.W,
                                                                    pady=4)
tk.Button(disaster_menu,
          text='Next',
          command=disaster_menu.quit).grid(row=6,
                                           column=1,
                                           sticky=tk.W,
                                           pady=4)

disaster_menu.mainloop()

relief_menu = tk.Tk()

tk.Label(relief_menu, text="Relief Type").grid(row=0)
tk.Label(relief_menu, text="Relief Day").grid(row=1)
tk.Label(relief_menu, text="Relief Amount").grid(row=2)

relief_list = [
    'None',
    'Stimulus',
    'Aid'
]

dropdown_result = tk.StringVar(relief_menu)
dropdown_result.set(relief_list[0])  # default value

relief_dropdown = tk.OptionMenu(relief_menu, dropdown_result, *relief_list)
e4 = tk.Entry(relief_menu)
e5 = tk.Entry(relief_menu)
e4.insert(10, '0')
e5.insert(10, '0')

relief_dropdown.grid(row=0, column=1)
e4.grid(row=1, column=1)
e5.grid(row=2, column=1)

tk.Button(relief_menu, text='Apply', command=relief_apply).grid(row=6,
                                                                column=0,
                                                                sticky=tk.W,
                                                                pady=4)
tk.Button(relief_menu,
          text='Next',
          command=relief_menu.quit).grid(row=6,
                                         column=1,
                                         sticky=tk.W,
                                         pady=4)

relief_menu.mainloop()

tk.mainloop()

# This section initializes all necessary variables for the simulation
pygame.init()
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
font = pygame.font.SysFont(None, 24)
clock = pygame.time.Clock()

all_sprites = pygame.sprite.Group()
homes = pygame.sprite.Group()
agents = pygame.sprite.Group()
businesses = pygame.sprite.Group()

num_homes = 40
num_agents = 40
num_businesses = 10
num_food_businesses = 0
num_water_businesses = 0
day_count = 0
for x in range(0, num_homes):
    create_home()
for x in range(0, num_agents):
    create_agent()
for x in range(0, num_businesses):
    if x % 2 == 0:
        create_business('food')
        num_food_businesses += 1
    else:
        create_business('water')
        num_water_businesses += 1

phase = 0
all_consumers_at_home = True
all_agents_at_work = True
all_agents_at_home = True

# This section handles the writing of headers to the csv files used to collect data
filename = "simulationdata.csv"
f = open(filename, "w+")
f.close()

filename = "agentdata.csv"
f = open(filename, "w+")
f.close()

with open('simulationdata.csv', 'a', newline='') as csvfile:
    mywriter = csv.writer(csvfile, delimiter=',', lineterminator='\n')
    mywriter.writerow(('average food price', 'average water price', 'average price'))

header = []
for x in range(0, num_agents):
    header.append('Agent_{}'.format(x))
    header.append('Agent_{}'.format(x))
    header.append('Agent_{}'.format(x))

header2 = []
for x in range(0, num_agents):
    header2.append('Money')
    header2.append('Food')
    header2.append('Water')

with open('agentdata.csv', 'a', newline='') as csvfile:
    mywriter = csv.writer(csvfile, delimiter=',', lineterminator='\n')
    mywriter.writerow(header)
    mywriter.writerow(header2)

# This section of the code is the main loop that handles all daily events in the simulation
running = True
while running:
    for event in pygame.event.get():
        if event.type == KEYDOWN:
            if event.key == K_ESCAPE: # This section handles data collection when the simulation ends
                running = False
                filepath = os.path.dirname(os.path.abspath(__file__))

                read_file = pd.read_csv(filepath + '/simulationdata.csv')
                read_file.to_excel(filepath + '/simulationdata.xlsx', index=None, header=True)
                read_file2 = pd.read_csv(filepath + '/agentdata.csv')
                read_file2.to_excel(filepath + '/agentdata.xlsx', index=None, header=True)
        elif event.type == QUIT:
            print('quitting successfully')
            running = False
            filepath = os.path.dirname(os.path.abspath(__file__))

            read_file = pd.read_csv(filepath + '/simulationdata.csv')
            read_file.to_excel(filepath + '/simulationdata.xlsx', index=None, header=True)
            read_file2 = pd.read_csv(filepath + '/agentdata.csv')
            read_file2.to_excel(filepath + '/agentdata.xlsx', index=None, header=True)
    if phase == 0: # Phase 0: Workers assigned to businesses
        day_count += 1
        for x in agents:
            x.get_highest_prio()
        for x in businesses:
            x.worker = x.find_worker(agents)
        phase = 1
    elif phase == 1: # Phase 1: Workers travel to businesses
        businesses.update()
        for x in agents:
            if not x.consumer:
                x.update(businesses)
        all_agents_at_work = True
        for x in agents:
            if not x.consumer and not x.rect.colliderect(x.destination.rect):
                all_agents_at_work = False
        for x in agents:
            if x.rect.colliderect(x.destination.rect) and x.consumer is False:
                x.rect.center = x.destination.rect.center
        if all_agents_at_work is True:
            phase = 2
            for x in agents:
                x.determine_business(businesses)
    elif phase == 2: # Phase 2: Consumers purchase from businesses
        for x in agents:
            if x.consumer:
                x.update(businesses)

        all_consumers_at_home = True
        for x in agents:
            if x.consumer:
                if x.is_at_home() is False:
                    all_consumers_at_home = False
        if all_consumers_at_home:
            phase = 3
    elif phase == 3: # Phase 3: Workers return home
        for x in agents:
            x.set_curr_prio('home')
            x.update(businesses)
        all_agents_at_home = True
        for x in agents:
            if x.is_at_home() is False:
                all_agents_at_home = False
        if all_agents_at_home:
            phase = 4
    elif phase == 4: # Phase 4: Simulation is prepared for new day
        food_sum = 0
        water_sum = 0
        for x in businesses:
            x.give_profits()
            x.clear_worker()
            x.price_change()
            x.produce()
            if x.product_type == 'food':
                food_sum += x.sell_price
            else:
                water_sum += x.sell_price
            print(x.product_amount, x.product_type)
        for x in agents:
            x.calc_prios()
            x.set_consumer()
            x.set_curr_prio('')
            x.lose_products()
            x.type = 0
        food_average = food_sum / num_food_businesses # Calculate stats to be recorded and write to csv files
        water_average = water_sum / num_water_businesses
        average_price = (food_average + water_average) / 2
        with open('simulationdata.csv', 'a', newline='') as csvfile:
            csv_writer = csv.writer(csvfile, delimiter=',', lineterminator='\n')
            csv_writer.writerow((food_average, water_average, average_price))

        row = []
        for x in agents:
            row.append(x.money)
            row.append(x.food)
            row.append(x.water)
        with open('agentdata.csv', 'a', newline='') as csvfile:
            csv_writer = csv.writer(csvfile, delimiter=',', lineterminator='\n')
            csv_writer.writerow(row)

        if day_count == disaster_stats['disaster_start_day']: # This section handles adjusting simulation parameters for disasters
            if disaster_stats['disaster_type'] == 'Pandemic':
                businesses_closed = 0
                for x in businesses:
                    if businesses_closed < disaster_stats['disaster_severity']:
                        x.can_be_worked = False
                        businesses_closed += 1

            if disaster_stats['disaster_type'] == 'Famine':
                for x in businesses:
                    if x.product_type == 'food':
                        x.set_production_amount(disaster_stats['disaster_severity'])

            if disaster_stats['disaster_type'] == 'Drought':
                for x in businesses:
                    if x.product_type == 'water':
                        x.set_production_amount(disaster_stats['disaster_severity'])

            if disaster_stats['disaster_type'] == 'Tornado':
                for x in businesses:
                    x.set_production_amount(disaster_stats['disaster_severity'])

        if day_count == disaster_stats['disaster_end_day']:
            if disaster_stats['disaster_type'] == 'Pandemic':
                for x in businesses:
                    x.can_be_worked = True
            if disaster_stats['disaster_type'] == 'Famine' or disaster_stats['disaster_type'] == 'Drought' or \
                    disaster_stats['disaster_type'] == 'Tornado':
                for x in businesses:
                    x.reset_production_amount()

        if day_count == relief_stats['relief_start_day']:
            if relief_stats['relief_type'] == 'Stimulus':
                for x in agents:
                    x.gain_money(relief_stats['relief_severity'])
            if relief_stats['relief_type'] == 'Aid':
                for x in agents:
                    x.gain_product('water', relief_stats['relief_severity'])
                    x.gain_product('food', relief_stats['relief_severity'])

        if day_count >= 200: # Ends the simulation if the day count reaches 200
            pygame.event.post(pygame.event.Event(QUIT))
        phase = 0

    screen.fill((255, 255, 255))

    for entity in all_sprites: # Draws all pygame sprites each game tick
        screen.blit(entity.surf, entity.rect)
    img = font.render(str(day_count), True, (0, 0, 0))
    screen.blit(img, (20, 20))
    clock.tick()

    pygame.display.flip()


import logging
import sys

class class_logger(object):
    """descriptor, used to make logger for attribute in first usage"""
    def __init__(self):
        pass

    def __get__(self, instance, owner):
        log = getattr(owner, '__log__', None)
        if log is None:
            name = owner.__module__ + '.' + owner.__name__
            log = logging.getLogger(name)
            owner.__log__ = log
            return log
        return log

def setup_log(args):
    root = logging.getLogger()
    if args.verbose:
        fs = "%(levelname)s:%(name)s:%(message)s"
        root.setLevel(logging.DEBUG)
    else:
        fs = "%(levelname)s %(message)s"
        root.setLevel(logging.INFO)
    handler = logging.StreamHandler(sys.stdout)
    fmt = logging.Formatter(fs)
    handler.setFormatter(fmt)
    root.addHandler(handler)
    #logging.basicConfig(level=logging.DEBUG)


from selenium import webdriver
from depot.manager import DepotManager
from PIL import Image
import os
import sys
sys.path.append("../app")
from settings import *

class Screenshot():

    def __init__(self, url, element):
        self.url = url
        self.element = element

    def get_original_diagram(self):
        depot = DepotManager.get()
        self.driver = webdriver.PhantomJS()
        self.driver.maximize_window()  # set the window size that you need
        self.driver.get(self.url)
        self.ele = self.driver.find_element_by_id(self.element)

        self.loc1 = self.ele.location
        # self.size1 = ele.size

        screenshot_name = self.generate_name()
        print(APP_ROOT + "/static/images")
        self.screenshot_file_location = APP_ROOT + \
            "/static/images/" + screenshot_name + '.png'
        self.driver.save_screenshot(self.screenshot_file_location)

    def crop_screenshot(self):
        image2 = Image.open(self.screenshot_file_location)

        # size is 700x400
        left = self.loc1['x']
        top = self.loc1['y']
        right = self.loc1['x'] + 700
        bottom1 = self.loc1['y'] + 400

        image2 = image2.crop((left, top, right, bottom1))
        os.remove(self.screenshot_file_location)
        image2.save(self.screenshot_file_location)

    def finish(self):
        self.driver.quit()

    def generate_name(self):
        tokens = self.url.split("/")
        print(tokens)
        return tokens[len(tokens) - 1]

    # scatterplot
    # radar_graph
    # lda_graph
    # ingredient_freq
    # word_cloud

if __name__ == "__main__":
    s = Screenshot("http://localhost:5000/lda_graph", "lda")
    s.get_original_diagram()
    s.crop_screenshot()
    s.finish()


import sqlite3, logging, sys

from flask import Flask, jsonify, json, render_template, request, url_for, redirect, flash, Response
from werkzeug.exceptions import abort

# Function to get a database connection.
# This function connects to database with the name `database.db`
def get_db_connection():
    connection = sqlite3.connect('database.db')
    connection.row_factory = sqlite3.Row
    return connection

# Function to get a post using its ID
def get_post(post_id):
    connection = get_db_connection()
    post = connection.execute('SELECT * FROM posts WHERE id = ?',
                        (post_id,)).fetchone()
    connection.close()
    return post

# Function to get total posts
def get_total_posts():
    connection = get_db_connection()
    total_posts = connection.execute('SELECT count(*) FROM posts').fetchone()
    connection.close()
    return total_posts[0]

# Function to get total db connections
def set_new_db_connection():
    connection = get_db_connection()
    connection.execute('UPDATE db_transactions set tx_count = tx_count + 1 WHERE id = 1')
    connection.commit()
    connection.close()

# Function to get total db connections
def get_total_db_connections():
    connection = get_db_connection()
    total_connections = connection.execute('SELECT tx_count FROM db_transactions where id = 1').fetchone()
    connection.close()
    return total_connections[0]

# Define the Flask application
app = Flask(__name__)
app.config['SECRET_KEY'] = 'your secret key'

# Define the main route of the web application
@app.route('/')
def index():
    connection = get_db_connection()
    posts = connection.execute('SELECT * FROM posts').fetchall()
    connection.close()
    return render_template('index.html', posts=posts)

# Define how each individual article is rendered
# If the post ID is not found a 404 page is shown
@app.route('/<int:post_id>')
def post(post_id):
    post = get_post(post_id)
    if post is None:
      app.logger.info('Post does not exist')
      return render_template('404.html'), 404
    else:
      app.logger.info("Article " + post['title'] + " retrieved!")
      set_new_db_connection()
      return render_template('post.html', post=post)

# Define the About Us page
@app.route('/about')
def about():
    app.logger.info('About us page was reached')
    return render_template('about.html')

# Define the Healthz page
@app.route("/healthz")
def healthz():
    app.logger.info('Healthz endpoint was reached')
    return Response(json.dumps({"result": "OK healthy"}), status = 200, mimetype='application/json')

# Define the Metrics page
@app.route("/metrics")
def metrics():
    total_posts = get_total_posts()
    total_connections = get_total_db_connections()
    app.logger.info('Metrics endpoint was reached')
    return Response(json.dumps({"status":"success","code":0, "data":{"db_connection_count": total_connections,"post_count": total_posts}}), status=200, mimetype='application/json')

# Define the post creation functionality
@app.route('/create', methods=('GET', 'POST'))
def create():
    if request.method == 'POST':
        title = request.form['title']
        content = request.form['content']

        if not title:
            flash('Title is required!')
        else:
            connection = get_db_connection()
            connection.execute('INSERT INTO posts (title, content) VALUES (?, ?)',
                         (title, content))
            connection.commit()
            connection.close()
            set_new_db_connection()
            app.logger.info('A new post was created. The title is ' + title)
            return redirect(url_for('index'))

    return render_template('create.html')

# set logger to handle STDOUT and STDERR
stdout_handler = logging.StreamHandler(sys.stdout)
stderr_handler = logging.StreamHandler(sys.stderr)
handlers = [stderr_handler, stdout_handler]

# format output
format_output = '%(levelname)s:%(name)s %(asctime)s : %(message)s'

logging.basicConfig(format=format_output, level=logging.DEBUG, handlers=handlers)

# start the application on port 3111
if __name__ == "__main__":
   app.run(host='0.0.0.0', port='3111', debug=True)