class MoneyModel(Model):
    """A simple model of an economy where agents exchange currency at random.

    All the agents begin with one unit of currency, and each time step can give
    a unit of currency to another agent. Note how, over time, this produces a
    highly skewed distribution of wealth.
    """

    def __init__(self, N, width, height):
        self.num_agents = N
        self.running = True
        self.grid = MultiGrid(height, width, True)
        self.schedule = RandomActivation(self)
        self.datacollector = DataCollector(
            model_reporters={"Gini": compute_gini},
            agent_reporters={"Wealth": lambda a: a.wealth}
        )
        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i, self)
            self.schedule.add(a)
            # Add the agent to a random grid cell
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()

    def run_model(self, n):
        for i in range(n):
            self.step()

class MoneyModel(Model):
    """A model with some number of agents."""
    def __init__(self, N, width, height):
        self.num_agents = N
        self.running = True
        self.grid = MultiGrid(height, width, True)
        self.schedule = RandomActivation(self)
        self.datacollector = DataCollector(model_reporters={"Gini": compute_gini},
                agent_reporters={"Wealth": lambda a: a.wealth})
        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i)
            self.schedule.add(a)
            # Add the agent to a random grid cell
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()
    
    def run_model(self, n):
        for i in range(n):
            self.step()

class SchellingModel(Model):
    '''
    Model class for the Schelling segregation model.
    '''

    def __init__(self, height, width, density, minority_pc, homophily):
        '''
        '''

        self.height = height
        self.width = width
        self.density = density
        self.minority_pc = minority_pc
        self.homophily = homophily

        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=True)

        self.happy = 0
        self.datacollector = DataCollector(
            {"happy": lambda m: m.happy}, # Model-level count of happy agents
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.x, "y": lambda a: a.y}) 

        self.running = True

        # Set up agents
        for x in range(self.width):
            for y in range(self.height):
                if random.random() < self.density:
                    if random.random() < self.minority_pc:
                        agent_type = 1
                    else:
                        agent_type = 0

                    agent = SchellingAgent((x,y), x, y, agent_type)
                    self.grid[y][x] = agent
                    self.schedule.add(agent)

    def get_empty(self):
        '''
        Get a list of coordinate tuples of currently-empty cells.
        '''
        empty_cells = []
        for x in range(self.width):
            for y in range(self.height):
                if self.grid[y][x] is None:
                    empty_cells.append((x, y))
        return empty_cells

    def step(self):
        '''
        Run one step of the model. If All agents are happy, halt the model.
        '''
        self.happy = 0 # Reset counter of happy agents
        self.schedule.step()
        self.datacollector.collect(self)

        if self.happy == self.schedule.get_agent_count():
            self.running = False

class WorldModel(Model):
    def __init__(self, N, width, height):
        self.grid = SingleGrid(height, width, True)

        self.schedule = RandomActivation(self)
        self.num_agents = N
        self.running = True
        
        for i in range(self.num_agents):
            ethnicity = random.choice(Ethnicities)
            a = PersonAgent(unique_id=i,
                            model=self,
                            ethnicity=int(ethnicity)
                            )
            self.schedule.add(a)
            # Add the agent to a random grid cell

            self.grid.position_agent(a)
            
        self.datacollector = DataCollector(
            agent_reporters={
                "Nationalism": lambda a: a.nationalism,
                "X": lambda a: a.pos[0],
                "Y": lambda a: a.pos[1]
            }
        )

    def step(self):
        self.datacollector.collect(self)
        self.schedule.step()

class Money_Model(Model):
    def __init__(self, N, width=50, height=50, torus=True):
        self.num_agents = N
        self.schedule = RandomActivation(self)
        self.grid = MultiGrid(height, width, torus)
        self.create_agents()
        self.dc = DataCollector({"Gini": lambda m: m.compute_gini()},
                               {"Wealth": lambda a: a.wealth})
        self.running = True

    def create_agents(self):
        for i in range(self.num_agents):
            a = Money_Agent(i)
            self.schedule.add(a)
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.dc.collect(self)
        self.schedule.step()
        
    def run_model(self, steps):
        for i in range(steps):
            self.step()
    
    def compute_gini(self):
        agent_wealths = [agent.wealth for agent in self.schedule.agents]
        x = sorted(agent_wealths)
        N = self.num_agents
        B = sum( xi * (N-i) for i,xi in enumerate(x) ) / (N*sum(x))
        return (1 + (1/N) - 2*B)

