def age_design(indices):
  tmp = np.hstack((sm.categorical(hrdat['sex'][indices])[:,2:],
                  sm.categorical(hrdat['educ'][indices])[:,2:],
                  sm.categorical(hrdat['PTFT'][indices])[:,2:],
                  hrdat['age'].reshape(n,1)[indices,:],
                  (hrdat['age']**2).reshape(n,1)[indices,:]))
  return sm.add_constant(tmp, prepend = True)

    def checkOLS(self, exog, endog, x, y):

        try:
            import scikits.statsmodels.api as sm
        except ImportError:
            import scikits.statsmodels as sm

        reference = sm.OLS(endog, sm.add_constant(exog)).fit()

        result = ols(y=y, x=x)

        assert_almost_equal(reference.params, result._beta_raw)
        assert_almost_equal(reference.df_model, result._df_model_raw)
        assert_almost_equal(reference.df_resid, result._df_resid_raw)
        assert_almost_equal(reference.fvalue, result._f_stat_raw[0])
        assert_almost_equal(reference.pvalues, result._p_value_raw)
        assert_almost_equal(reference.rsquared, result._r2_raw)
        assert_almost_equal(reference.rsquared_adj, result._r2_adj_raw)
        assert_almost_equal(reference.resid, result._resid_raw)
        assert_almost_equal(reference.bse, result._std_err_raw)
        assert_almost_equal(reference.t(), result._t_stat_raw)
        assert_almost_equal(reference.cov_params(), result._var_beta_raw)
        assert_almost_equal(reference.fittedvalues, result._y_fitted_raw)

        _check_non_raw_results(result)

def test_HC_use():
    np.random.seed(0)
    nsample = 100
    x = np.linspace(0,10, 100)
    X = sm.add_constant(np.column_stack((x, x**2)), prepend=False)
    beta = np.array([1, 0.1, 10])
    y = np.dot(X, beta) + np.random.normal(size=nsample)

    results = sm.OLS(y, X).fit()

    #test cov_params
    idx = np.array([1,2])
    #need to call HC0_se to have cov_HC0 available
    results.HC0_se
    cov12 = results.cov_params(column=[1,2], cov_p=results.cov_HC0)
    assert_almost_equal(cov12, results.cov_HC0[idx[:,None], idx], decimal=15)

    #test t_test
    tvals = results.params/results.HC0_se
    ttest = results.t_test(np.eye(3), cov_p=results.cov_HC0)
    assert_almost_equal(ttest.tvalue, tvals, decimal=14)
    assert_almost_equal(ttest.sd, results.HC0_se, decimal=14)

    #test f_test
    ftest = results.f_test(np.eye(3)[:-1], cov_p=results.cov_HC0)
    slopes = results.params[:-1]
    idx = np.array([0,1])
    cov_slopes = results.cov_HC0[idx[:,None], idx]
    fval = np.dot(slopes, np.linalg.inv(cov_slopes).dot(slopes))/len(idx)
    assert_almost_equal(ftest.fvalue, fval, decimal=12)

 def setupClass(cls):
     data = sm.datasets.spector.load()
     data.exog = sm.add_constant(data.exog)
     res2 = Spector()
     res2.probit()
     cls.res2 = res2
     cls.res1 = Probit(data.endog, data.exog).fit(method="ncg", disp=0, avextol=1e-8)

    def __init__(self):

        # generate artificial data
        np.random.seed(98765678)
        nobs = 200
        rvs = np.random.randn(nobs,6)
        data_exog = rvs
        data_exog = sm.add_constant(data_exog)
        xbeta = 1 + 0.1*rvs.sum(1)
        data_endog = np.random.poisson(np.exp(xbeta))

        #estimate discretemod.Poisson as benchmark
        self.res_discrete = Poisson(data_endog, data_exog).fit(disp=0)

        mod_glm = sm.GLM(data_endog, data_exog, family=sm.families.Poisson())
        self.res_glm = mod_glm.fit()

        #estimate generic MLE
        #self.mod = PoissonGMLE(data_endog, data_exog)
        #res = self.mod.fit()
        offset = self.res_discrete.params[0] * data_exog[:,0]  #1d ???
        #self.res = PoissonOffsetGMLE(data_endog, data_exog[:,1:], offset=offset).fit(start_params = np.ones(6)/2., method='nm')
        modo = PoissonOffsetGMLE(data_endog, data_exog[:,1:], offset=offset)
        self.res = modo.fit(start_params = 0.9*self.res_discrete.params[1:],
                            method='nm', disp=0)

 def setupClass(cls):
     data = sm.datasets.spector.load()
     data.exog = sm.add_constant(data.exog)
     cls.res1 = Logit(data.endog, data.exog).fit(method="newton", disp=0)
     res2 = Spector()
     res2.logit()
     cls.res2 = res2

 def __init__(self):
     data = sm.datasets.spector.load()
     data.exog = sm.add_constant(data.exog)
     #mod = sm.Probit(data.endog, data.exog)
     self.mod = sm.Logit(data.endog, data.exog)
     #res = mod.fit(method="newton")
     self.params = [np.array([1,0.25,1.4,-7])]

 def test_qqplot(self):
   #just test that it runs
   data = sm.datasets.longley.load()
   data.exog = sm.add_constant(data.exog)
   mod_fit = sm.OLS(data.endog, data.exog).fit()
   res = mod_fit.resid
   fig = sm.qqplot(res)
   plt.close(fig)

