def generate_autocorrelated_data(theta, mu, sigma, N):
    X = np.zeros((N, 1))
    for t in range(1, N):
        X[t] = theta * X[t-1] + np.random.normal(mu, sigma)
    return X

def newey_west_SE(data):
    ind = range(0, len(data))
    ind = sm.add_constant(ind)
    model = regression.linear_model.OLS(data, ind).fit(cov_type='HAC',cov_kwds={'maxlags':1})
    return model.bse[0]

def newey_west_matrix(data):
    ind = range(0, len(data))
    ind = sm.add_constant(ind)
    model = regression.linear_model.OLS(data, ind).fit()
    return sw.cov_hac(model)

# Useful Libraries
import numpy as np
import seaborn as sns
from scipy import stats
import matplotlib.pyplot as plt
from statsmodels.stats.stattools import jarque_bera
import statsmodels.stats.sandwich_covariance as sw
from statsmodels import regression
import statsmodels.api as sm

"""
Explanation: Exercises: Confidence Intervals - Answer Key
By Christopher Fenaroli and Delaney Mackenzie
Lecture Link:
https://www.quantopian.com/lectures/confidence-intervals
IMPORTANT NOTE:
This lecture corresponds to the Confidence Intervals lecture, which is part of the Quantopian lecture series. This homework expects you to rely heavily on the code presented in the corresponding lecture. Please copy and paste regularly from that lecture when starting to work on the problems, as trying to do them from scratch will likely be too difficult.
When you feel comfortable with the topics presented here, see if you can create an algorithm that qualifies for the Quantopian Contest. Participants are evaluated on their ability to produce risk-constrained alpha and the top 10 contest participants are awarded cash prizes on a daily basis.
https://www.quantopian.com/contest
Part of the Quantopian Lecture Series:

www.quantopian.com/lectures
github.com/quantopian/research_public


Key Concepts
End of explanation
"""


np.random.seed(11)
POPULATION_MU = 105
POPULATION_SIGMA = 20
sample_size = 50

"""
Explanation: Data
End of explanation
"""


sample = np.random.normal(POPULATION_MU, POPULATION_SIGMA, sample_size)

#Your code goes here

Mean = np.mean(sample)

print "Mean:", Mean

"""
Explanation: Exercise 1: Determining Confidence Intervals
a. Mean
Determine the mean of the following artificial data in sample.
End of explanation
"""


#Your code goes here

SD = np.std(sample)

print "Standard Deviation:", SD

"""
Explanation: b. Standard Deviation
Determine standard deviation of the sample.
End of explanation
"""


#Your code goes here

SE = SD / np.sqrt(sample_size)

print "Standard Error:", SE

"""
Explanation: c. Standard Error
Using the standard deviation and sample_size, determine the standard error for the sample.
End of explanation
"""


#Your code goes here

print "95% Confidence Interval:", (-1.96 * SE + Mean, 1.96 * SE + Mean)
print "90% Confidence Interval:", (-1.64 * SE + Mean, 1.64 * SE + Mean)
print "80% Confidence Interval:", (-1.28 * SE + Mean, 1.28 * SE + Mean)

"""
Explanation: d. Confidence Intervals
Using the standard error and mean, determine 95% (Z = 1.96), 90% (Z = 1.64), and 80% (Z = 1.28) confidence intervals for the sample.
End of explanation
"""


n = 1000
correct = 0
samples = [np.random.normal(loc=POPULATION_MU, scale=POPULATION_SIGMA, size=sample_size) for i in range(n)]

#Your code goes here

for i in range(n):
    sample_mean = np.mean(samples[i])
    sample_SE = stats.sem(samples[i])
    if ((POPULATION_MU >= -1.96 * sample_SE + sample_mean) and (POPULATION_MU <= 1.96 * sample_SE + sample_mean)):
        correct += 1;

print "Expected Result:", .95 * 1000
print "Actual Result", correct

"""
Explanation: Exercise 2: Interpreting Confidence Intervals
Assuming our interval was correctly calculated and that the underlying data was independent, if we take many samples and make many 95% confidence intervals, the intervals will contain the true mean 95% of the time. Run 1000 samples and measure how many of their confidence intervals actually contain the true mean.
End of explanation
"""


