from os.path import basename, exists


def download(url):
    filename = basename(url)
    if not exists(filename):
        from urllib.request import urlretrieve

        local, _ = urlretrieve(url, filename)
        print("Downloaded " + local)


download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/thinkstats2.py")
download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/thinkplot.py")

import numpy as np

import random

import thinkstats2
import thinkplot

"""
Explanation: Chapter 10
Examples and Exercises from Think Stats, 2nd Edition
http://thinkstats2.com
Copyright 2016 Allen B. Downey
MIT License: https://opensource.org/licenses/MIT
End of explanation
"""


download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/nsfg.py")
download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/first.py")
download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/2002FemPreg.dct")
download(
    "https://github.com/AllenDowney/ThinkStats2/raw/master/code/2002FemPreg.dat.gz"
)

import first
live, firsts, others = first.MakeFrames()
live = live.dropna(subset=['agepreg', 'totalwgt_lb'])
ages = live.agepreg
weights = live.totalwgt_lb

"""
Explanation: Least squares
One more time, let's load up the NSFG data.
End of explanation
"""


from thinkstats2 import Mean, MeanVar, Var, Std, Cov

def LeastSquares(xs, ys):
    meanx, varx = MeanVar(xs)
    meany = Mean(ys)

    slope = Cov(xs, ys, meanx, meany) / varx
    inter = meany - slope * meanx

    return inter, slope

"""
Explanation: The following function computes the intercept and slope of the least squares fit.
End of explanation
"""


inter, slope = LeastSquares(ages, weights)
inter, slope

"""
Explanation: Here's the least squares fit to birth weight as a function of mother's age.
End of explanation
"""


inter + slope * 25

"""
Explanation: The intercept is often easier to interpret if we evaluate it at the mean of the independent variable.
End of explanation
"""


slope * 10

"""
Explanation: And the slope is easier to interpret if we express it in pounds per decade (or ounces per year).
End of explanation
"""


def FitLine(xs, inter, slope):
    fit_xs = np.sort(xs)
    fit_ys = inter + slope * fit_xs
    return fit_xs, fit_ys

"""
Explanation: The following function evaluates the fitted line at the given xs.
End of explanation
"""


fit_xs, fit_ys = FitLine(ages, inter, slope)

"""
Explanation: And here's an example.
End of explanation
"""


thinkplot.Scatter(ages, weights, color='blue', alpha=0.1, s=10)
thinkplot.Plot(fit_xs, fit_ys, color='white', linewidth=3)
thinkplot.Plot(fit_xs, fit_ys, color='red', linewidth=2)
thinkplot.Config(xlabel="Mother's age (years)",
                 ylabel='Birth weight (lbs)',
                 axis=[10, 45, 0, 15],
                 legend=False)

"""
Explanation: Here's a scatterplot of the data with the fitted line.
End of explanation
"""


def Residuals(xs, ys, inter, slope):
    xs = np.asarray(xs)
    ys = np.asarray(ys)
    res = ys - (inter + slope * xs)
    return res

"""
Explanation: Residuals
The following functon computes the residuals.
End of explanation
"""


live['residual'] = Residuals(ages, weights, inter, slope)

"""
Explanation: Now we can add the residuals as a column in the DataFrame.
End of explanation
"""


bins = np.arange(10, 48, 3)
indices = np.digitize(live.agepreg, bins)
groups = live.groupby(indices)

age_means = [group.agepreg.mean() for _, group in groups][1:-1]
age_means

"""
Explanation: To visualize the residuals, I'll split the respondents into groups by age, then plot the percentiles of the residuals versus the average age in each group.
First I'll make the groups and compute the average age in each group.
End of explanation
"""


cdfs = [thinkstats2.Cdf(group.residual) for _, group in groups][1:-1]

"""
Explanation: Next I'll compute the CDF of the residuals in each group.
End of explanation
"""


def PlotPercentiles(age_means, cdfs):
    thinkplot.PrePlot(3)
    for percent in [75, 50, 25]:
        weight_percentiles = [cdf.Percentile(percent) for cdf in cdfs]
        label = '%dth' % percent
        thinkplot.Plot(age_means, weight_percentiles, label=label)

"""
Explanation: The following function plots percentiles of the residuals against the average age in each group.
End of explanation
"""


PlotPercentiles(age_means, cdfs)

thinkplot.Config(xlabel="Mother's age (years)",
                 ylabel='Residual (lbs)',
                 xlim=[10, 45])

