#######################
# Gregory S. Warrington
# 06/05/2018

# all of the metrics we're considering in this paper
# note that Lop and EVW include corresponding statistical data
# see populate_metrics function for how this is recorded
# metric_dict = {'Dec': get_declin,\
#                'BDec':get_bdec,\
#                'EG':get_EG,\
#                'LG':get_EG_loss_only,\
#                'DG':get_EG_difference,\
#                'SG':get_EG_surplus_only,\
#                'VC1':get_EG_vote_centric_one,\
#                'VC2':get_EG_vote_centric_two,\
#                'Bias':get_bias,\
#                'MM':get_mean_median,\
#                'Lop':get_lop_list,\
#                'EVW':get_evw}

#########################################################################
# summary of metrics defined
# 
#
#########################################################################
# get_EG_loss_only = lambda=0; 1-proportionality
# get_EG_original = lambda=1; 2-proportionality
# get_EG_difference = lambda=2; 3-proportionality
# get_EG_surplus_only = Cho's suggestion
#
# get_EG_vote_centric_one - Nagle's version of EG
# get_EG_vote_centric_two - Nagle's version of EG with lambda=2 - Tapp's preferred version
#
# get_dec = declination
# get_bias = partisan bias
# get_mean_median = mean-median
# get_lopsided_means_sig
# get_equal_vote_weight

import math
import numpy as np
import scipy.stats as stats

#########################################################################
# code to compute tau-gap, EG and declination from vote distribution data
#########################################################################

def get_tau_gap(vals,tau):
    """ compute tau-gap. 

    Note that tau-gap when tau=0 is twice the EG
    TODO: Review code for when tau < 0.
    """  
    ans = 0.0
    m = 0.0
    N = len(vals)
    if N == 0:
        return float('nan')
    for i in range(N):
        ai = 2.0*vals[i]-1
        if ai > 0:
            m += 1
        tmp = 0.0
        if tau >= 0: # votes close to 50% are weighed more ("traditional")
            if ai >= 0:  
                tmp = pow(ai,tau+1)
            else:
                tmp = -pow(-ai,tau+1)
        else:             # votes close to 50% are weighed less 
            if ai >= 0:
                tmp = pow((1-ai),-tau+1)
            else:
                tmp = -pow((1+ai),-tau+1)
        ans += tmp
                    
    if tau >= 0:
        return 2.0*(ans/N + 0.5 - m/N)
    else:
        return -2.0*(ans/N + 0.5 - m/N)

def get_EG(vals):
    """ return the efficiency gap
    """
    return get_tau_gap(vals,0)/2

#################################################################
# variants on the efficiency gap obtained by varying definition
# and/or relative weight of wasted votes.
#################################################################
def _EG_lam(vals,lam):
    """ weight excess votes by lambda - lambda=1 is usual EG
    """
    N = len(vals)
    if N == 0:
        return float('nan')
    ans = 0.0
    dem_surp = 0.0
    rep_surp = 0.0
    dem_loss = 0.0
    rep_loss = 0.0
    for i in range(N):
        # dems lose
        if vals[i] <= 0.5:
            dem_loss += vals[i]
            rep_surp += lam*(0.5 - vals[i])
        else:
            dem_surp += lam*(vals[i] - 0.5)
            rep_loss += (1-vals[i])
    return ((dem_surp + dem_loss) - (rep_surp + rep_loss))/N

def get_EG_loss_only(vals):
    """ weight excess votes by lambda=0
        only counts losing votes as wasted
        corresponds to 1-proportionality
    """
    return _EG_lam(vals,0)

def get_EG_original(vals):
    """ weight excess votes by lam=1
        this is usual EG - should return same answer as get_EG
        corresponds to 2-proportionality
    """
    return _EG_lam(vals,1)

def get_EG_difference(vals):
    """ weight excess votes by lambda=2
        corresponds to 3-proportionality
        this is what Griesbach suggests
        Nagle lambda=2?
    """
    return _EG_lam(vals,2)

def get_EG_surplus_only(vals):
    """ only looks at surplus as wasted
        - see Table 2, "Winning Efficiency" column in Cho's UPenn essay
    """
    if len(vals) == 0:
        return float('nan')

    dem_surp = 0.0
    rep_surp = 0.0
    for i in range(len(vals)):
        # dems lose
        if vals[i] <= 0.5:
            rep_surp += (0.5 - vals[i])
        else:
            dem_surp += (vals[i] - 0.5)
    # negate given correlations
    return (dem_surp - rep_surp)/(len(vals))

