import numpy as np
import pandas as pd
from scipy.stats import multivariate_normal
from scipy.special import expit  # Sigmoid function
import statsmodels.api as sm
from sklearn.metrics import brier_score_loss
from sklearn.calibration import calibration_curve
import matplotlib.pyplot as plt


# -log P(yt|yp) = -(yt log(yp) + (1 - yt) log(1 - yp))

# Parameters
n_obs = 100   # Number of observations (within one cluster)
beta = np.array([1, 2, 3])  # Coefficients (Intercept, x1, x2)
rho = 0.5  # Correlation between responses

# Generate Covariates
np.random.seed(12038)
x1 = np.random.normal(size=n_obs)  # Continuous predictor
x2 = np.random.binomial(1, 0.5, size=n_obs)  # Binary predictor
X = np.column_stack((np.ones_like(x1), x1, x2))  # Design matrix with intercept

# Generate Correlated Errors
cov_matrix = rho * np.ones((n_obs, n_obs)) + (1 - rho) * np.eye(n_obs)  # Exchangeable correlation
correlated_errors = multivariate_normal.rvs(mean=np.zeros(n_obs), cov=cov_matrix)

# Compute Probabilities
eta = X @ beta  # Linear predictor
p = expit(eta)  # Convert to probability

# Generate Correlated Binary Outcomes using Thresholding
y_null = (p + correlated_errors > np.median(p + correlated_errors)).astype(int)

# Create DataFrame
df = pd.DataFrame({"x1": x1, "x2": x2, "y": y_null})

# Fit Models
logit_result_null = sm.Logit(y_null, sm.add_constant(df[["x1", "x2"]])).fit()
gee_result_null = sm.GEE(y_null, sm.add_constant(df[["x1", "x2"]]), groups=np.ones(len(df)), 
                         family=sm.families.Binomial(), cov_struct=sm.cov_struct.Exchangeable()).fit()

# Extract p-values
logit_pvalues = logit_result_null.pvalues
gee_pvalues = gee_result_null.pvalues

# Compare Results
print("Simple Logistic Regression Coefficients & Standard Errors:")
print(logit_result_null.summary())

print("\nGEE Coefficients & Robust Standard Errors:")
print(gee_result_null.summary())

# Calculate Brier Score
logit_probs = logit_result_null.predict()
gee_probs = gee_result_null.predict()
brier_score_logit = brier_score_loss(y_null, logit_probs)
brier_score_gee = brier_score_loss(y_null, gee_probs)
print("\nBrier Score for Logistic Regression:", brier_score_logit)
print("Brier Score for GEE:", brier_score_gee)

# Plot Calibration Curves
plt.figure(figsize=(7, 5))
prob_true_logit, prob_pred_logit = calibration_curve(df["y"], logit_probs, n_bins=10)
prob_true_gee, prob_pred_gee = calibration_curve(df["y"], gee_probs, n_bins=10)

plt.plot(prob_pred_logit, prob_true_logit, "s-", label="Logistic Regression")
plt.plot(prob_pred_gee, prob_true_gee, "o-", label="GEE")
plt.plot([0, 1], [0, 1], "k--")  # Perfect calibration
plt.xlabel("Predicted Probability")
plt.ylabel("True Probability")
plt.title("Calibration Curve")
plt.legend()
plt.show()