import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report
import statsmodels.api as sm
from statsmodels.formula.api import glm
from scipy import stats

# Create the dataset from Table 2.21
data = {
    'patient_id': list(range(1, 51)),
    'X1': [2.5, 173.0, 119.0, 10.0, 502.2, 4.0, 14.4, 2.0, 40.0, 6.6, 
           21.4, 2.8, 2.5, 6.0, 3.5, 1.2, 3.5, 39.7, 62.4, 2.4, 
           34.7, 28.4, 0.9, 30.6, 5.8, 6.1, 2.7, 4.7, 128.0, 35.0, 
           62.2, 10.8, 21.6, 2.0, 3.4, 5.1, 2.4, 1.7, 1.1, 12.8, 
           2.0, 8.5, 2.0, 2.0, 4.3, 244.8, 4.0, 5.1, 32.0, 1.4],
    'X2': [0, 2, 2, 2, 2, 0, 0, 2, 2, 0, 
           2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 
           0, 2, 0, 2, 0, 0, 2, 0, 2, 0, 
           0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 
           0, 0, 2, 0, 0, 2, 0, 0, 0, 0],
    'X3': [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 
           1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 
           0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 
           0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 
           0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    'Y': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 
          1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
}

# Create a pandas DataFrame
df = pd.DataFrame(data)

print("Dataset shape:", df.shape)
print("\nData summary:")
print(df.describe())

# 1. Full model with all variables - using statsmodels for detailed statistics
print("\n=== FULL MODEL (X1, X2, X3) ===")
formula_full = 'Y ~ X1 + X2 + X3'
model_full = glm(formula=formula_full, data=df, family=sm.families.Binomial()).fit()
print(model_full.summary())

# Global significance of the model - manually calculate LR test p-value
ll_full = model_full.llf
ll_null = model_full.llnull
lr_stat = 2 * (ll_full - ll_null)
# Degrees of freedom = number of predictors
df_diff = 3
lr_pvalue = stats.chi2.sf(lr_stat, df_diff)

print("\nModel Log-Likelihood:", ll_full)
print("Null Log-Likelihood:", ll_null)
print("LR test statistic:", lr_stat)
print("LR p-value:", lr_pvalue)

# Calculate odds ratios for the full model
print("\nOdds Ratios for Full Model:")
odds_ratios = np.exp(model_full.params)
conf_intervals = np.exp(model_full.conf_int())
odds_ratio_df = pd.DataFrame({
    'OR': odds_ratios,
    'Lower CI': conf_intervals.iloc[:, 0],
    'Upper CI': conf_intervals.iloc[:, 1]
})
print(odds_ratio_df)

# Make predictions with the full model
df['pred_prob_full'] = model_full.predict()
df['pred_class_full'] = (df['pred_prob_full'] >= 0.5).astype(int)
accuracy_full = accuracy_score(df['Y'], df['pred_class_full'])
print(f"\nAccuracy of full model: {accuracy_full:.4f}")
print("\nClassification Report for Full Model:")
print(classification_report(df['Y'], df['pred_class_full']))

# 2. Create reduced models by removing one variable at a time
# Model without X1
print("\n=== MODEL WITHOUT X1 (ONLY X2, X3) ===")
formula_no_x1 = 'Y ~ X2 + X3'
model_no_x1 = glm(formula=formula_no_x1, data=df, family=sm.families.Binomial()).fit()
print(model_no_x1.summary())

# Model without X2
print("\n=== MODEL WITHOUT X2 (ONLY X1, X3) ===")
formula_no_x2 = 'Y ~ X1 + X3'
model_no_x2 = glm(formula=formula_no_x2, data=df, family=sm.families.Binomial()).fit()
print(model_no_x2.summary())

# Model without X3
print("\n=== MODEL WITHOUT X3 (ONLY X1, X2) ===")
formula_no_x3 = 'Y ~ X1 + X2'
model_no_x3 = glm(formula=formula_no_x3, data=df, family=sm.families.Binomial()).fit()
print(model_no_x3.summary())

# Predictions and accuracy for each reduced model
models = {
    'Full Model': model_full,
    'No X1': model_no_x1,
    'No X2': model_no_x2, 
    'No X3': model_no_x3
}

# Compare model accuracies
print("\n=== MODEL COMPARISON ===")
accuracies = {}
for name, model in models.items():
    predictions = (model.predict() >= 0.5).astype(int)
    accuracy = accuracy_score(df['Y'], predictions)
    accuracies[name] = accuracy
    print(f"Accuracy of {name}: {accuracy:.4f}")

# Create a bar chart to compare accuracies
plt.figure(figsize=(10, 6))
plt.bar(accuracies.keys(), accuracies.values())
plt.ylabel('Accuracy')
plt.title('Model Accuracy Comparison')
plt.ylim(0, 1)
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.savefig('model_comparison.png')

# Find best model based on AIC and BIC
print("\nModel selection criteria:")
for name, model in models.items():
    print(f"{name}: AIC = {model.aic:.2f}, BIC = {model.bic:.2f}")

# Visualize the predicted probabilities from the full model
plt.figure(figsize=(10, 6))
sns.histplot(df[df['Y'] == 0]['pred_prob_full'], color='blue', alpha=0.5, label='Y=0', bins=10)
sns.histplot(df[df['Y'] == 1]['pred_prob_full'], color='orange', alpha=0.5, label='Y=1', bins=10)
plt.axvline(x=0.5, color='r', linestyle='--', label='Decision boundary')
plt.xlabel('Predicted Probability')
plt.ylabel('Frequency')
plt.title('Distribution of Predicted Probabilities (Full Model)')
plt.legend()
plt.savefig('predicted_probabilities.png')

# Simple visualization of the relationship between X3 and Y
plt.figure(figsize=(8, 5))
cross_tab = pd.crosstab(df['X3'], df['Y'])
cross_tab_pct = cross_tab.div(cross_tab.sum(axis=1), axis=0)
cross_tab_pct.plot(kind='bar', stacked=True)
plt.xlabel('Consolidation Therapy (X3)')
plt.ylabel('Proportion')
plt.title('Relationship between Consolidation Therapy and Survival')
plt.legend(title='Survival', labels=['<1 year', '≥1 year'])
plt.tight_layout()
plt.savefig('x3_y_relationship.png')

# Test significance of removing variables - Likelihood Ratio Tests
print("\n=== TESTING SIGNIFICANCE OF INDIVIDUAL VARIABLES ===")
# Test removing X1
lr_stat_x1 = 2 * (model_full.llf - model_no_x1.llf)
p_value_x1 = stats.chi2.sf(lr_stat_x1, 1)  # df=1 for one variable difference
print(f"Removing X1: LR stat = {lr_stat_x1:.4f}, p-value = {p_value_x1:.4f}")

# Test removing X2
lr_stat_x2 = 2 * (model_full.llf - model_no_x2.llf)
p_value_x2 = stats.chi2.sf(lr_stat_x2, 1)
print(f"Removing X2: LR stat = {lr_stat_x2:.4f}, p-value = {p_value_x2:.4f}")

# Test removing X3
lr_stat_x3 = 2 * (model_full.llf - model_no_x3.llf)
p_value_x3 = stats.chi2.sf(lr_stat_x3, 1)
print(f"Removing X3: LR stat = {lr_stat_x3:.4f}, p-value = {p_value_x3:.4f}")

print("\nAnalysis complete!")