import numpy as np
import pandas as pd
from scipy.stats import multivariate_normal

# from scipy.special import expit
import statsmodels.api as sm
import matplotlib.pyplot as plt

def plot_dataset(X, y, axes):
    plt.plot(X[:, 0][y==0], X[:, 1][y==0], "bs")
    plt.plot(X[:, 0][y==1], X[:, 1][y==1], "g^")
    plt.axis(axes)
    plt.grid(True, which='major', linestyle='--', linewidth='0.5')
    plt.xlabel(r"$x_1$", fontsize=20)
    plt.ylabel(r"$x_2$", fontsize=20, rotation=0)


# Parameters
n_obs = 1000   # Total observations
rho = 0.5  # Intracluster correlation (exchangeable structure)

# Generate Two Clouds of Points
np.random.seed(10039)

first_half = n_obs // 2

x1 = np.random.normal(loc=1 , scale=1, size=(first_half, 2))  # Spread out around different centers
x2 = np.random.normal(loc=-1, scale=1, size=(n_obs - first_half, 2))  # Spread out around different centers
X = np.concatenate((x1, x2))  # uncorrelated features


vcv = rho * np.ones((n_obs, n_obs)) + (1 - rho) * np.eye(n_obs)  # Exchangeable correlation
E,Q = np.linalg.eig( vcv )
L = Q @ np.diag( E**(1/2) )

X = L @ X  # Now observations are correlated
X = np.real(X)  # Ensure no complex numbers
b = 0
eta = X[:,0]*X[:,1] + b  # Linear predictor
p = 1/(1 + np.exp(-(eta + b)))  # Linear predictor
correlated_errors = multivariate_normal.rvs(mean=np.zeros(n_obs), cov=vcv)

y = p + correlated_errors >= 0.5   # Binary outcome based on x1 and x2
# y = p >= 0.5   # Binary outcome based on x1 and x2
# plot_dataset(X, y, np.array([-1, 1, -1, 1])*5)

# Create DataFrame
df = pd.DataFrame({"x1": X[:,0], "x2": X[:,1], "y": y})

# Fit GEE Model
model = sm.GEE(df["y"], sm.add_constant(df[["x1", "x2"]]), groups=np.ones(n_obs), 
               family=sm.families.Binomial(), cov_struct=sm.cov_struct.Exchangeable())
result = model.fit()

# Display Results
# print(result.summary())

# get pearson residuals
gee_err = np.mean(result.resid_pearson)
print("Pearson residuals: ", gee_err)

y_prob_gee = model.predict(result.params, exog=sm.add_constant(df[["x1", "x2"]]))

y_pred_gee = (y_prob_gee >= .5).astype(float)
loss_gee = np.mean(y_pred_gee != df["y"])  # Accuracy
print('loss GEE: ', loss_gee)



from sklearn.linear_model import LogisticRegression
log_reg = LogisticRegression()
X_l = np.concatenate(( np.ones((X.shape[0], 1)), X), axis=1)  # Add intercept

log_reg.fit(X_l, y)
log_reg_loss = 1 - log_reg.score(X_l, y)
print('log_reg_loss: ', log_reg_loss)
log_reg_probs = log_reg.predict_proba(X_l)[:, 1]  # Get probabilities

# make auc curve
from sklearn.metrics import roc_curve, auc

log_reg_fpr, log_reg_tpr, _ = roc_curve(y, log_reg_probs)
gee_fpr, gee_tpr, _ = roc_curve(y, y_prob_gee)

log_reg_auc = auc(log_reg_fpr, log_reg_tpr)
gee_auc = auc(gee_fpr, gee_tpr)

print('log_reg_auc: ', log_reg_auc)
print('gee_auc: ', gee_auc)

# plt.figure(figsize=(8, 6))
# plt.plot(log_reg_fpr, log_reg_tpr, 'b-', label='Logistic Regression (area = %0.4f)' % log_reg_auc)
# plt.plot(gee_fpr, gee_tpr, 'r-', label='Logistic KRGEE  (area = %0.4f)' % gee_auc)
# plt.plot([0, 1], [0, 1], 'k--')
# plt.xlabel('False Positive Rate')
# plt.ylabel('True Positive Rate')
# plt.legend()
# plt.grid(True)
# plt.tight_layout()