"""
Explanation: The following figure shows the 25th, 50th, and 75th percentiles.
Curvature in the residuals suggests a non-linear relationship.
End of explanation
"""


def SampleRows(df, nrows, replace=False):
    """Choose a sample of rows from a DataFrame.

    df: DataFrame
    nrows: number of rows
    replace: whether to sample with replacement

    returns: DataDf
    """
    indices = np.random.choice(df.index, nrows, replace=replace)
    sample = df.loc[indices]
    return sample

def ResampleRows(df):
    """Resamples rows from a DataFrame.

    df: DataFrame

    returns: DataFrame
    """
    return SampleRows(df, len(df), replace=True)

"""
Explanation: Sampling distribution
To estimate the sampling distribution of inter and slope, I'll use resampling.
End of explanation
"""


def SamplingDistributions(live, iters=101):
    t = []
    for _ in range(iters):
        sample = ResampleRows(live)
        ages = sample.agepreg
        weights = sample.totalwgt_lb
        estimates = LeastSquares(ages, weights)
        t.append(estimates)

    inters, slopes = zip(*t)
    return inters, slopes

"""
Explanation: The following function resamples the given dataframe and returns lists of estimates for inter and slope.
End of explanation
"""


inters, slopes = SamplingDistributions(live, iters=1001)

"""
Explanation: Here's an example.
End of explanation
"""


def Summarize(estimates, actual=None):
    mean = Mean(estimates)
    stderr = Std(estimates, mu=actual)
    cdf = thinkstats2.Cdf(estimates)
    ci = cdf.ConfidenceInterval(90)
    print('mean, SE, CI', mean, stderr, ci)

"""
Explanation: The following function takes a list of estimates and prints the mean, standard error, and 90% confidence interval.
End of explanation
"""


Summarize(inters)

"""
Explanation: Here's  the summary for inter.
End of explanation
"""


Summarize(slopes)

"""
Explanation: And for slope.
End of explanation
"""


for slope, inter in zip(slopes, inters):
    fxs, fys = FitLine(age_means, inter, slope)
    thinkplot.Plot(fxs, fys, color='gray', alpha=0.01)
    
thinkplot.Config(xlabel="Mother's age (years)",
                 ylabel='Residual (lbs)',
                 xlim=[10, 45])

"""
Explanation: Exercise: Use ResampleRows and generate a list of estimates for the mean birth weight.  Use Summarize to compute the SE and CI for these estimates.
Visualizing uncertainty
To show the uncertainty of the estimated slope and intercept, we can generate a fitted line for each resampled estimate and plot them on top of each other.
End of explanation
"""


def PlotConfidenceIntervals(xs, inters, slopes, percent=90, **options):
    fys_seq = []
    for inter, slope in zip(inters, slopes):
        fxs, fys = FitLine(xs, inter, slope)
        fys_seq.append(fys)

    p = (100 - percent) / 2
    percents = p, 100 - p
    low, high = thinkstats2.PercentileRows(fys_seq, percents)
    thinkplot.FillBetween(fxs, low, high, **options)

"""
Explanation: Or we can make a neater (and more efficient plot) by computing fitted lines and finding percentiles of the fits for each value of the dependent variable.
End of explanation
"""


PlotConfidenceIntervals(age_means, inters, slopes, percent=90, 
                        color='gray', alpha=0.3, label='90% CI')
PlotConfidenceIntervals(age_means, inters, slopes, percent=50,
                        color='gray', alpha=0.5, label='50% CI')

thinkplot.Config(xlabel="Mother's age (years)",
                 ylabel='Residual (lbs)',
                 xlim=[10, 45])

"""
Explanation: This example shows the confidence interval for the fitted values at each mother's age.
End of explanation
"""


def CoefDetermination(ys, res):
    return 1 - Var(res) / Var(ys)

"""
Explanation: Coefficient of determination
The coefficient compares the variance of the residuals to the variance of the dependent variable.
End of explanation
"""


inter, slope = LeastSquares(ages, weights)
res = Residuals(ages, weights, inter, slope)
r2 = CoefDetermination(weights, res)
r2

"""
Explanation: For birth weight and mother's age $R^2$ is very small, indicating that the mother's age predicts a small part of the variance in birth weight.
End of explanation
"""


print('rho', thinkstats2.Corr(ages, weights))
print('R', np.sqrt(r2))    

"""
Explanation: We can confirm that $R^2 = \rho^2$:
End of explanation
"""


print('Std(ys)', Std(weights))
print('Std(res)', Std(res))