#################################################################
# versions of nagle's vote-centric Efficiency Gap
#################################################################
def get_EG_vote_centric_lam(vals,lam):
    """ compare shares of wasted votes a la Nagle
    """
    ans = 0.0
    dem_surp = 0.0
    rep_surp = 0.0
    dem_loss = 0.0
    rep_loss = 0.0
    dem_tot = 0.0
    rep_tot = 0.0
    for i in range(len(vals)):
        # dems lose
        if vals[i] <= 0.5:
            dem_loss += vals[i]
            rep_surp += lam*(0.5 - vals[i])
            rep_tot += (1-vals[i])
            dem_tot += vals[i]
        else:
            dem_surp += lam*(vals[i]-0.5)
            rep_loss += (1-vals[i])
            dem_tot += vals[i]
            rep_tot += (1-vals[i])
    if dem_tot == 0 or rep_tot == 0:
        return float('nan')
    return (dem_surp + dem_loss)/dem_tot - (rep_surp + rep_loss)/rep_tot

def get_EG_vote_centric_one(vals):
    """ Nagle's vote-centric version with lam=1
        This is the closest to the original EG
    """
    return get_EG_vote_centric_lam(vals,1)

def get_EG_vote_centric_two(vals):
    """ Nagle's vote-centric version with lam=2
        This is the vote-centric version of Griesbach's suggestion
    """
    return get_EG_vote_centric_lam(vals,2)

###########################################################################

def get_declination(st,vals):
    """ Get declination.

    Non-pandas version
    Expressed as a fraction of 90 degrees
    """
    bel = sorted(filter(lambda x: x <=  0.5, vals))
    abo = sorted(filter(lambda x: x > 0.5, vals))
    if len(bel) < 1 or len(abo) < 1:
        return -2.0

    theta = np.arctan((1-2*np.mean(bel))*len(vals)/len(bel))
    gamma = np.arctan((2*np.mean(abo)-1)*len(vals)/len(abo))

    return 2.0*(gamma-theta)/3.1415926535 # Enough precision for you?

def get_declin(vals):
    """ Get declination.
    Modified for pandas

    Expressed as a fraction of 90 degrees
    """
    bel = sorted(filter(lambda x: x <=  0.5, vals))
    abo = sorted(filter(lambda x: x > 0.5, vals))
    if len(bel) < 1 or len(abo) < 1:
        return float('nan')

    theta = np.arctan((1-2*np.mean(bel))*len(vals)/len(bel))
    gamma = np.arctan((2*np.mean(abo)-1)*len(vals)/len(abo))

    return 2.0*(gamma-theta)/3.1415926535 # Enough precision for you?

def get_declin_tilde(vals):
    """ compatibility routine for getting declination; for pandas
    """
    declin = get_declin(vals)
    if math.isnan(declin):
        return float('nan')
    else:
        return declin*math.log(len(vals))/2

def get_dec(vals):
    """ compatibility routine for getting declination
    """
    return get_declination('',vals)

def get_dec_tilde(vals):
    """ compatibility routine for getting declination
    """
    return get_declination('',vals)*math.log(len(vals))/2

def get_bdec(vals):
    """ Get declination but with extra districts added in to mute response
    Modified for pandas
    Expressed as a fraction of 90 degrees
    """
    # This is obviously somewhat arbitrary for large elections
    xtra_num = 1 + int(math.ceil(len(vals)/20))
    # add in two points so each slope is decreased slightly.
    # only significant effect should be when one side dominates.
    bel = sorted(filter(lambda x: x <=  0.5, vals) + [0.5]*xtra_num) 
    abo = sorted(filter(lambda x: x > 0.5, vals) + [0.5001]*xtra_num)
    # need at least one real district
    if len(bel) <= xtra_num or len(abo) <= xtra_num:
        return float('nan')

    theta = np.arctan((1-2*np.mean(bel))*(len(vals)+2*xtra_num)/len(bel))
    gamma = np.arctan((2*np.mean(abo)-1)*(len(vals)+2*xtra_num)/len(abo))

    # print theta,gamma,2.0*(gamma-theta)/3.1415926535
    return 2.0*(gamma-theta)/3.1415926535 # Enough precision for you?