n = 500
normal_samples = [np.mean(np.random.normal(loc=POPULATION_MU, scale=POPULATION_SIGMA, size=sample_size)) for i in range(n)]

#Your code goes here

plt.hist(normal_samples, 10)

_, pvalue, _, _ = jarque_bera(normal_samples)

print pvalue
if pvalue > 0.05:
    print 'The distribution of sample means is likely normal.'
else:
    print 'The distribution of sample means is likely not normal.'

"""
Explanation: Exercise 3: Central Limit Theorem
a. Plotting Sample Means - Normal
Assuming our samples are independent, the distribution of the sample means should be normally distributed, regardless of the underlying distribution. 
Draw 500 samples of size sample_size from the same normal distribution from question 1, plot the means of each of the samples, and check to see if the distribution of the sample means is normal using the jarque_bera function (see here more information on the Jarque-Bera test)
End of explanation
"""


n = 500
expo_samples = [np.mean(np.random.exponential(POPULATION_MU, sample_size)) for i in range(n)]

#Your code goes here

plt.hist(expo_samples, 10)

_, pvalue, _, _ = jarque_bera(expo_samples)

print pvalue
if pvalue > 0.05:
    print 'The distribution of sample means is likely normal, despite the underlying distribution being non-normal (exponential).'
else:
    print 'The distribution of sample means is likely not normal.'

"""
Explanation: b. Plotting Sample Means - Exponential
Draw 500 samples of size sample_size from a new exponential distribution, plot the means of each of the samples, and check to see if the distribution of the sample means is normal.
End of explanation
"""


n = 500
autocorrelated_samples = [(generate_autocorrelated_data(0.5, 0, 1, sample_size) + POPULATION_MU) for i in range(n)]
autocorrelated_means = [np.mean(autocorrelated_samples[i]) for i in range(n)]

#Your code goes here

plt.hist(autocorrelated_means, 10)

_, pvalue, _, _ = jarque_bera(autocorrelated_means)

print pvalue
if pvalue > 0.05:
    print 'The distribution of sample means is likely normal, despite an autocorrelated underlying distribution.'
else:
    print 'The distribution of sample means is likely not normal.'

"""
Explanation: c.i Plotting Sample Means - Autocorrelated
Draw 500 samples of size sample_size from a new autocorrelated (dependent) distribution, plot the means of each of the samples, and check to see if the distribution of the sample means is normal.
End of explanation
"""


n = 500
autocorrelated_samples = [(generate_autocorrelated_data(0.5, 0, 1, sample_size) + POPULATION_MU) for i in range(n)]
autocorrelated_stds = [np.std(autocorrelated_samples[i]) for i in range(n)]

#Your code goes here

plt.hist(autocorrelated_stds, 10)

_, pvalue, _, _ = jarque_bera(autocorrelated_stds)

print pvalue
if pvalue > 0.05:
    print 'The distribution of sample standard deviations is likely normal.'
else:
    print 'The distribution of sample standard deviations is likely not normal, due to the autocorrelated underlying distribution and the different assumptions for the CLT for means and for standard deviations.'
    

"""
Explanation: c.ii Plotting Sample Standard Deviations - Autocorrelated
Draw 500 samples of size sample_size from the same autocorrelated distribution, plot the standard deviations of each of the samples, and check to see if the distribution of the sample standard deviations is normal.
End of explanation
"""


n = 100
small_size = 3
correct = 0
samples = [np.random.normal(loc=POPULATION_MU, scale=POPULATION_SIGMA, size=small_size) for i in range(n)]
    
#Your code goes here
    
for i in range(n):
    sample_mean = np.mean(samples[i])
    sample_SE = stats.sem(samples[i])
    if ((POPULATION_MU >= -1.96 * sample_SE + sample_mean) and (POPULATION_MU <= 1.96 * sample_SE + sample_mean)):
        correct += 1

print "Expected Result:", .95 * n
print "Actual Result:", correct
print "Due to the small sample size, the actual number of confidence intervals containing the population mean is much lower than what we would expect given a correctly calibrated interval."


"""
Explanation: Exercise 4: Small Sample Sizes
a. Error Due to Small Sample Size
Run 100 samples of size small_size and measure how many of their 95% confidence intervals actually contain the true mean.
End of explanation
"""