"""
Explanation: To express predictive power, I think it's useful to compare the standard deviation of the residuals to the standard deviation of the dependent variable, as a measure RMSE if you try to guess birth weight with and without taking into account mother's age.
End of explanation
"""


var_ys = 15**2
rho = 0.72
r2 = rho**2
var_res = (1 - r2) * var_ys
std_res = np.sqrt(var_res)
std_res

"""
Explanation: As another example of the same idea, here's how much we can improve guesses about IQ if we know someone's SAT scores.
End of explanation
"""


class SlopeTest(thinkstats2.HypothesisTest):

    def TestStatistic(self, data):
        ages, weights = data
        _, slope = thinkstats2.LeastSquares(ages, weights)
        return slope

    def MakeModel(self):
        _, weights = self.data
        self.ybar = weights.mean()
        self.res = weights - self.ybar

    def RunModel(self):
        ages, _ = self.data
        weights = self.ybar + np.random.permutation(self.res)
        return ages, weights

"""
Explanation: Hypothesis testing with slopes
Here's a HypothesisTest that uses permutation to test whether the observed slope is statistically significant.
End of explanation
"""


ht = SlopeTest((ages, weights))
pvalue = ht.PValue()
pvalue

"""
Explanation: And it is.
End of explanation
"""


ht.actual, ht.MaxTestStat()

"""
Explanation: Under the null hypothesis, the largest slope we observe after 1000 tries is substantially less than the observed value.
End of explanation
"""


sampling_cdf = thinkstats2.Cdf(slopes)

"""
Explanation: We can also use resampling to estimate the sampling distribution of the slope.
End of explanation
"""


thinkplot.PrePlot(2)
thinkplot.Plot([0, 0], [0, 1], color='0.8')
ht.PlotCdf(label='null hypothesis')

thinkplot.Cdf(sampling_cdf, label='sampling distribution')

thinkplot.Config(xlabel='slope (lbs / year)',
                   ylabel='CDF',
                   xlim=[-0.03, 0.03],
                   legend=True, loc='upper left')

"""
Explanation: The distribution of slopes under the null hypothesis, and the sampling distribution of the slope under resampling, have the same shape, but one has mean at 0 and the other has mean at the observed slope.
To compute a p-value, we can count how often the estimated slope under the null hypothesis exceeds the observed slope, or how often the estimated slope under resampling falls below 0.
End of explanation
"""


pvalue = sampling_cdf[0]
pvalue

"""
Explanation: Here's how to get a p-value from the sampling distribution.
End of explanation
"""


def ResampleRowsWeighted(df, column='finalwgt'):
    weights = df[column]
    cdf = thinkstats2.Cdf(dict(weights))
    indices = cdf.Sample(len(weights))
    sample = df.loc[indices]
    return sample

"""
Explanation: Resampling with weights
Resampling provides a convenient way to take into account the sampling weights associated with respondents in a stratified survey design.
The following function resamples rows with probabilities proportional to weights.
End of explanation
"""


iters = 100
estimates = [ResampleRowsWeighted(live).totalwgt_lb.mean()
             for _ in range(iters)]
Summarize(estimates)

"""
Explanation: We can use it to estimate the mean birthweight and compute SE and CI.
End of explanation
"""


estimates = [thinkstats2.ResampleRows(live).totalwgt_lb.mean()
             for _ in range(iters)]
Summarize(estimates)

"""
Explanation: And here's what the same calculation looks like if we ignore the weights.
End of explanation
"""


download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/brfss.py")
download("https://github.com/AllenDowney/ThinkStats2/raw/master/code/CDBRFS08.ASC.gz")

import brfss

df = brfss.ReadBrfss(nrows=None)
df = df.dropna(subset=['htm3', 'wtkg2'])
heights, weights = df.htm3, df.wtkg2
log_weights = np.log10(weights)

"""
Explanation: The difference is non-negligible, which suggests that there are differences in birth weight between the strata in the survey.
Exercises
Exercise: Using the data from the BRFSS, compute the linear least squares fit for log(weight) versus height. How would you best present the estimated parameters for a model like this where one of the variables is log-transformed? If you were trying to guess someone’s weight, how much would it help to know their height?
Like the NSFG, the BRFSS oversamples some groups and provides a sampling weight for each respondent. In the BRFSS data, the variable name for these weights is totalwt. Use resampling, with and without weights, to estimate the mean height of respondents in the BRFSS, the standard error of the mean, and a 90% confidence interval. How much does correct weighting affect the estimates?
Read the BRFSS data and extract heights and log weights.
End of explanation
"""


