from __future__ import print_function, division

import numpy as np
import pandas as pd

import first
import thinkstats2
import thinkplot

%matplotlib inline

"""
Explanation: Regression
Copyright 2015 Allen Downey
License: Creative Commons Attribution 4.0 International
End of explanation
"""


live, firsts, others = first.MakeFrames()
live.shape

"""
Explanation: Let's load up the NSFG data again.
End of explanation
"""


live = live[live.prglngth>=37]
live.shape

"""
Explanation: And select live, full-term births.
End of explanation
"""


live = live.dropna(subset=['agepreg', 'totalwgt_lb'])
live.shape

"""
Explanation: And drop rows with missing data (just for the variables we want).
End of explanation
"""


live.head()

"""
Explanation: Check a few rows:
End of explanation
"""


live[['agepreg', 'totalwgt_lb']].describe()

"""
Explanation: And summarize a few variables.
End of explanation
"""


ages = live.agepreg
weights = live.totalwgt_lb
thinkplot.Scatter(ages, weights, alpha=0.1, s=15)
thinkplot.Config(xlabel='age (years)',
                 ylabel='weight (lbs)',
                 xlim=[10, 45],
                 ylim=[0, 15],
                 legend=False)

"""
Explanation: Here's a scatterplot of age and birthweight, with parameters tuned to avoid saturation.
End of explanation
"""


live['agepreg'].mean()

"""
Explanation: Mean of mother's age:
End of explanation
"""


live['totalwgt_lb'].mean(), live['totalwgt_lb'].std()

"""
Explanation: Mean and standard deviation of birthweight:
End of explanation
"""


thinkstats2.Corr(ages, weights)

"""
Explanation: And the coefficient of correlation:
End of explanation
"""


live['totalwgt_lb'].corr(live['agepreg'])

"""
Explanation: The Pandas corr function gets the same result:
End of explanation
"""


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

ages = [group.agepreg.mean() for i, group in groups][1:-1]
cdfs = [thinkstats2.Cdf(group.totalwgt_lb) for i, group in groups][1:-1]

thinkplot.PrePlot(5)
for percent in [90, 75, 50, 25, 10]:
    weights = [cdf.Percentile(percent) for cdf in cdfs]
    label = '%dth' % percent
    thinkplot.Plot(ages, weights, label=label)

thinkplot.Config(xlabel="mother's age (years)",
                 ylabel='birth weight (lbs)',
                 xlim=[14, 50],
                 legend=True)

"""
Explanation: To see the relationship more clearly, we can group mother's age into 3-year bins and plot percentiles of birth weight for each bin.
End of explanation
"""


def ReadFemResp(dct_file='2002FemResp.dct',
                dat_file='2002FemResp.dat.gz',
                nrows=None):
    """Reads the NSFG respondent data.

    dct_file: string file name
    dat_file: string file name

    returns: DataFrame
    """
    dct = thinkstats2.ReadStataDct(dct_file)
    df = dct.ReadFixedWidth(dat_file, compression='gzip', nrows=nrows)
    return df

"""
Explanation: The first and last points are not very reliable, because they represent fewer data points.
It looks like there is a generally positive relationshop between birth weight and mother's age, possibly leveling or dropping for older mothers.
We can get more information about the mothers by reading the respondents file, which contains one row per respondent.
End of explanation
"""


resp = ReadFemResp()
resp.shape

"""
Explanation: There are 7643 respondents and 3087 variables about each.
End of explanation
"""


resp.index = resp.caseid
resp.head()

"""
Explanation: If we use the caseid variable as the index, we can look up respondents efficiently by id.
Here's what the first few rows look like:
End of explanation
"""


join = live.join(resp, on='caseid', rsuffix='_r')
join.shape

"""
Explanation: Now we can join the tables, using the caseid from each pregnancy record to find the corresponding respondent and (abstractly) copy over the additional variables.
So the joined table has one row for each pregnancy and all the columns from both tables.
End of explanation
"""


join.screentime.head()

"""
Explanation: The encoding for screentime is a colon-separated timestamp.
End of explanation
"""


join.screentime = pd.to_datetime(join.screentime)
join.screentime.head()

"""
Explanation: If we convert to a datetime object, we avoid some processing problems later.
End of explanation
"""


ages = join.agepreg
weights = join.totalwgt_lb

inter, slope = thinkstats2.LeastSquares(ages, weights)
inter, slope, slope*16*10

"""
Explanation: To estimate the effect of mother's age on birthweight, we can use a simple least squares fit.
End of explanation
"""


import statsmodels.formula.api as smf

formula = ('totalwgt_lb ~ agepreg')
results = smf.ols(formula, data=join).fit()
results.summary()

"""
Explanation: The slope is almost 3 ounces per decade.
We can do the same thing using Ordinary Least Squares from statsmodels:
End of explanation
"""


