import numpy as np
import pandas as pd
import statsmodels.api as sm

"""
Explanation: GEE nested covariance structure simulation study
This notebook is a simulation study that illustrates and evaluates the performance of the GEE nested covariance structure.
A nested covariance structure is based on a nested sequence of groups, or "levels".  The top level in the hierarchy is defined by the groups argument to GEE.  Subsequent levels are defined by the dep_data argument to GEE.
End of explanation
"""


p = 5

"""
Explanation: Set the number of covariates.
End of explanation
"""


groups_var = 1
level1_var = 2
level2_var = 3
resid_var = 4

"""
Explanation: These parameters define the population variance for each level of grouping.
End of explanation
"""


n_groups = 100

"""
Explanation: Set the number of groups
End of explanation
"""


group_size = 20
level1_size = 10
level2_size = 5

"""
Explanation: Set the number of observations at each level of grouping.  Here, everything is balanced, i.e. within a level every group has the same size.
End of explanation
"""


n = n_groups * group_size * level1_size * level2_size

"""
Explanation: Calculate the total sample size.
End of explanation
"""


xmat = np.random.normal(size=(n, p))

"""
Explanation: Construct the design matrix.
End of explanation
"""


groups_ix = np.kron(np.arange(n // group_size), np.ones(group_size)).astype(int)
level1_ix = np.kron(np.arange(n // level1_size), np.ones(level1_size)).astype(int)
level2_ix = np.kron(np.arange(n // level2_size), np.ones(level2_size)).astype(int)

"""
Explanation: Construct labels showing which group each observation belongs to at each level.
End of explanation
"""


groups_re = np.sqrt(groups_var) * np.random.normal(size=n // group_size)
level1_re = np.sqrt(level1_var) * np.random.normal(size=n // level1_size)
level2_re = np.sqrt(level2_var) * np.random.normal(size=n // level2_size)

"""
Explanation: Simulate the random effects.
End of explanation
"""


y = groups_re[groups_ix] + level1_re[level1_ix] + level2_re[level2_ix]
y += np.sqrt(resid_var) * np.random.normal(size=n)

"""
Explanation: Simulate the response variable.
End of explanation
"""


df = pd.DataFrame(xmat, columns=["x%d" % j for j in range(p)])
df["y"] = y + xmat[:, 0] - xmat[:, 3]
df["groups_ix"] = groups_ix
df["level1_ix"] = level1_ix
df["level2_ix"] = level2_ix

"""
Explanation: Put everything into a dataframe.
End of explanation
"""


cs = sm.cov_struct.Nested()
dep_fml = "0 + level1_ix + level2_ix"
m = sm.GEE.from_formula(
    "y ~ x0 + x1 + x2 + x3 + x4",
    cov_struct=cs,
    dep_data=dep_fml,
    groups="groups_ix",
    data=df,
)
r = m.fit()

"""
Explanation: Fit the model.
End of explanation
"""


r.cov_struct.summary()

"""
Explanation: The estimated covariance parameters should be similar to groups_var, level1_var, etc. as defined above.
End of explanation
"""


