# exec(open(".//09_Index_Arrays.py"))

import numpy as np
print "Executing: 09_Index_Arrays.py"

'''
NumPy Index Arrays
'''

# Change False to True for each block of code to see what it does

# Using index arrays
if False:
    a = np.array([1, 2, 3, 4])
    b = np.array([True, True, False, False])
    
    print a[b]
    print a[np.array([True, False, True, False])]

# Creating the index array using vectorized operations
if False:
    a = np.array([1, 2, 3, 2, 1])
    b = (a >= 2)
    
    print a[b]
    print a[a >= 2]
   
    
# Creating the index array using vectorized operations on another array
if False:
    a = np.array([1, 2, 3, 4, 5])
    b = np.array([1, 2, 3, 2, 1])
    
    print b == 2
    print a[b == 2]

# Time spent in the classroom in the first week for 20 students
time_spent = np.array([
       12.89697233,    0.        ,   64.55043217,    0.        ,
       24.2315615 ,   39.991625  ,    0.        ,    0.        ,
      147.20683783,    0.        ,    0.        ,    0.        ,
       45.18261617,  157.60454283,  133.2434615 ,   52.85000767,
        0.        ,   54.9204785 ,   26.78142417,    0.
])

# Days to cancel for 20 students
days_to_cancel = np.array([
      4,   5,  37,   3,  12,   4,  35,  38,   5,  37,   3,   3,  68,
     38,  98,   2, 249,   2, 127,  35
])


# This is my version which is also correct

# def mean_time_for_paid_students(time_spent, days_to_cancel):
#    '''
#    Fill in this function to calculate the mean time spent in the classroom
#    for students who stayed enrolled at least (greater than or equal to) 7 days.
#    Unlike in Lesson 1, you can assume that days_to_cancel will contain only
#    integers (there are no students who have not canceled yet).
    
#    The arguments are NumPy arrays. time_spent contains the amount of time spent
#    in the classroom for each student, and days_to_cancel contains the number
#    of days until each student cancel. The data is given in the same order
#    in both arrays.
#    '''
#    return time_spent[days_to_cancel >= 7].mean()

# This is udacity's version

def mean_time_for_paid_students(time_spent, days_to_cancel):
    # create the index array
    is_paid = days_to_cancel >= 7
    paid_time = time_spent[is_paid]
    return paid_time.mean()

