#Some useful functions for more advanced analytics

#Joins in spark take RDD[K,V] x RDD[K,U] => RDD[K, [U,V] ]
#This function returns U
def join_lhs(t):
    return t[1][0]

#Joins in spark take RDD[K,V] x RDD[K,U] => RDD[K, [U,V] ]
#This function returns V
def join_rhs(t):
    return t[1][1]

#Add a key/value to a dictionary and return the dictionary
def annotate_dict(d, k, v):
    d[k] = v
    return d

#Plots a density plot of a set of points representing inliers and outliers
#A rugplot is used to indicate the points and the outliers are marked in red.
def plot_outliers(inliers, outliers, reason):
    fig, ax = plt.subplots(nrows=1)
    sns.distplot(inliers + outliers, ax=ax, rug=True, hist=False)
    ax.plot(outliers, np.zeros_like(outliers), 'ro', clip_on=False)
    fig.suptitle('Distribution for {} Values'.format(reason), size=14)

#Outlier analysis using Median Absolute Deviation

#Using reservoir sampling, uniformly sample N points
#requires O(N) memory
def sample_points(points, N):
    sample = [];
    for i,point in enumerate(points):
        if i < N:
            sample.append(point)
        elif i >= N and random.random() < N/float(i+1):
            replace = random.randint(0,len(sample)-1)
            sample[replace] = point
    return sample

#Returns a function which will extract the median at location 'key'
#a list of dictionaries.
def median_func(key):
    #Right now it uses numpy's median, but probably a quickselect
    #implementation is called for as I expect this doesn't scale
    return lambda partition_value : ( partition_value[0]\
                                    , np.median( 
                                        [ d[key] \
                                         for d in partition_value[1]\
                                        ]\
                                               )\
                                    )

#Compute the modified z-score for use by  as per Iglewicz and Hoaglin:
#Boris Iglewicz and David Hoaglin (1993),
#"Volume 16: How to Detect and Handle Outliers",
#The ASQC Basic References in Quality Control: Statistical Techniques
#, Edward F. Mykytka, Ph.D., Editor.
def get_z_score(reason_to_diff):
    med = join_rhs(reason_to_diff)
    if med > 0:
        return 0.6745 * join_lhs(reason_to_diff)['diff'] / med
    else:
        return 0

def is_outlier(thresh):
    return lambda reason_to_diff : get_mad(reason_to_diff) > thresh

#Return a RDD of a uniform random sample of a specified size per key
def get_inliers(reason_amount_pairs, size=2000):
    group_by_reason = reason_amount_pairs.groupByKey()
    return group_by_reason.map(lambda t : (t[0], sample_points(t[1], size)))


#Return the outliers based on Median Absolute Deviation
#See http://www.itl.nist.gov/div898/handbook/eda/section3/eda35h.htm 
#for more info.
#The input key structure is reason_specialty => dict(amount
#                                                   , physician
#                                                   , payer
#                                                   , specialty
#                                                   )
def get_outliers(reason_amount_pairs, thresh=3.5):
    """
        This uses the median absolute deviation (MAD) statistic to find
        outliers for each reason x specialty partitions.
        
        Outliers are computed as follows: 
        * Let X be all the payments for a given specialty, reason pair
        * Let x_i be a payment in X
        * Let MAD be the median absolute deviation, defined as
          MAD = median( for all x in X, | x - median(X)| )
        * Let M_i be the modified z-score for payment x_i, defined as
          0.6745*(x_i − median(X) )/MAD
        
        As per the recommendations by Iglewicz and Hoaglin, a payment is
        considered an outlier if the modified z-score, M_i > thresh, which
        is 3.5 by default.
        
        
        REFERENCE:
        Boris Iglewicz and David Hoaglin (1993),
        "Volume 16: How to Detect and Handle Outliers",
        The ASQC Basic References in Quality Control: Statistical Techniques,
        Edward F. Mykytka, Ph.D., Editor.
    """
    
    group_by_reason = reason_amount_pairs.groupByKey()
   
    #Filter by only reason/specialty's with more than 1k entries
    #and compute the median of the amounts across the partition.
    
    #NOTE: There may be some scalability challenges around median,
    #so some care should be taken to reimplement this if partitioning
    #by (reason, specialty) does not yield small enough numbers to 
    #handle in an individual map function.
    reason_to_median = group_by_reason.filter(lambda t: len(t[1]) > 1000) \
                                      .map(median_func('amount'))
   
    #Join the base, non-grouped data, with the median per key,
    #consider just the payments more than the median
    #since we're looking for large money outliers and annotate 
    #the dictionary for each entry x_i with the following:
    # * diff = |x_i - median(X)| in the parlance of the comment above.
    #   NOTE: Strictly speaking I can drop the absolute value since 
    #         x_i > median(X), but I choose not to.
    # * median = median(X)
    # 
    reason_abs_dist_from_median = \
        reason_amount_pairs.join(reason_to_median) \
                           .filter(lambda t : \
                              join_lhs(t)['amount'] > join_rhs(t)\
                                  ) \
                           .map(lambda t :\
                                 ( t[0]\
                                 , dict( diff=abs(\
                                      join_lhs(t)['amount'] - join_rhs(t)\
                                                 )\
                                       , row=annotate_dict(join_lhs(t) \
                                                          , 'median' \
                                                          , join_rhs(t) \
                                                          )\
                                       )\
                                 )\
                               )
                
    # Given diff cached per element, we need only compute the median 
    # of the diffs to compute the MAD.  
    #Remember, MAD = median( for all x in X, | x - median(X)| )
    reason_to_MAD = reason_abs_dist_from_median.groupByKey() \
                                               .map(median_func('diff'))
    reason_to_MAD.take(1)
    
    # Joining the grouped data to get both | x_i - median(X) | 
    # and MAD in the same place, we can compute the modified z-score
    # , 0.6475*| x_i - median(X)| / MAD, and filter by the ones which 
    # are more than threshold we can then do some pivoting of keys and 
    #sort by that threshold to give us the ranked list of outliers.
    return reason_abs_dist_from_median\
              .join(reason_to_MAD) \
              .filter(is_outlier(thresh))\
              .map(lambda t: (get_z_score(t)\
                             , annotate_dict(join_lhs(t)['row']\
                                            , 'key'\
                                            , t[0]\
                                            )\
                             )\
                  ) \
              .sortByKey(False) \
              .map(lambda t: ( t[1]['key']\
                             , annotate_dict( t[1]\
                                            , 'mad'\
                                            , t[0]\
                                            )\
                             )\
                  )

