"""examples for sandwich estimators of covariance

Author: Josef Perktold

"""

from statsmodels.compat.python import lzip

import numpy as np
from numpy.testing import assert_almost_equal

import statsmodels.api as sm
import statsmodels.stats.sandwich_covariance as sw

nobs = 100
kvars = 4  # including constant
x = np.random.randn(nobs, kvars - 1)
exog = sm.add_constant(x)
params_true = np.ones(kvars)
y_true = np.dot(exog, params_true)
sigma = 0.1 + np.exp(exog[:, -1])
endog = y_true + sigma * np.random.randn(nobs)

self = sm.OLS(endog, exog).fit()

print(self.HC3_se)
print(sw.se_cov(sw.cov_hc3(self)))
# test standalone refactoring
assert_almost_equal(sw.se_cov(sw.cov_hc0(self)), self.HC0_se, 15)
assert_almost_equal(sw.se_cov(sw.cov_hc1(self)), self.HC1_se, 15)
assert_almost_equal(sw.se_cov(sw.cov_hc2(self)), self.HC2_se, 15)
assert_almost_equal(sw.se_cov(sw.cov_hc3(self)), self.HC3_se, 15)
print(self.HC0_se)
print(sw.se_cov(sw.cov_hac_simple(self, nlags=0, use_correction=False)))
# test White as HAC with nlags=0, same as nlags=1 ?
bse_hac0 = sw.se_cov(sw.cov_hac_simple(self, nlags=0, use_correction=False))
assert_almost_equal(bse_hac0, self.HC0_se, 15)
print(bse_hac0)
# test White as HAC with nlags=0, same as nlags=1 ?
bse_hac0c = sw.se_cov(sw.cov_hac_simple(self, nlags=0, use_correction=True))
assert_almost_equal(bse_hac0c, self.HC1_se, 15)

bse_w = sw.se_cov(sw.cov_white_simple(self, use_correction=False))
print(bse_w)
# test White
assert_almost_equal(bse_w, self.HC0_se, 15)

bse_wc = sw.se_cov(sw.cov_white_simple(self, use_correction=True))
print(bse_wc)
# test White
assert_almost_equal(bse_wc, self.HC1_se, 15)


groups = np.repeat(np.arange(5), 20)

idx = np.nonzero(np.diff(groups))[0].tolist()
groupidx = lzip([0] + idx, idx + [len(groups)])
ngroups = len(groupidx)

print(sw.se_cov(sw.cov_cluster(self, groups)))
# two strange looking corner cases BUG?
print(sw.se_cov(sw.cov_cluster(self, np.ones(len(endog), int), use_correction=False)))
print(sw.se_cov(sw.cov_crosssection_0(self, np.arange(len(endog)))))
# these results are close to simple (no group) white, 50 groups 2 obs each
groups = np.repeat(np.arange(50), 100 // 50)
print(sw.se_cov(sw.cov_cluster(self, groups)))
# 2 groups with 50 obs each, what was the interpretation again?
groups = np.repeat(np.arange(2), 100 // 2)
print(sw.se_cov(sw.cov_cluster(self, groups)))

"http://www.kellogg.northwestern.edu/faculty/petersen/htm/papers/se/test_data.txt"
"""
test <- read.table(
      url(paste("http://www.kellogg.northwestern.edu/",
            "faculty/petersen/htm/papers/se/",
            "test_data.txt",sep="")),
    col.names=c("firmid", "year", "x", "y"))
"""