def run_WLS():
    import scikits.statsmodels.api as sm
    res = sm.WLS(y, sm.add_constant(x, prepend=True),
                 weights=1. / sigma ** 2).fit()
    print ('statsmodels.api.WLS')
    print('popt: {0}'.format(res.params))
    print('perr: {0}'.format(res.bse))
    return res

def quadratic_term(list_of_mean, list_of_var):
    """Fit a quadratic term and return its p-value"""
    # Remove records with 0 variance
    log_var = [np.log(x) for x in list_of_var if x > 0]
    log_mean = [np.log(list_of_mean[i]) for i in range(len(list_of_mean)) if list_of_var[i] > 0]
    log_mean_quad = [x ** 2 for x in log_mean]
    indep_var = np.column_stack((log_mean, log_mean_quad))
    indep_var = sm.add_constant(indep_var, prepend = True)
    quad_res = sm.OLS(log_var, indep_var).fit()
    return quad_res.pvalues[2]

def age_design(indices):
    tmp = np.hstack(
        (
            sm.categorical(hrdat["sex"][indices])[:, 2:],
            sm.categorical(hrdat["educ"][indices])[:, 2:],
            sm.categorical(hrdat["PTFT"][indices])[:, 2:],
            hrdat["age"].reshape(n, 1)[indices, :],
            (hrdat["age"] ** 2).reshape(n, 1)[indices, :],
        )
    )
    return sm.add_constant(tmp, prepend=True)

def explain_rseq_by_rfreq_and_copy():
    r_rseqs = [motif_ic(getattr(Escherichia_coli,tf)) for tf in Escherichia_coli.tfs
               if tf in copy_numbers]
    r_rfreqs = [log2(4.6*10**6/len(getattr(Escherichia_coli,tf)))
          for tf in Escherichia_coli.tfs
                if tf in copy_numbers]
    copies = [copy_numbers[tf] for tf in Escherichia_coli.tfs if tf in copy_numbers]
    log_copies = map(log2,copies)
    X = sm.add_constant(np.column_stack((r_rfreqs,log_copies)),prepend=True)
    res = sm.OLS(r_rseqs,X).fit()
    print res.summary()

def test_perfect_prediction():
    cur_dir = os.path.dirname(os.path.abspath(__file__))
    iris_dir = os.path.join(cur_dir, "..", "..", "genmod", "tests", "results")
    iris_dir = os.path.abspath(iris_dir)
    iris = np.genfromtxt(os.path.join(iris_dir, "iris.csv"), delimiter=",", skip_header=1)
    y = iris[:, -1]
    X = iris[:, :-1]
    X = X[y != 2]
    y = y[y != 2]
    X = sm.add_constant(X, prepend=True)
    mod = Logit(y, X)
    assert_raises(PerfectSeparationError, mod.fit)

def cm_test(X):
    """
    Conditional moment test.  X is a flat numpy array.
    """
    betahat, alphahat, shat = ar1_functions.fit(X)
    n = len(X)
    xL = X[:(n-1)]  #  All but the last one
    xF = X[1:]      #  All but the first one
    Z = (xF - betahat - alphahat * xL)**2 
    XX = sm.add_constant(xL)
    out = sm.OLS(Z, XX).fit()
    return np.abs(out.tvalues[0]) > 1.96

    def setup(self):
        nsample = 100
        sig = 0.5
        x1 = np.linspace(0, 20, nsample)
        x2 = 5 + 3* np.random.randn(nsample)
        X = np.c_[x1, x2, np.sin(0.5*x1), (x2-5)**2, np.ones(nsample)]
        beta = [0.5, 0.5, 1, -0.04, 5.]
        y_true = np.dot(X, beta)
        y = y_true + sig * np.random.normal(size=nsample)
        exog0 = sm.add_constant(np.c_[x1, x2], prepend=False)
        res = sm.OLS(y, exog0).fit()

        self.res = res