class BoltzmannWealthModelNetwork(Model):
    """A model with some number of agents."""

    def __init__(self, num_agents=7, num_nodes=10):

        self.num_agents = num_agents
        self.num_nodes = num_nodes if num_nodes >= self.num_agents else self.num_agents
        self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=0.5)
        self.grid = NetworkGrid(self.G)
        self.schedule = RandomActivation(self)
        self.datacollector = DataCollector(
            model_reporters={"Gini": compute_gini},
            agent_reporters={"Wealth": lambda _: _.wealth}
        )

        list_of_random_nodes = self.random.sample(self.G.nodes(), self.num_agents)

        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i, self)
            self.schedule.add(a)
            # Add the agent to a random node
            self.grid.place_agent(a, list_of_random_nodes[i])

        self.running = True
        self.datacollector.collect(self)

    def step(self):
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

    def run_model(self, n):
        for i in range(n):
            self.step()

class ShapesModel(Model):
    def __init__(self, N, width=20, height=10):
        self.running = True
        self.N = N    # num of agents
        self.headings = ((1, 0), (0, 1), (-1, 0), (0, -1))  # tuples are fast
        self.grid = SingleGrid(width, height, torus=False)
        self.schedule = RandomActivation(self)
        self.make_walker_agents()

    def make_walker_agents(self):
        unique_id = 0
        while True:
            if unique_id == self.N:
                break
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            pos = (x, y)
            heading = random.choice(self.headings)
            # heading = (1, 0)
            if self.grid.is_cell_empty(pos):
                print("Creating agent {2} at ({0}, {1})"
                      .format(x, y, unique_id))
                a = Walker(unique_id, self, pos, heading)
                self.schedule.add(a)
                self.grid.place_agent(a, pos)
                unique_id += 1

    def step(self):
        self.schedule.step()

class WalkerWorld(Model):
    '''
    Random walker world.
    '''
    height = 10
    width = 10

    def __init__(self, height, width, agent_count):
        '''
        Create a new WalkerWorld.

        Args:
            height, width: World size.
            agent_count: How many agents to create.
        '''
        self.height = height
        self.width = width
        self.grid = MultiGrid(self.height, self.width, torus=True)
        self.agent_count = agent_count

        self.schedule = RandomActivation(self)
        # Create agents
        for i in range(self.agent_count):
            x = random.randrange(self.width)
            y = random.randrange(self.height)
            a = WalkerAgent((x, y), self, True)
            self.schedule.add(a)
            self.grid.place_agent(a, (x, y))

    def step(self):
        self.schedule.step()

class ForestFire(Model):
    '''
    Simple Forest Fire model.
    '''
    def __init__(self, height, width, density):
        '''
        Create a new forest fire model.

        Args:
            height, width: The size of the grid to model
            density: What fraction of grid cells have a tree in them.
        '''
        # Initialize model parameters
        self.height = height
        self.width = width
        self.density = density

        # Set up model objects
        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=False)

        self.datacollector = DataCollector(
            {"Fine": lambda m: self.count_type(m, "Fine"),
             "On Fire": lambda m: self.count_type(m, "On Fire"),
             "Burned Out": lambda m: self.count_type(m, "Burned Out")})

        # Place a tree in each cell with Prob = density
        for (contents, x, y) in self.grid.coord_iter():
            if random.random() < self.density:
                # Create a tree
                new_tree = TreeCell((x, y))
                # Set all trees in the first column on fire.
                if x == 0:
                    new_tree.condition = "On Fire"
                self.grid._place_agent((x, y), new_tree)
                self.schedule.add(new_tree)
        self.running = True

    def step(self):
        '''
        Advance the model by one step.
        '''
        self.schedule.step()
        self.datacollector.collect(self)

        # Halt if no more fire
        if self.count_type(self, "On Fire") == 0:
            self.running = False

    @staticmethod
    def count_type(model, tree_condition):
        '''
        Helper method to count trees in a given condition in a given model.
        '''
        count = 0
        for tree in model.schedule.agents:
            if tree.condition == tree_condition:
                count += 1
        return count