n = 100
small_size = 5
correct = 0
samples = [np.random.normal(loc=POPULATION_MU, scale=POPULATION_SIGMA, size=small_size) for i in range(n)]

#Your code goes here
    
for i in range(n):
    sample_mean = np.mean(samples[i])
    sample_SE = stats.sem(samples[i])
    h = sample_SE * stats.t.ppf((1+0.95) / 2, len(samples[i])-1)
    if ((POPULATION_MU >= sample_mean - h) and (POPULATION_MU <= sample_mean + h)):
        correct += 1

print "Expected Result:", .95 * n
print "Actual Result:", correct
print "After using the t-distribution to correct for the smaller sample size, the actual number of confidence intervals containing the population mean is about what we expected."


"""
Explanation: b. T-distribution Correction
Run 100 samples of size small_size, this time accouting for the small sample size using a t-distribution, and measure how many of their 95% confidence intervals actually contain the true mean.
End of explanation
"""


n = 100
correct = 0
theta = 0.5
noise_mu = 0
noise_sigma = 1

#Your code goes here

for i in range(n):
    X = generate_autocorrelated_data(theta, noise_mu, noise_sigma, sample_size) + POPULATION_MU
    sample_mean = np.mean(X)
    sample_SE = np.std(X) / np.sqrt(sample_size)
    if ((POPULATION_MU >= -1.96 * sample_SE + sample_mean) and (POPULATION_MU <= 1.96 * sample_SE + sample_mean)):
        correct += 1
    

print "Expected Result:", .95 * n
print "Actual Result:", correct
print "Because the underlying data was autocorrelated, the actual number of confidence intervals containing the population mean is much lower than what we expected."


"""
Explanation: Exercise 5: Dependence
a. Error due to Dependence
Run 100 samples of the following autocorrelated distribution and measure how many of their 95% confidence intervals actually contain the true mean. (Use the helper function generate_autocorrelated_data(theta, noise_mu, noise_sigma, sample_size) to generate the samples)
End of explanation
"""


n = 100
correct = 0

#Your code goes here

for i in range(n):
    X = generate_autocorrelated_data(theta, noise_mu, noise_sigma, sample_size) + POPULATION_MU
    sample_mean = np.mean(X)
    sample_SE = np.std(X) / np.sqrt(sample_size)
    h = sample_SE * stats.t.ppf((1+0.95) / 2, len(X)-1)
    if ((POPULATION_MU >= sample_mean - h) and (POPULATION_MU <= sample_mean + h)):
        correct += 1
    

print "Expected Result:", .95 * n
print "Actual Result:", correct
print "We did not see a significant improvement in the actual number of confidence intervals containing the population mean. This is because a t-distribution only corrects for small sample sizes, not autocorrelation."


"""
Explanation: b. T-distribution Correction
Run 100 samples from the autocorrelated distribution, this time attempting to account for the autocorrelation using a t-distribution, and measure how many of their 95% confidence intervals actually contain the true mean to see if the correction works.
End of explanation
"""


X = generate_autocorrelated_data(theta, noise_mu, noise_sigma, sample_size) + POPULATION_MU

#Your code goes here

print newey_west_matrix(X)

"""
Explanation: c. Newey-West Matrix
Use the newey_west_matrix helper function to compute an adjusted (robust) covariance matrix for a single sample of the autocorrelated data.
End of explanation
"""


n = 100
correct = 0

#Your code goes here

for i in range(n):
    X = generate_autocorrelated_data(theta, noise_mu, noise_sigma, sample_size) + POPULATION_MU
    sample_mean = np.mean(X)
    sample_SE = newey_west_SE(X)
    if ((POPULATION_MU >= -1.96 * sample_SE + sample_mean) and (POPULATION_MU <= 1.96 * sample_SE + sample_mean)):
        correct += 1

print "New Standard Error:", sample_SE
print "Expected Result:", .95 * n
print "Actual Result:", correct
print "After accounting for autocorrelation by finding a Newey-West standard error, the actual number of confidence intervals containing the population mean is about what we expected."


"""
Explanation: d. Newey-West Correction
Run 100 samples of the following autocorrelated distribution, this time accounting for the autocorrelation by using a Newey-West correction on the standard error, and measure how many of their 95% confidence intervals actually contain the true mean to see if the correction works. (Use the helper function newey_west_SE to find the corrected standard error)
End of explanation
"""