def linear_fit_robust(x, y, return_coef=False):
    """
    Fit a straight-line by robust regression (M-estimate).

    If `return_coef=True` returns the slope (m) and intercept (c).
    """
    import scikits.statsmodels.api as sm
    ind, = np.where((~np.isnan(x)) & (~np.isnan(y)))
    x, y = x[ind], y[ind]
    X = sm.add_constant(x, prepend=False)
    y_model = sm.RLM(y, X, M=sm.robust.norms.HuberT())
    y_fit = y_model.fit()
    if return_coef:
        if len(y_fit.params) < 2: return (y_fit.params[0], 0.)
        else: return y_fit.params[:]
    else:
        return (x, y_fit.fittedvalues)

    def _check_wls(self, x, y, weights):
        result = ols(y=y, x=x, weights=1/weights)

        combined = x.copy()
        combined['__y__'] = y
        combined['__weights__'] = weights
        combined = combined.dropna()

        endog = combined.pop('__y__').values
        aweights = combined.pop('__weights__').values
        exog = sm.add_constant(combined.values, prepend=False)

        sm_result = sm.WLS(endog, exog, weights=1/aweights).fit()

        assert_almost_equal(sm_result.params, result._beta_raw)
        assert_almost_equal(sm_result.resid, result._resid_raw)

        self.checkMovingOLS('rolling', x, y, weights=weights)
        self.checkMovingOLS('expanding', x, y, weights=weights)

def regression_analysis(play_arr,dataFunction1,dataFunction2):

	totalBefore = []
	totalAfter = []
	for weekNum in range(10,15):
		Before,After = regression_weekly(play_arr,weekNum,dataFunction1, dataFunction2)

		totalBefore = np.concatenate([totalBefore, Before])
		totalAfter = np.concatenate([totalAfter, After])

	slope, intercept, r_value, p_value, err = stats.linregress(totalBefore, totalAfter)
	results = sm.OLS(totalAfter, sm.add_constant(totalBefore)).fit()

	print results.summary()

	plt.plot(totalBefore, totalAfter, '.')
	X_plot = np.linspace(0, 1, 100)
	plt.plot(X_plot, X_plot * results.params[0] + results.params[1])
	plt.show()

    def calc_factors(self, x=None, keepdim=0, addconst=True):
        '''get factor decomposition of exogenous variables

        This uses principal component analysis to obtain the factors. The number
        of factors kept is the maximum that will be considered in the regression.
        '''
        if x is None:
            x = self.exog
        else:
            x = np.asarray(x)
        xred, fact, evals, evecs  = pca(x, keepdim=keepdim, normalize=1)
        self.exog_reduced = xred
        #self.factors = fact
        if addconst:
            self.factors = sm.add_constant(fact, prepend=True)
            self.hasconst = 1  #needs to be int
        else:
            self.factors = fact
            self.hasconst = 0  #needs to be int

        self.evals = evals
        self.evecs = evecs

    def checkOLS(self, exog, endog, x, y):
        reference = sm.OLS(endog, sm.add_constant(exog, prepend=False)).fit()
        result = ols(y=y, x=x)

        # check that sparse version is the same
        sparse_result = ols(y=y.to_sparse(), x=x.to_sparse())
        _compare_ols_results(result, sparse_result)

        assert_almost_equal(reference.params, result._beta_raw)
        assert_almost_equal(reference.df_model, result._df_model_raw)
        assert_almost_equal(reference.df_resid, result._df_resid_raw)
        assert_almost_equal(reference.fvalue, result._f_stat_raw[0])
        assert_almost_equal(reference.pvalues, result._p_value_raw)
        assert_almost_equal(reference.rsquared, result._r2_raw)
        assert_almost_equal(reference.rsquared_adj, result._r2_adj_raw)
        assert_almost_equal(reference.resid, result._resid_raw)
        assert_almost_equal(reference.bse, result._std_err_raw)
        assert_almost_equal(reference.tvalues, result._t_stat_raw)
        assert_almost_equal(reference.cov_params(), result._var_beta_raw)
        assert_almost_equal(reference.fittedvalues, result._y_fitted_raw)

        _check_non_raw_results(result)