inter, slope = results.params
inter, slope

"""
Explanation: The results object contains the parameters (and all the other info in the table):
End of explanation
"""


slope * 16 * 10           # slope in ounces per decade

"""
Explanation: And the results are consistent with my implementation:
End of explanation
"""


join['isfirst'] = (join.birthord == 1)

formula = 'totalwgt_lb ~ isfirst'
results = smf.ols(formula, data=join).fit()
results.summary()

"""
Explanation: We can use a boolean variable as a predictor:
End of explanation
"""


results.params['isfirst[T.True]'] * 16

"""
Explanation: First babies are lighter by about 1.5 ounces.
End of explanation
"""


formula = 'totalwgt_lb ~ agepreg + isfirst'
results = smf.ols(formula, data=join).fit()
results.summary()

"""
Explanation: And we can make a model with multiple predictors.
End of explanation
"""


results.params['isfirst[T.True]'] * 16

"""
Explanation: If we control for mother's age, the difference in weight for first babies is cut to about 0.5 ounces (and no longer statistically significant).
End of explanation
"""


join['age2'] = join.agepreg**2

formula = 'totalwgt_lb ~ agepreg + age2'
results = smf.ols(formula, data=join).fit()
results.summary()

"""
Explanation: The relationship with age might be non-linear.  Adding a quadratic term helps a little, although note that the $R^2$ values for all of these models are very small.
End of explanation
"""


formula = 'totalwgt_lb ~ agepreg + age2 + isfirst'
results = smf.ols(formula, data=join).fit()
results.summary()

"""
Explanation: Now we can combine the quadratic age model with isfirst
End of explanation
"""


results.params['isfirst[T.True]'] * 16

"""
Explanation: Now the effect is cut to less that a third of an ounce, and very plausibly due to chance.
End of explanation
"""


formula = ('totalwgt_lb ~ agepreg + age2 + C(race) + '
           'nbrnaliv>1 + paydu==1 + totincr')
results = smf.ols(formula, data=join).fit()
results.summary()

"""
Explanation: Here's the best model I found, combining all variables that seemed plausibly predictive.
End of explanation
"""


live['isboy'] = (live.babysex==1).astype(int)
model = smf.logit('isboy ~ agepreg', data=live)    
results = model.fit()
results.summary()

"""
Explanation: All predictors are statistically significant, so the effects could be legit, but the $R^2$ value is still very small: this model doesn't provide much help for the office pool.
Logistic regression
Let's say we want to predict the sex of a baby based on information about the mother.
I'll start by creating a binary dependent variable, isboy, and checking for dependence on mother's age:
End of explanation
"""


log_odds_ratio = results.params['agepreg'] * 10
odds_ratio = np.exp(log_odds_ratio)
odds_ratio

"""
Explanation: The estimated parameter is 0.0016, which is small and not statistically significant.  So the apparent relationship might be due to chance.
But for the sake of the example, I'll take it at face value and work out the effect on the prediction.
A parameter in a logistic regression is a log odds ratio, so we can compute the odds ratio for a difference of 10 years in mother's age:
End of explanation
"""


p = 0.51
prior_odds = p / (1-p)
post_odds = prior_odds * odds_ratio
p = post_odds / (post_odds + 1)
p

"""
Explanation: And we can use the odds ratio to update a prior probability.  A mother at the mean age has a 51% chance of having a boy.
In the case a mother who is 10 years older has a 51.4% chance.
End of explanation
"""


formula = 'isboy ~ agepreg + hpagelb + birthord + C(race)'
model = smf.logit(formula, data=live)
results = model.fit()
results.summary()

"""
Explanation: I searched for other factors that might be predictive.  The most likely candidates turn out not to be statistically significant.
End of explanation
"""


exog = pd.DataFrame(model.exog, columns=model.exog_names)
endog = pd.DataFrame(model.endog, columns=[model.endog_names])

actual = endog['isboy']
baseline = actual.mean()
baseline

"""
Explanation: Again, taking these parameters at face values, we can use the model to make predictions.
The baseline strategy is to always guess boy, which yields accuracy of 50.8%
End of explanation
"""


predict = (results.predict() >= 0.5)
true_pos = predict * actual
true_neg = (1 - predict) * (1 - actual)

acc = (sum(true_pos) + sum(true_neg)) / len(actual)
acc

"""
Explanation: results.predict uses the model to generate predictions for the data.
Adding up the correct positive and negative predictions, we get accuracy 51.3%
End of explanation
"""


columns = ['agepreg', 'hpagelb', 'birthord', 'race']
new = pd.DataFrame([[39, 30, 1, 2]], columns=columns)
y = results.predict(new)
y

"""
Explanation: And we can use the model to generate a prediction for the office pool.
Suppose your hypothetical coworker is is 39 years old and white, her husband is 30, and they are expecting their first child.
End of explanation
"""


