# define the font styles
title_font = fm.FontProperties(family='serif', style='normal', size=19, weight='normal', stretch='normal')
label_font = fm.FontProperties(family='serif', style='normal', size=16, weight='normal', stretch='normal')
ticks_font = fm.FontProperties(family='serif', style='normal', size=10, weight='normal', stretch='normal')
annotation_font = fm.FontProperties(family='serif', style='normal', size=10, weight='normal', stretch='normal')
axis_bgcolor = 'white'
# make a re-usable function to draw nice bar charts
def plot_hist(data, title='', xlabel='Test Statistic Values', ylabel='Frequency', color='red'):
    
    if type(data) == list:
        data = pd.Series(data)
    ax = data.hist(figsize=[9, 6], width=0.35, alpha=0.5, 
                   color=color, edgecolor='k', grid=False, rwidth=.9)

    ax.set_xticklabels(data.index, rotation=45, rotation_mode='anchor', ha='right')
    ax.yaxis.grid(True)
    #for label in ax.get_yticklabels():
    #    label.set_fontproperties(ticks_font)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.grid(False)
    ax.set_facecolor(axis_bgcolor)   
    ax.set_title(title, fontproperties=title_font)
    ax.set_xlabel(xlabel, fontproperties=label_font)
    ax.set_ylabel(ylabel, fontproperties=label_font)

    plt.show()

ax2 = fig.add_subplot(2,1,2)
    ax2.hist(dem_electoral, 22, alpha=0.7, label='Democrat',rwidth=.9, color='#0099ff')
    ax2.axvline(x=269, c='k', linewidth=7)
    ax2.set_title("Histogram of Electoral Votes: Democrats", fontsize=20, **hfont)
    ax2.grid(False)
    ax2.set_xlabel('Number of Electoral Votes', fontsize=15)
    ax2.set_ylabel('Frequency', fontsize=14)
    ax2.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.set_xlim(200,335)


"""
Explanation: utils functions
End of explanation
"""


def diff_means3(data, response, treatment):
    trt1 = data[treatment].cat.categories[0]
    trt2 = data[treatment].cat.categories[1]
    mean_diff = data[response][data[treatment]==trt1].mean() - data[response][data[treatment]==trt2].mean()
    return mean_diff

"""
Explanation: Now our data is formatted in the format we want.
We just need to write pd.category specific test statistics functions.
That said, this is a big limitation. We should see if we can write it in other ways. It would be better to have test statistic functions that just use 'group_a and group_b as inputs.
The above can be done. I should re-write format so that it includes
End of explanation
"""


class PermTest(object):
    
    
    def __init__(self, data):
        """"""
        self.data = data
        self.formatted_data = None
        
        
        
    def format_data(self, response1=None, response2=None):
        """
        Converts data from a 2-column format (where each column refers to a treatment type) to
        a 2-column format where response is one column, treatment is the other.

        This is done for both efficiency of storage and test runtime.

        :param data: pandas DataFrame. each column must be of same length
        :return: pandas DataFrame in response, treatment format
        """
        # create new treatment Series
        
        if (response1==None) and (response2==None):
            # No columns specified, so default to first 2
            
            nrow, ncol = self.data.iloc[:,[0,1]].shape
            treatment_column = pd.Series([self.data.iloc[:,[0,1]].columns.values[i]
                            for i in range(ncol)
                            for j in range(nrow)], dtype='category')
            response_column = pd.concat([self.data.iloc[:,0],
                                         self.data.iloc[:,1]]).tolist()
        else:
            
            nrow, ncol = self.data.loc[:,[response1,response2]].shape
            treatment_column = pd.Series([self.data.loc[:,[response1,response2]].columns.values[i]
                            for i in range(ncol)
                            for j in range(nrow)], dtype='category')
            response_column = pd.concat([self.data.loc[:,response1],
                                         self.data.loc[:,response2]]).tolist()

        self.formatted_data = pd.DataFrame({"response": response_column,
                                      "treatment": treatment_column})

        
                
            
    def diff_means(self, response="response", treatment="treatment", data=None):
        """"""
        if data is None:
            trt1 = self.formatted_data[treatment].cat.categories[0]
            trt2 = self.formatted_data[treatment].cat.categories[1]

            mean_diff = self.formatted_data[response][self.formatted_data[treatment]==trt1].mean() \
            - self.formatted_data[response][self.formatted_data[treatment]==trt2].mean()

            return mean_diff
        else:
            trt1 = data[treatment].cat.categories[0]
            trt2 = data[treatment].cat.categories[1]

            mean_diff = data[response][data[treatment]==trt1].mean() \
            - data[response][data[treatment]==trt2].mean()

            return mean_diff
    
    
    def simPermDsn(self, response="response", treatment="treatment", test="mean", k=100):
        """
        Creates permutation dsn for our data
        :param data:
        :param hasTrt:
        :param testStat:
        :param k:
        :return:
        """
        df = self.formatted_data[:]
        testStatistics = []
        if test == "mean":
            for simulation in range(k):
                df[treatment] = pd.Series(np.random.permutation(self.formatted_data[treatment]),
                                            dtype='category')
                testStatistics.append(self.diff_means(response=response, treatment=treatment, data=df))

            return testStatistics
        
    def pvalue(self, response="response", treatment="treatment", test="mean", k=100):
        """Returns p-value for test"""
        if test == "mean":
            current_stat = self.diff_means(response=response, treatment=treatment)

            simulated_stats = self.simPermDsn(response=response, treatment=treatment, test=test, k=k)
            simulated_stats = np.array(simulated_stats)
            p_value = 1.0 * sum(abs(simulated_stats) >= abs(current_stat))/k

            return p_value
        
    def plot_dsn(self, response="response", treatment="treatment", test="mean", k=100):
        """Plots distribution of test statistics"""
        simulated_stats = self.simPermDsn(response=response, treatment=treatment, test=test, k=k)
        tit = "Permutated Test Statistic Distribution for k={} Simulations".format(k)
        plot_hist(simulated_stats, title=tit, xlabel="Test Statistic Value", ylabel="Frequency")

    