def get_bdec_tilde(vals):
    """ for pandas
    """
    bdec = get_bdec(vals)
    if math.isnan(bdec):
        return float('nan')
    else:
        return bdec*math.log(len(vals))/2

def get_bias(vals):
    """ Compute partisan bias. 02.15.18
    """
    if len(vals) == 0:
        return float('nan')
    # uniformly shift election to be tied statewide; make sure between 0 and 1
    nvals = map(lambda x: min(max(x + (0.5-np.mean(vals)),0),1), vals) 
    # find fraction of seats won by democrats in this tied election
    demfrac = len(filter(lambda x: x > 0.5, nvals))*1.0/len(nvals)
    return 0.5 - demfrac # negated so positive for republicans for consistency

def get_mean_median(vals):
    """ Compute the mean-median difference. 02.15.18
    """
    return np.mean(vals)-np.median(vals)

def get_t_test(vals):
    """ Compute the t-test 
    """
    bel = sorted(filter(lambda x: x <=  0.5, vals))
    abo = sorted(filter(lambda x: x > 0.5, vals))
    dem_arr = map(lambda x: x, abo)
    rep_arr = map(lambda x: 1 - x, bel)
    dem_margin = np.mean(dem_arr)
    rep_margin = np.mean(rep_arr)
    # print "diff: %.2f dem  %.2f rep %.2f %s" % (dem_margin-rep_margin,dem_margin,rep_margin,stats.ttest_ind(dem_arr,rep_arr))
    return stats.ttest_ind(dem_arr,rep_arr)

def get_lop_list(vals,verbose=False):
    """ Compute the t-test for lopsided means and decide if significant
        Return data as a tuple to be parsed out
    """
    bel = sorted(filter(lambda x: x <=  0.5, vals))
    abo = sorted(filter(lambda x: x > 0.5, vals))
    # need at least two for each
    if len(bel) < 2 or len(abo) < 2:
        return [float('nan'),float('nan'),float('nan'),False]

    dem_arr = map(lambda x: x - 0.5, abo)
    rep_arr = map(lambda x: 0.5 - x, bel)
    dem_margin = np.mean(dem_arr)
    rep_margin = np.mean(rep_arr)
    ttest = stats.ttest_ind(dem_arr,rep_arr)
    return [dem_margin-rep_margin,ttest[0],ttest[1],ttest < 0.05]

def get_lopsided_means_sig(vals,verbose=False):
    """ Compute the t-test for lopsided means and decide if significant
    """
    bel = sorted(filter(lambda x: x <=  0.5, vals))
    abo = sorted(filter(lambda x: x > 0.5, vals))
    dem_arr = map(lambda x: x - 0.5, abo)
    rep_arr = map(lambda x: 0.5 - x, bel)
    dem_margin = np.mean(dem_arr)
    rep_margin = np.mean(rep_arr)
    ttest = stats.ttest_ind(dem_arr,rep_arr)
    print "diff: %.2f dem  %.2f rep %.2f %s" % (dem_margin-rep_margin,dem_margin,rep_margin,stats.ttest_ind(dem_arr,rep_arr))
    if ttest[1] < 0.05:
        if dem_margin > rep_margin:
            return 1 # negated - larger margin for dems is greater republican advantage
        else:
            return -1
    else:
        return 0

def get_equal_vote_weight(vals):
    """ Compute the Best et al. equal vote weight statistic
    essentially mean-median but only positive if majority rule is violated
    """
    demvotes = np.mean(vals)
    demseats = len(filter(lambda x: x > 0.5, vals))*1.0/len(vals)
    if ((demvotes > 0.5 and demseats < 0.5) or (demvotes <= 0.5 and demseats > 0.5)):
        return np.mean(vals)-np.median(vals)
    else:
        return 0

def get_evw(vals):
    """ Compute the Best et al. equal vote weight statistic
    essentially mean-median but only positive if majority rule is violated
    modified to work with pandas
    """
    if len(vals) == 0:
        return [float('nan'),False]
    demvotes = np.mean(vals)
    demseats = len(filter(lambda x: x > 0.5, vals))*1.0/len(vals)
    return [np.mean(vals)-np.median(vals),(demvotes > 0.5 and demseats < 0.5) or \
            (demvotes <= 0.5 and demseats > 0.5)]