class Charts(Model):

    # grid height
    grid_h = 20
    # grid width
    grid_w = 20

    """init parameters "init_people", "rich_threshold", and "reserve_percent"
       are all UserSettableParameters"""
    def __init__(self, height=grid_h, width=grid_w, init_people=2, rich_threshold=10,
                 reserve_percent=50,):
        self.height = height
        self.width = width
        self.init_people = init_people
        self.schedule = RandomActivation(self)
        self.grid = MultiGrid(self.width, self.height, torus=True)
        # rich_threshold is the amount of savings a person needs to be considered "rich"
        self.rich_threshold = rich_threshold
        self.reserve_percent = reserve_percent
        # see datacollector functions above
        self.datacollector = DataCollector(model_reporters={
                                           "Rich": get_num_rich_agents,
                                           "Poor": get_num_poor_agents,
                                           "Middle Class": get_num_mid_agents,
                                           "Savings": get_total_savings,
                                           "Wallets": get_total_wallets,
                                           "Money": get_total_money,
                                           "Loans": get_total_loans},
                                           agent_reporters={
                                           "Wealth": lambda x: x.wealth})

        # create a single bank for the model
        self.bank = Bank(1, self, self.reserve_percent)

        # create people for the model according to number of people set by user
        for i in range(self.init_people):
            # set x, y coords randomly within the grid
            x = self.random.randrange(self.width)
            y = self.random.randrange(self.height)
            p = Person(i, (x, y), self, True, self.bank, self.rich_threshold)
            # place the Person object on the grid at coordinates (x, y)
            self.grid.place_agent(p, (x, y))
            # add the Person object to the model schedule
            self.schedule.add(p)

        self.running = True
        self.datacollector.collect(self)

    def step(self):
        # tell all the agents in the model to run their step function
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

    def run_model(self):
        for i in range(self.run_time):
            self.step()

class Schelling(Model):
    '''
    Model class for the Schelling segregation model.
    '''

    def __init__(self, height=20, width=20, density=0.8, minority_pc=0.2, homophily=3):
        '''
        '''

        self.height = height
        self.width = width
        self.density = density
        self.minority_pc = minority_pc
        self.homophily = homophily

        self.schedule = RandomActivation(self)
        self.grid = SingleGrid(height, width, torus=True)

        self.happy = 0
        self.datacollector = DataCollector(
            {"happy": "happy"},  # Model-level count of happy agents
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.pos[0], "y": lambda a: a.pos[1]})

        # Set up agents
        # We use a grid iterator that returns
        # the coordinates of a cell as well as
        # its contents. (coord_iter)
        for cell in self.grid.coord_iter():
            x = cell[1]
            y = cell[2]
            if self.random.random() < self.density:
                if self.random.random() < self.minority_pc:
                    agent_type = 1
                else:
                    agent_type = 0

                agent = SchellingAgent((x, y), self, agent_type)
                self.grid.position_agent(agent, (x, y))
                self.schedule.add(agent)

        self.running = True
        self.datacollector.collect(self)

    def step(self):
        '''
        Run one step of the model. If All agents are happy, halt the model.
        '''
        self.happy = 0  # Reset counter of happy agents
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

        if self.happy == self.schedule.get_agent_count():
            self.running = False

class SchellingModel(Model):
    """
    Model class for the Schelling segregation model.
    """

    def __init__(self, height, width, density, minority_pc, homophily):
        """
        """

        self.height = height
        self.width = width
        self.density = density
        self.minority_pc = minority_pc
        self.homophily = homophily

        self.schedule = RandomActivation(self)
        self.grid = SingleGrid(height, width, torus=True)

        self.happy = 0
        self.total_agents = 0
        self.datacollector = DataCollector(
            {"unhappy": lambda m: m.total_agents - m.happy},
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.pos[X], "y": lambda a: a.pos[Y]},
        )

        self.running = True

        # Set up agents
        # We use a grid iterator that returns
        # the coordinates of a cell as well as
        # its contents. (coord_iter)
        for cell, x, y in self.grid.coord_iter():
            if random.random() < self.density:
                if random.random() < self.minority_pc:
                    agent_type = 1
                else:
                    agent_type = 0

                agent = SchellingAgent(self.total_agents, agent_type)
                self.grid.position_agent(agent, x, y)
                self.schedule.add(agent)
                self.total_agents += 1

    def step(self):
        """
        Run one step of the model. If All agents are happy, halt the model.
        """
        self.happy = 0  # Reset counter of happy agents
        self.schedule.step()
        self.datacollector.collect(self)

        if self.happy == self.total_agents:
            self.running = False

