# Experiment 4: Logistic Regression for Heart Disease Prediction
# Requirements: pip install pandas matplotlib scikit-learn

# Heart Disease Prediction using Logistic Regression

import pandas as pd
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score,
    f1_score, roc_auc_score, confusion_matrix,
    ConfusionMatrixDisplay, RocCurveDisplay
)

# 1. Load the dataset
# Replace the file name with the location of your dataset.
# CHANGE THIS PATH TO YOUR DATASET
data = pd.read_csv("heart_disease.csv")

# 2. Select input features and target variable
features = [
    "Age",
    "Resting_Blood_Pressure",
    "Cholesterol",
    "Maximum_Heart_Rate",
    "Exercise_Induced_Angina"
]

X = data[features]
y = data["Heart_Disease"]   # 0 = No Heart Disease, 1 = Heart Disease

# 3. Split data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.20,
    random_state=42,
    stratify=y
)

# 4. Standardize numerical features
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 5. Train Logistic Regression model
model = LogisticRegression(random_state=42)
model.fit(X_train_scaled, y_train)

# 6. Predict class labels and probabilities
y_pred = model.predict(X_test_scaled)
y_prob = model.predict_proba(X_test_scaled)[:, 1]

# 7. Calculate evaluation metrics
accuracy = accuracy_score(y_test, y_pred)
precision = precision_score(y_test, y_pred, zero_division=0)
recall = recall_score(y_test, y_pred, zero_division=0)
f1 = f1_score(y_test, y_pred, zero_division=0)
roc_auc = roc_auc_score(y_test, y_prob)

print("Model Performance")
print("-----------------")
print("Accuracy :", round(accuracy, 4))
print("Precision:", round(precision, 4))
print("Recall   :", round(recall, 4))
print("F1-Score :", round(f1, 4))
print("ROC-AUC  :", round(roc_auc, 4))

# 8. Display confusion matrix
cm = confusion_matrix(y_test, y_pred)
print("\nConfusion Matrix:\n", cm)

ConfusionMatrixDisplay(confusion_matrix=cm).plot()
plt.title("Confusion Matrix - Heart Disease Prediction")
plt.show()

# 9. Plot ROC curve
RocCurveDisplay.from_predictions(y_test, y_prob)
plt.title("ROC Curve - Heart Disease Prediction")
plt.show()

# 10. Predict for a new patient
new_patient = pd.DataFrame([{
    "Age": 55,
    "Resting_Blood_Pressure": 140,
    "Cholesterol": 240,
    "Maximum_Heart_Rate": 150,
    "Exercise_Induced_Angina": 1
}])

new_patient_scaled = scaler.transform(new_patient)
prediction = model.predict(new_patient_scaled)[0]
probability = model.predict_proba(new_patient_scaled)[0, 1]

print("\nNew Patient Prediction")
print("----------------------")
print("Probability of Heart Disease:", round(probability, 4))

if prediction == 1:
    print("Prediction: Heart Disease Present")
else:
    print("Prediction: No Heart Disease")
