# -*- coding: utf-8 -*-
"""
Created on Tue Oct 11 2016

@author: pcutter

Modified by: (Your name(s) here)
With Assistance from:  (Anyone who helped)

Population.py
This module provides a computational model to examining a specific system that
is discrete time, discrete space, and model of predator_prey population dynamics.

This model will make the following simplifying assumptions.
1. The environment is divided in to a finite grid of discrete cells that can
contain at most one individual.
2. Individuals never move from one cell to another.
3. Individuals can only interact with their immediate neighbors.
4. Interactions occur one at a time at randomly selected positions.

In the simulation, the populations of preys and predators of each generation will
be counted to plot a dynamic changes.
"""

import numpy as np
import matplotlib.pyplot as plt
import random

"""
    Display the environment as an image with three colors.
    Purple represents empty cells, green/teal is prey, yellow is predator
    PARAM:
        env: the current environment
   """
def display_env(env):
    plt.imshow(env)
    plt.show()
    plt.pause(0.5)
    print env   # this can be commented out when you know your code is working

"""
    Randomly select a row and column in a 2d array.
    PARAMS:
    env: the two dimensional array.
    RETURNS: 
        (r, c): randomly selected row and column.
    """
def rand_cell(env):
    """
    Randomly select a row and column in a 2d array.

    :param env: the two dimensional array.
    :return: (r, c) randomly selected row and column.
    """
    row, col = env.shape
    r = random.randint(0, row - 1)
    c = random.randint(0, col - 1)
    return r, c

"""
    Randomly select one of the four neighbors of a cell.
    PARAMSL
        env: the two dimensional array.
        row: the row index of the cell
        col: the column index of the cell
   RETURNS:
        (row, col): row and column of the randomly selected neighbor.
    """    
def rand_neighbor(env, r, c):
    
    rows, cols = env.shape
    nei = [(-1,0),(1,0),(0,1),(0,-1)]    
    i = nei[random.randint(0, 3)]
    n_r = (r + i[0])%rows
    n_c = (c + i[1])%cols

    return n_r, n_c

"""
    Randomly initialize a 2d array to represent predator and prey
    distribution, 1's will represent prey, 2's will represent predators,
    and 0's will represent empty space
    PARAMS:
        width: width of the environment(2D array).
        height: height of the environment(2D array).
        prop_prey: beginning proportions of prey.
        prop_pred: beginning proportions of predator.
    RETURNS:
        env - the initialized environment(2D array).
    """
def initial_env(width, height, prop_prey, prop_pred):
   
    if (prop_pred + prop_prey) > 1:
        raise ValueError('Percentage prey plus percentage predators for initial ' +
                         'environment is greater than 100%!')
    else:
        arr = np.random.choice([0,1,2],size=(width,height),p=[1-(prop_prey+prop_pred),prop_prey,prop_pred])
        return arr

"""
    Give a randomly selected cell a chance to reproduce if it contains prey
    PARAMS:
        env: A two dimensional array representing the environment
        r: The probability that prey will reproduce if a randomly selected 
            neighboring cell is empty
    """
def reproduction(env,r):
    env = env # DELETE THIS LINE ONCE YOU IMPLEMENT THIS FUNCTION!
    # STEPS TO COMPLETE:
    # RANDOMLY SELECT A CELL (USE rand_cell function)
    # RANDOMLY SELECT ONE OF THE 4 NEIGHBORS OF THAT CELL (USE rand_neighbor function)
    
    # IF THE RANDOMLY SELECTED CELL CONTAINS PREY:
    #   IF THE RANDOMLY SELECTED NEIGHBOR IS EMPTY:
    #       POPULATE THE EMPTY CELL WITH PREY WITH PROBABILITY r
    
    # HINT: the statement random.random() < r will be true with probability r
 
"""
    This function contains the main loop of a simulation of predator prey
    interactions on a two dimensional grid.
    PARAMS: 
        n_gens: The number of generations to run the simulation
        r: Probability of reproduction.
        p: Probability of predation.
        s: Probability of starvation.
        prop_prey: beginning proportions of prey.
        prop_pred: beginning proportions of predator.
        env_rows: height of the environment(2D array).
        env_cols: width of the environment(2D array).
 """           
def predator_prey(n_gens,r,p,d, prop_prey, prop_pred, env_rows, env_cols):
    # Initialize the environment with prey and predator.
    env = initial_env(env_rows, env_cols, prop_prey, prop_pred)

    # Build lists to record the prey and predator populations.
    # ADD CODE HERE TO CREATE TWO LISTS

    # Calculate the initial populations and append to the lists you just created.
    #ADD CODE HERE TO COUNT # OF PREY, # OF PREDATORS (USE A FUNCTION IF POSSIBLE)
    
    # Simulation for n generations.
    for gen in range(n_gens):
        for i in range((env_rows * env_cols)//3):
            reproduction(env,r)
            #THERE SHOULD BE TWO ADDITIONAL FUNCTION CALLS HERE
            # REPRESENTING PREDATION AND STARVATION
        # ADD FUNCTION CALLS HERE TO GET THE COUNTS OF DIFFERENT SPECIES
        # ADD STATEMENTS TO APPEND THESE COUNTS TO THE APPROPRIATE LISTS
        display_env(env)
        print   # Put a blank line between environments
    # CONSIDER PRINTING YOUR LISTS OF POPULATION VALUES FOR EACH SPECIES HERE