class BoidModel(Model):
    '''
    Flocker model class. Handles agent creation, placement and scheduling.
    '''

    def __init__(self,
                 population=100,
                 width=100,
                 height=100,
                 speed=1,
                 vision=10,
                 separation=2,
                 cohere=0.025,
                 separate=0.25,
                 match=0.04):
        '''
        Create a new Flockers model.

        Args:
            population: Number of Boids
            width, height: Size of the space.
            speed: How fast should the Boids move.
            vision: How far around should each Boid look for its neighbors
            separation: What's the minimum distance each Boid will attempt to
                    keep from any other
            cohere, separate, match: factors for the relative importance of
                    the three drives.        '''
        self.population = population
        self.vision = vision
        self.speed = speed
        self.separation = separation
        self.schedule = RandomActivation(self)
        self.space = ContinuousSpace(width, height, True,
                                     grid_width=10, grid_height=10)
        self.factors = dict(cohere=cohere, separate=separate, match=match)
        self.make_agents()
        self.running = True

    def make_agents(self):
        '''
        Create self.population agents, with random positions and starting headings.
        '''
        for i in range(self.population):
            x = random.random() * self.space.x_max
            y = random.random() * self.space.y_max
            pos = np.array((x, y))
            velocity = np.random.random(2) * 2 - 1
            boid = Boid(i, self, pos, self.speed, velocity, self.vision,
                        self.separation, **self.factors)
            self.space.place_agent(boid, pos)
            self.schedule.add(boid)

    def step(self):
        self.schedule.step()

class MoneyModel(Model):
    """A model with some number of agents."""
    def __init__(self, N):
        self.num_agents = N
        self.schedule = RandomActivation(self)
        # Create agents
        for i in range(self.num_agents):
            a = MoneyAgent(i)
            self.schedule.add(a)

    def step(self):
        '''Advance the model by one step.'''
        self.schedule.step()

class InspectionModel(Model):
    '''
    Simple Restaurant Inspection model.
    '''
    def __init__(self, height, width, density):
        '''
        Create a new restaurant inspection model.

        Args:
            height, width: The size of the grid to model
            density: What fraction of grid cells have a restaurant in them.
        '''
        # Initialize model parameters
        self.height = height
        self.width = width
        self.density = density

        # Set up model objects
        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=False)

        self.datacollector = DataCollector(
            {"Good": lambda m: self.count_type(m, "Good"),
             "Bad": lambda m: self.count_type(m, "Bad")})

        # Place a restaurant in each cell with Prob = density
        for (contents, x, y) in self.grid.coord_iter():
            if random.random() < self.density:
                # Create a restaurant
                new_restaurant = RestaurantCell((x, y))
                self.grid._place_agent((x, y), new_restaurant)
                self.schedule.add(new_restaurant)
        self.running = True

    def step(self):
        '''
        Advance the model by one step.
        '''
        self.schedule.step()
        self.datacollector.collect(self)

    @staticmethod
    def count_type(model, restaurant_hygiene):
        '''
        Helper method to count restaurants in a given condition in a given model.
        '''
        count = 0
        for restaurant in model.schedule.agents:
            if restaurant.hygiene == restaurant_hygiene and restaurant.rating != 'Closed':
                count += 1
        return count

class VirusOnNetwork(Model):
    """A virus model with some number of agents"""

    def __init__(self, num_nodes=10, avg_node_degree=3, initial_outbreak_size=1, virus_spread_chance=0.4,
                virus_check_frequency=0.4, recovery_chance=0.3, gain_resistance_chance=0.5):

        self.num_nodes = num_nodes
        prob = avg_node_degree / self.num_nodes
        self.G = nx.erdos_renyi_graph(n=self.num_nodes, p=prob)
        self.grid = NetworkGrid(self.G)
        self.schedule = RandomActivation(self)
        self.initial_outbreak_size = initial_outbreak_size if initial_outbreak_size <= num_nodes else num_nodes
        self.virus_spread_chance = virus_spread_chance
        self.virus_check_frequency = virus_check_frequency
        self.recovery_chance = recovery_chance
        self.gain_resistance_chance = gain_resistance_chance

        self.datacollector = DataCollector({"Infected": number_infected,
                                            "Susceptible": number_susceptible,
                                            "Resistant": number_resistant})

        # Create agents
        for i, node in enumerate(self.G.nodes()):
            a = VirusAgent(i, self, State.SUSCEPTIBLE, self.virus_spread_chance, self.virus_check_frequency,
                           self.recovery_chance, self.gain_resistance_chance)
            self.schedule.add(a)
            # Add the agent to the node
            self.grid.place_agent(a, node)

        # Infect some nodes
        infected_nodes = self.random.sample(self.G.nodes(), self.initial_outbreak_size)
        for a in self.grid.get_cell_list_contents(infected_nodes):
            a.state = State.INFECTED

        self.running = True
        self.datacollector.collect(self)

    def resistant_susceptible_ratio(self):
        try:
            return number_state(self, State.RESISTANT) / number_state(self, State.SUSCEPTIBLE)
        except ZeroDivisionError:
            return math.inf

    def step(self):
        self.schedule.step()
        # collect data
        self.datacollector.collect(self)

    def run_model(self, n):
        for i in range(n):
            self.step()

