import pandas as pd
import joblib

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.tree import DecisionTreeClassifier, export_text

from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix
)

# =====================================
# LOAD DATASET
# =====================================

df = pd.read_csv("diabetes_dataset.csv")

print("\nDataset Shape:")
print(df.shape)

# =====================================
# RISK DISTRIBUTION
# =====================================

print("\nRisk Distribution")
print(df["Risk_Level"].value_counts())

print("\nRisk Percentage")
print(
    round(
        df["Risk_Level"].value_counts(normalize=True) * 100,
        2
    )
)

# =====================================
# ENCODE CATEGORICAL VARIABLES
# =====================================

gender_encoder = LabelEncoder()
df["Gender"] = gender_encoder.fit_transform(df["Gender"])

family_encoder = LabelEncoder()
df["Family_History"] = family_encoder.fit_transform(
    df["Family_History"]
)

risk_encoder = LabelEncoder()
df["Risk_Level"] = risk_encoder.fit_transform(
    df["Risk_Level"]
)

# =====================================
# FEATURES
# =====================================

features = [
    "Age",
    "Gender",
    "BMI",
    "Glucose_Level",
    "Blood_Pressure",
    "Insulin",
    "Physical_Activity",
    "Family_History"
]

X = df[features]

y = df["Risk_Level"]

# =====================================
# FEATURE CORRELATION
# =====================================

print("\n===================================================")
print("FEATURE CORRELATION WITH RISK LEVEL")
print("===================================================")

corr_df = X.copy()
corr_df["Risk_Level"] = y

print(
    corr_df.corr()["Risk_Level"]
    .sort_values(ascending=False)
)

# =====================================
# TRAIN TEST SPLIT
# =====================================

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.20,
    random_state=42,
    stratify=y
)

# =====================================
# MODEL
# =====================================

model = DecisionTreeClassifier(
    criterion="gini",
    max_depth=5,
    min_samples_split=15,
    min_samples_leaf=8,
    class_weight="balanced",
    random_state=42
)

# =====================================
# TRAIN MODEL
# =====================================

model.fit(X_train, y_train)

# =====================================
# PREDICTIONS
# =====================================

predictions = model.predict(X_test)

# =====================================
# ACCURACY
# =====================================

print("\n===================================================")
print("MODEL PERFORMANCE")
print("===================================================")

print("\nAccuracy:")
print(
    round(
        accuracy_score(y_test, predictions),
        4
    )
)

print("\nClassification Report:")
print(
    classification_report(
        y_test,
        predictions,
        target_names=risk_encoder.classes_
    )
)

print("\nConfusion Matrix:")
print(
    confusion_matrix(
        y_test,
        predictions
    )
)

# =====================================
# FEATURE IMPORTANCE
# =====================================

importance = pd.DataFrame({
    "Feature": features,
    "Importance": model.feature_importances_
})

importance = (
    importance
    .sort_values(
        by="Importance",
        ascending=False
    )
    .reset_index(drop=True)
)

print("\n===================================================")
print("FEATURE IMPORTANCE")
print("===================================================")

print(importance)

# =====================================
# RISK LABEL MAPPING
# =====================================

print("\n===================================================")
print("RISK LABEL MAPPING")
print("===================================================")

for index, label in enumerate(risk_encoder.classes_):
    print(f"{index} = {label}")

# =====================================
# DECISION TREE RULES
# =====================================

print("\n===================================================")
print("DECISION TREE RULES")
print("===================================================")

rules = export_text(
    model,
    feature_names=features
)

print(rules)

# =====================================
# PROBABILITY TESTING
# =====================================

print("\n===================================================")
print("SAMPLE PREDICTION PROBABILITIES")
print("===================================================")

probabilities = model.predict_proba(
    X_test.head(5)
)

print(probabilities)

# =====================================
# SAVE EVERYTHING
# =====================================

model_package = {
    "model": model,
    "gender_encoder": gender_encoder,
    "family_encoder": family_encoder,
    "risk_encoder": risk_encoder,
    "features": features
}

joblib.dump(
    model_package,
    "diabetes_risk_model.pkl"
)

print("\n===================================================")
print("MODEL SAVED SUCCESSFULLY")
print("Saved File: diabetes_risk_model.pkl")
print("===================================================")