#Filter the outliers by reason and return a RDD with just the outliers 
#of a specified reason.
def get_by_reason(outliers, reason):
    return outliers.filter(lambda t: str.startswith(t[0].encode('ascii'\
                                                               , 'ignore'\
                                                               )\
                                                   ,reason\
                                                   )\
                          )

#Grab data using Spark-SQL and filter with spark core RDD operations 
#to only yield the data we want, ones with physicians, payers and reasons
reason_amount_pairs = \
    sqlContext.sql("""select reason
                           , physician_specialty
                           , amount
                           , physician_id
                           , payer 
                      from payments"""\
                  )\
              .filter(lambda row:len(row.reason) > 3 \
                              and len(row.physician_id) > 3\
                              and len(row.payer) > 3\
                     ) \
              .map(lambda row: ( "{}_{}".format( row.reason\
                                               , row.physician_specialty\
                                               )\
                             , dict(amount=row.amount\
                                   ,physician_id=row.physician_id\
                                   ,payer=row.payer\
                                   ,specialty=row.physician_specialty\
                                   )\
                             )\
                   )

#Get the outliers based on a modified z-score threshold of 3.5
outliers = get_outliers(reason_amount_pairs, 3.5)
#Get a sample per specialty/reason partition
inliers = get_inliers(reason_amount_pairs)

#display the top k outliers in a table and a distribution plot
#of an inlier sample along with the outliers rug-plotted in red
def display_info(inliers_raw, outliers_raw_tmp, reason, k=None):
    outliers_raw = []
    if k is None:
        outliers_raw = sorted(outliers_raw_tmp\
                             , key=lambda d:d[1]['amount']\
                             , reverse=True\
                             )
    else:
        outliers_raw = sorted(outliers_raw_tmp\
                             , key=lambda d:d[1]['amount']\
                             , reverse=True\
                             )[0:k]
    inlier_pts = []
    for i in [d[1] for d in inliers_raw]:
        for j in i:
            inlier_pts.append(j['amount'])
    outlier_pts= [d[1]['amount'] for d in outliers_raw]
    plot_outliers(inlier_pts[0:1500], outlier_pts, reason)

    print_table(['Physician','Specialty', 'Payer', 'Amount']\
               , [d[1]['physician_id'] for d in outliers_raw]\
               , [ [ d[1]['specialty'] \
                   , d[1]['payer'].encode('ascii', 'ignore') \
                   , '$' + locale.format('%0.2f', d[1]['amount'], grouping=True)\
                   ] \
                 for d in outliers_raw]\
               )

#outliers for food and beverage purchases
food_outliers = get_by_reason(outliers, 'Food and Beverage').collect()
food_inliers = get_by_reason(inliers, 'Food and Beverage').collect()
display_info(food_inliers, food_outliers, 'Food and Beverage', 4)

travel_outliers = get_by_reason(outliers, 'Travel and Lodging').collect()
travel_inliers = get_by_reason(inliers, 'Travel and Lodging').collect()
display_info(travel_inliers, travel_outliers, 'Travel and Lodging', 10)

consulting_outliers = get_by_reason(outliers, 'Consulting Fee').collect()
consulting_inliers = get_by_reason(inliers, 'Consulting Fee').collect()
display_info(consulting_inliers, consulting_outliers, 'Consulting Fee', 10)

gift_outliers = get_by_reason(outliers, 'Gift').collect()
gift_inliers = get_by_reason(inliers, 'Gift').collect()
display_info(gift_inliers, gift_outliers, 'Gift', 10)