class SchellingModel(Model):
    '''
    Model class for the Schelling segregation model.
    '''

    def __init__(self, height, width, density, type_pcs=[.2, .2, .2, .2, .2]):
        '''
        '''

        self.height = height
        self.width = width
        self.density = density
        self.type_pcs = type_pcs

        self.schedule = RandomActivation(self)
        self.grid = SingleGrid(height, width, torus=False)

        self.happy = 0
        self.datacollector = DataCollector(
            {"happy": lambda m: m.happy},  # Model-level count of happy agents
            # For testing purposes, agent's individual x and y
            {"x": lambda a: a.pos[0], "y": lambda a: a.pos[1]})

        self.running = True

        # Set up agents
        # We use a grid iterator that returns
        # the coordinates of a cell as well as
        # its contents. (coord_iter)

        total_agents = self.height * self.width * self.density
        agents_by_type = [total_agents*val for val in self.type_pcs]

        for loc, types in enumerate(agents_by_type):
            for i in range(int(types)):
                pos = self.grid.find_empty()
                agent = SchellingAgent(pos, self, loc)
                self.grid.position_agent(agent, pos)
                self.schedule.add(agent)

    def step(self):
        '''
        Run one step of the model. If All agents are happy, halt the model.
        '''
        self.happy = 0  # Reset counter of happy agents
        self.schedule.step()
        self.datacollector.collect(self)

        if self.happy == self.schedule.get_agent_count():
            self.running = False

    def __init__(self, height, width, density):
        '''
        Create a new restaurant inspection model.

        Args:
            height, width: The size of the grid to model
            density: What fraction of grid cells have a restaurant in them.
        '''
        # Initialize model parameters
        self.height = height
        self.width = width
        self.density = density

        # Set up model objects
        self.schedule = RandomActivation(self)
        self.grid = Grid(height, width, torus=False)

        self.datacollector = DataCollector(
            {"Good": lambda m: self.count_type(m, "Good"),
             "Bad": lambda m: self.count_type(m, "Bad")})

        # Place a restaurant in each cell with Prob = density
        for (contents, x, y) in self.grid.coord_iter():
            if random.random() < self.density:
                # Create a restaurant
                new_restaurant = RestaurantCell((x, y))
                self.grid._place_agent((x, y), new_restaurant)
                self.schedule.add(new_restaurant)
        self.running = True

 def __init__(self, N, width=20, height=10):
     self.running = True
     self.N = N    # num of agents
     self.headings = ((1, 0), (0, 1), (-1, 0), (0, -1))  # tuples are fast
     self.grid = SingleGrid(width, height, torus=False)
     self.schedule = RandomActivation(self)
     self.make_walker_agents()

 def __init__(self, height, width, density):
     '''
     Create a new forest fire model.
     
     Args:
         height, width: The size of the grid to model
         density: What fraction of grid cells have a tree in them.
     '''
     # Initialize model parameters
     self.height = height
     self.width = width
     self.density = density
     
     # Set up model objects
     self.schedule = RandomActivation(self)
     self.grid = Grid(height, width, torus=False)
     
     # Place a tree in each cell with Prob = density
     for x in range(self.width):
         for y in range(self.height):
             if random.random() < self.density:
                 # Create a tree
                 new_tree = TreeCell(x, y)
                 # Set all trees in the first column on fire.
                 if x == 0:
                     new_tree.condition = "On Fire"
                 self.grid[y][x] = new_tree
                 self.schedule.add(new_tree)
     self.running = True