"""
Explanation: Permutation Test
To do:
    Format data initialization so that the formatted data is automatic.
    This will require changing default arguments in format_data_from_two_response
End of explanation
"""


data = pd.DataFrame({"treatment1":[1.5, 2, 5, 0.5, 4, 6.6, 4.5, 3, 4, 2.5],
                     "treatment2":[1, 5, 4, 3, 5, 6, 4, 3, 4.5, 3.2]})
data

"""
Explanation: Create data
End of explanation
"""


pt = PermTest(data)
pt.format_data()
pt.formatted_data

"""
Explanation: Create PermTest instance and formate data
End of explanation
"""


print(pt.diff_means('response', 'treatment', data=None))

"""
Explanation: Calculate difference of means
End of explanation
"""


testsimdata = pt.simPermDsn(k=1000)

"""
Explanation: Create simulated distribution of test statistics (differences of means)
End of explanation
"""


pt.plot_dsn(k=1000)

"""
Explanation: Visualize simulated distribution
End of explanation
"""


pt.test()

"""
Explanation: Obtain p-value for test
End of explanation
"""


chi_squared_values = []
for i in range(100):
    nums = np.random.random((32561,))
    ex = 32561/2
    nums[nums < 0.5] = 0
    nums[nums >= 0.5] = 1
    male = sum(nums==0)
    female = 32561-male
    male_diff = (male - ex)**2 / ex
    female_diff = (female - ex)**2 / ex
    chi_squared_values.append(male_diff + female_diff)
plt.hist(chi_squared_values)
plt.show()

# return error if unequal lengths

def chi_squared_statistic(o,e):
    """"""
    
    o = np.array(o)
    e = np.array(e)
    return np.sum(1.0*(o - e)**2 / e)


class ChiSquaredTest():
    
    def __init__(self, observed=None, expected=None):
        self.observed = observed
        self.expected = expected
        
        
    def plot_dsn(self, k=100):
        simulated_stats = self.simPermDsn(k=k)
        tit = "Chi-Squared Test Statistic Distribution for k={} Simulations".format(k)
        plot_hist(simulated_stats, title=tit, xlabel="Test Statistic Value", ylabel="Frequency")
    
    
    def testStat(self, input_observed=None, input_expected=None):
        if input_observed or input_expected:
            chisq_teststat = chi_squared_statistic(input_observed, input_expected)
        else:
            chisq_teststat = chi_squared_statistic(self.observed, self.expected)
        return chisq_teststat
    
    def simPermDsn(self, input_observed=None, input_expected=None, k=100):
        if input_observed or input_expected:
            test_stat = self.testStat(input_observed, input_expected)
            n = len(input_observed)
        else:
            test_stat = self.testStat()
            n = len(self.observed)
            
        # create sampling distribution
        chisqrd_vals = []
        for i in range(k):
            values = np.random.random((n,))
            ex = 1.0*n/2
            values[values<.5]=0
            values[values>=.5]=1
            diff1 = chi_squared_statistic(sum(values==0),ex)
            diff2 = chi_squared_statistic(sum(values==1),ex)
            chisqrd_vals.append(diff1+diff2)
        return chisqrd_vals
    
    def pvalue(self, input_observed=None, input_expected=None):
        
        current_stat = self.testStat(input_observed, input_expected)
        simulated_stats = self.simPermDsn(input_observed, input_expected)
        
        # p-value = proportion of test stats greater than ours
        p_value = 1.0*sum(simulated_stats >= current_stat) / len(simulated_stats)
        return p_value

"""
Explanation: Chi Squared Test
End of explanation
"""


obs = [26816, 3624, 1039, 311, 771]
ex = [26146.5, 3939.9, 1044.3, 310.5, 1119.8]
print(sum(obs))
print(sum(ex))

obs1 = [1,2,3,4,5]
ex1 = [1,2,2,4,5]

"""
Explanation: Create data
End of explanation
"""


chisq = ChiSquaredTest(obs, ex)

"""
Explanation: Create instance of ChiSquaredTest
End of explanation
"""


chisq.pvalue(obs1,ex1)

"""
Explanation: Get pvalue
End of explanation
"""


chisq.plot_dsn(k=250)

z=None

# To Implement

# Levene's Test for Unequal Variances
# Mann-Whitney Test    http://www.stat.wisc.edu/~st571-1/Fall2005/lec18-21.1.pdf
# t-test
# ANOVA  http://www.marsja.se/four-ways-to-conduct-one-way-anovas-using-python/
z==None

"""
Explanation: View bootstrapped distribution
End of explanation
"""


