# Define X and Y sets, using train_test_split to save 25% for test and use 75% of the data to fit our model

# Initialize label encoder to be able to encode the HeartDisease column
label_encoder = LabelEncoder()

# Apply label encoder to the 'HeartDisease' column to ensure XGBoost handles it 
heart_2020['HeartDisease'] = label_encoder.fit_transform(heart_2020['HeartDisease'])

# Define our feature set
X = heart_2020.drop(["HeartDisease","Race"], axis = 1)
y = heart_2020["HeartDisease"]

# Use stratified splitting to ensure y has a representative % of the target class 
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, stratify = y)

# Model 1 - Dummify everything needed and scale numerical columns, All Predictors
ct_all = ColumnTransformer(
  [
    ("dummify", OneHotEncoder(sparse_output = False, handle_unknown='ignore'),
    make_column_selector(dtype_include=object)),
    ("standardize", StandardScaler(), make_column_selector(dtype_include=np.number))  #keep int64 as is
  ],
  remainder = "passthrough" 
).set_output(transform = "pandas")

# Use target class distribution to determine values to tune for pos_weight (which takes into account severe class imbalances)
target_class_distribution = heart_2020['HeartDisease'].value_counts() # 0 - 292422  
                                                                      # 1 - 27373  

# Define a grid of possible values that are close to 292,422 / 27,373 = 10.68
pos_weight_values = [10.6,10.7]
    
    
# Setup pipeline #1 for XGBoost
xgb_pipeline = Pipeline([
    ("preprocessor", ct_all),
    ("classifier", XGBClassifier(use_label_encoder=False, eval_metric='logloss'))  # Parameters needed to suppress warnings
])

# Parameter grid for XGBoost with different parameters being tuned
param_grid_xgb = {
    'classifier__n_estimators': [100],            # Number of boosted trees, reduced from 300 to 100 to run 
    'classifier__learning_rate': [0.1, 0.5],    # Boosting learning rate (xgb’s “eta”)
    'classifier__max_depth': [3,4],                  # Maximum depth of a tree (used 3,4 to have code run)
    'classifier__scale_pos_weight': pos_weight_values 
}

# Setup GridSearchCV for XGBoost                               # n_jobs -1 tells Jupyter to user all CPU's for gridsearch
grid_search_xgb = GridSearchCV(xgb_pipeline, param_grid_xgb, cv=5, scoring='precision', n_jobs=-1)

# Fitting our model on training data
grid_search_xgb.fit(X_train, y_train)

# Predict on the test set
test_xgb_predictions = grid_search_xgb.predict(X_test)

# Use precision function to compute how often our predictions match the actual Y's
test_xgb_precision = precision_score(y_test, test_xgb_predictions)

print("Test XGBoost Precision: ", test_xgb_precision)

# Define feature set

X = heart_2020.drop(["HeartDisease"], axis=1)
y = heart_2020["HeartDisease"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, stratify=y)

# Model 1 - Dummify everything needed and scale numerical columns, All Predictors
ct_all = ColumnTransformer(
    [("dummify", OneHotEncoder(sparse_output=False, handle_unknown='ignore'),
      make_column_selector(dtype_include=object)),
     ("standardize", StandardScaler(), make_column_selector(dtype_include=np.number))],
    remainder="passthrough"
)

# Setup pipeline for AdaBoost with decision tree as the base estimator
ada_pipeline = ImblearnPipeline([
    ("preprocessor", ct_all),
    ("smote", SMOTE(random_state=42)),
    ("classifier", AdaBoostClassifier(estimator=DecisionTreeClassifier(max_depth=1), random_state=2, algorithm='SAMME'))
])

# Parameter grid for AdaBoost with different parameters being tuned
param_grid_ada = {
    'classifier__n_estimators': [100],  # Number of boosted trees, set to 100 to run
    'classifier__learning_rate': [0.1]      # Boosting learning rate (set to 0.1 to run)
}

# Setup GridSearchCV for AdaBoost
grid_search_ada = GridSearchCV(ada_pipeline, param_grid_ada, cv=5, scoring='precision', n_jobs=-1)

# Fitting our model on training data
grid_search_ada.fit(X_train, y_train)

# Predict on the test set
test_ada_predictions = grid_search_ada.predict(X_test)

# Use precision function to compute how often our predictions match the actual Y's, which are coded as 1's
test_ada_precision = precision_score(y_test, test_ada_predictions, pos_label = 1)

print("Test AdaBoost Precision: ", test_ada_precision)

# Initialize label encoder to be able to encode the HeartDisease column
label_encoder = LabelEncoder()

# Apply label encoder to the 'HeartDisease' column to ensure XGBoost handles it 
heart_2022_final['HadHeartAttack'] = label_encoder.fit_transform(heart_2022_final['HadHeartAttack'])

# Define the new X test and Y sets, denoted as X_2 and y_2 to avoid confusion
X_2 = heart_2022_final.drop(["HadHeartAttack","Race"], axis=1)
y_2 = heart_2022_final["HadHeartAttack"]

# Use stratified splitting to ensure y has a representative % of the target class
X2_train, X2_test, y2_train, y2_test = train_test_split(X_2, y_2, test_size=0.3, stratify=y_2, random_state = 2)

# Model 1 - Dummify everything needed and scale numerical columns, All Predictors
ct_all = ColumnTransformer(
    [("dummify", OneHotEncoder(sparse_output=False, handle_unknown='ignore'),
      make_column_selector(dtype_include=object)),
     ("standardize", StandardScaler(), make_column_selector(dtype_include=np.number))],
    remainder="passthrough"
)

# Setup pipeline for AdaBoost with decision tree as the base estimator
ada_pipeline = ImblearnPipeline([
    ("preprocessor", ct_all),
    ("smote", SMOTE(random_state=42)),
    ("classifier", AdaBoostClassifier(estimator=DecisionTreeClassifier(max_depth=1), random_state=2, algorithm='SAMME'))
])

# Parameter grid for AdaBoost with different parameters being tuned
param_grid_ada = {
    'classifier__n_estimators': [100],  # Number of boosted trees, set to 100 to run
    'classifier__learning_rate': [0.1]      # Boosting learning rate (set to 0.1 to run)
}

# Setup GridSearchCV for AdaBoost
grid_search_ada = GridSearchCV(ada_pipeline, param_grid_ada, cv=5, scoring='precision', n_jobs=-1)

# Fitting our model on training data
grid_search_ada.fit(X2_train, y2_train)

# Predict on the test set
test_ada_predictions = grid_search_ada.predict(X2_test)

# Use precision function to compute how often our predictions match the actual Y's
test_ada_precision = precision_score(y2_test, test_ada_predictions, pos_label = 1)

print("Test AdaBoost 2022 Precision: ", test_ada_precision)