# ---------------------------------------------------------------------------
# safe_water.py
#
# Capstone project: predicting, from a quick field-kit test, whether a
# drinking-water source will fail the laboratory's safety test - so unsafe
# sources can be flagged the same day instead of two days later.
#
# CBSE Class XII, Artificial Intelligence (843), Part C - Capstone Project.
# SDG 6 (Clean Water and Sanitation) and SDG 3 (Good Health and Well-being).
#
# Data Science Methodology, 10 steps. CLASSIFICATION (Safe / Unsafe). The
# field data has gaps, so the project compares four ways of handling
# missing values - the Unit 1 topic - with stratified 5-fold cross
# validation, before training the final Random Forest.
#
# Libraries: pandas, numpy, matplotlib, scikit-learn, joblib.
# ---------------------------------------------------------------------------

# %% Step 1 - Business Understanding
print("=" * 64)
print("IS THIS WATER SAFE TO DRINK?")
print("=" * 64)
print("Problem : the lab's verdict takes two days; people keep drinking meanwhile")
print("Goal    : predict the lab verdict from field-kit readings, the same day")
print()

# %% Step 2 - Analytic Approach
print("Analytic approach : predictive - classification (Safe or Unsafe)")
print()

# %% Step 3 - Data Requirements
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.impute import SimpleImputer
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, recall_score, f1_score
import joblib

pd.set_option("display.width", 140)
pd.set_option("display.max_columns", 20)

# %% Step 4 - Data Collection
data = pd.read_csv("water_samples.csv")
print("Samples collected :", len(data))
print(data.head(6))
print()

# %% Step 5 - Data Understanding
print("--- Step 5: Data Understanding ---")
missing = data.isnull().sum()
print("Missing values in each column:")
print(missing.to_string())
print()
print("Lab verdicts:")
print(data["verdict"].value_counts().to_string())
print("Share unsafe :", round((data["verdict"] == "Unsafe").mean() * 100, 1), "%")
print()
print("Share unsafe by kind of source (%):")
by_source = ((data["verdict"] == "Unsafe").groupby(data["source"]).mean() * 100).sort_values()
print(by_source.round(1).to_string())
print()

# Chart 1: how much is missing in each column.
m = (missing[missing > 0] / len(data) * 100).sort_values()
plt.figure(figsize=(7, 3.8))
plt.barh(m.index, m.values, color="#f28e2b")
plt.title("Missing values in the field data")
plt.xlabel("% of samples with no reading")
plt.tight_layout()
plt.savefig("chart1_missing.png")
plt.close()

# Chart 2: unsafe share by kind of source.
plt.figure(figsize=(7, 4))
plt.barh(by_source.index, by_source.values, color="#4e9fd1")
plt.title("Share of samples the lab found unsafe")
plt.xlabel("% unsafe")
plt.tight_layout()
plt.savefig("chart2_by_source.png")
plt.close()

# Chart 3: turbidity for safe and unsafe samples (log scale).
turb = pd.to_numeric(data["turbidity_ntu"])
plt.figure(figsize=(7, 4.2))
plt.boxplot([turb[data["verdict"] == "Safe"].dropna(), turb[data["verdict"] == "Unsafe"].dropna()],
            labels=["Safe", "Unsafe"])
plt.yscale("log")
plt.axhline(5, color="#e15759", linestyle="--", label="IS 10500 limit: 5 NTU")
plt.title("Turbidity of safe and unsafe samples")
plt.ylabel("Turbidity (NTU, log scale)")
plt.legend()
plt.tight_layout()
plt.savefig("chart3_turbidity.png")
plt.close()

# %% Step 6 - Data Preparation
print("--- Step 6: Data Preparation ---")
before = len(data)
data = data.drop_duplicates()
print("Duplicate samples removed :", before - len(data))
for col in ["fluoride_mg_l", "nitrate_mg_l", "turbidity_ntu"]:
    data[col] = pd.to_numeric(data[col])
data["unsafe"] = (data["verdict"] == "Unsafe").astype(int)
data = pd.get_dummies(data, columns=["source"], dtype=int)
features = ["ph", "tds_mg_l", "turbidity_ntu", "nitrate_mg_l", "fluoride_mg_l", "chloride_mg_l",
            "residual_chlorine_mg_l", "source_Hand pump", "source_Open well", "source_Piped supply",
            "source_Pond", "source_Tube well"]
X = data[features]
y = data["unsafe"]

# Four ways of handling the missing values, each judged by stratified
# 5-fold cross validation. The filling happens INSIDE each pipeline, so it
# is learned from the training folds only and never sees the test fold.
folds = StratifiedKFold(n_splits=5, shuffle=True, random_state=2)
forest = RandomForestClassifier(n_estimators=300, min_samples_leaf=2, random_state=0)
complete = data.dropna(subset=features)
strategies = [
    ["A. Drop rows with a gap", None],
    ["B. Fill with the mean", make_pipeline(SimpleImputer(strategy="mean"), forest)],
    ["C. Fill with the median", make_pipeline(SimpleImputer(strategy="median"), forest)],
    ["D. Median + a 'was missing' column", make_pipeline(SimpleImputer(strategy="median", add_indicator=True), forest)],
]
print("Rows left if every row with a gap is dropped :", len(complete), "of", len(data))
recalls = []
for name, model in strategies:
    if model is None:
        r = cross_val_score(forest, complete[features], complete["unsafe"], cv=folds, scoring="recall")
        f = cross_val_score(forest, complete[features], complete["unsafe"], cv=folds, scoring="f1")
    else:
        r = cross_val_score(model, X, y, cv=folds, scoring="recall")
        f = cross_val_score(model, X, y, cv=folds, scoring="f1")
    recalls.append([name, r.mean(), f.mean()])
    print(name.ljust(36), "CV recall", round(r.mean(), 3), "  CV F1", round(f.mean(), 3))
print()

# Chart 4: the four strategies compared.
x = np.arange(4)
plt.figure(figsize=(8, 4.3))
plt.bar(x - 0.2, [r[1] for r in recalls], width=0.4, label="Recall (unsafe found)", color="#e15759")
plt.bar(x + 0.2, [r[2] for r in recalls], width=0.4, label="F1 score", color="#4e79a7")
plt.xticks(x, ["A. Drop", "B. Mean", "C. Median", "D. Median +\nindicator"])
plt.ylim(0.5, 1)
plt.ylabel("Mean over 5 folds")
plt.title("Four ways of handling missing values")
plt.legend()
plt.tight_layout()
plt.savefig("chart4_strategies.png")
plt.close()

# %% Step 7 - Modelling
print("--- Step 7: Modelling ---")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=5)
# Strategy A scored best, but only on the 797 complete rows, and a deployed
# model cannot refuse every sample with a missing test. B, C and D tied, and
# the indicator columns added nothing, so the simplest - C, the median - is used.
model = make_pipeline(SimpleImputer(strategy="median"),
                      RandomForestClassifier(n_estimators=300, min_samples_leaf=2, random_state=0))
model.fit(X_train, y_train)
print("Strategy used          : C (fill with the median)")
print("Training samples       :", len(X_train), " Testing samples :", len(X_test))
print()

# %% Step 8 - Evaluation
print("--- Step 8: Evaluation ---")
pred = model.predict(X_test)
tn, fp, fn, tp = confusion_matrix(y_test, pred).ravel()
print("TP, FN, FP, TN :", tp, fn, fp, tn)
print("Accuracy       :", round(accuracy_score(y_test, pred), 3))
print("Precision      :", round(precision_score(y_test, pred), 3))
print("Recall         :", round(recall_score(y_test, pred), 3))
print("F1 score       :", round(f1_score(y_test, pred), 3))
print()
cv_recall = cross_val_score(model, X, y, cv=folds, scoring="recall")
print("Recall in each of the 5 folds :", np.round(cv_recall, 3))
print()

# A missed unsafe source is far worse than an extra lab test, so the
# threshold is lowered from 0.5 to 0.3.
threshold = 0.3
prob_test = model.predict_proba(X_test)[:, 1]
flag = (prob_test >= threshold).astype(int)
print("Threshold for 'unsafe'        :", threshold)
print("Recall at threshold 0.3       :", round(recall_score(y_test, flag), 3))
print("Precision at threshold 0.3    :", round(precision_score(y_test, flag), 3))
print("Unsafe sources missed at 0.3  :", int(((flag == 0) & (y_test == 1)).sum()), "of", int(y_test.sum()))
print()

# Chart 5: confusion matrix.
cm = confusion_matrix(y_test, pred, labels=[1, 0])
plt.figure(figsize=(5.5, 4.5))
plt.imshow(cm, cmap="Blues")
labels = [["TP", "FN"], ["FP", "TN"]]
for r in range(2):
    for c in range(2):
        plt.text(c, r, labels[r][c] + "\n" + str(cm[r][c]), ha="center", va="center", fontsize=13,
                 color="white" if cm[r][c] > cm.max() / 2 else "black")
plt.xticks([0, 1], ["Predicted unsafe", "Predicted safe"])
plt.yticks([0, 1], ["Unsafe", "Safe"])
plt.title("Confusion matrix (test samples)")
plt.tight_layout()
plt.savefig("chart5_confusion.png")
plt.close()

# Chart 6: which readings the forest relied on.
importance = pd.Series(model[-1].feature_importances_, index=features).sort_values()
print("Feature importance:")
print(importance.round(3).to_string())
print()
plt.figure(figsize=(7.5, 5.2))
plt.barh(importance.index, importance.values, color="#59a14f")
plt.title("What the Random Forest relied on")
plt.xlabel("Importance")
plt.tight_layout()
plt.savefig("chart6_importance.png")
plt.close()

# %% Step 9 - Deployment
print("--- Step 9: Deployment ---")
joblib.dump(model, "water_model.pkl")
model = joblib.load("water_model.pkl")
new = pd.DataFrame({
    "ph": [7.2, 7.8, 6.9], "tds_mg_l": [380, 1250, 520], "turbidity_ntu": [0.4, 2.1, 9.5],
    "nitrate_mg_l": [8.0, 30.0, np.nan], "fluoride_mg_l": [0.3, np.nan, 0.6], "chloride_mg_l": [110, 380, 150],
    "residual_chlorine_mg_l": [0.4, 0.0, 0.0], "source_Hand pump": [0, 1, 0], "source_Open well": [0, 0, 1],
    "source_Piped supply": [1, 0, 0], "source_Pond": [0, 0, 0], "source_Tube well": [0, 0, 0],
})
names = ["Piped supply, chlorinated", "Hand pump, fluoride not tested", "Open well, cloudy, nitrate not tested"]
prob = model.predict_proba(new)[:, 1]
for i in range(len(names)):
    if prob[i] >= threshold:
        verdict = "LIKELY UNSAFE - warn users, send to the lab first"
    else:
        verdict = "likely safe - confirm with the lab test"
    print(names[i].ljust(38), round(prob[i], 2), "->", verdict)
print()

# %% Step 10 - Feedback
print("--- Step 10: Feedback ---")
# After the monsoon, 250 more sources are tested. Rain washes dirt and
# sewage into wells and ponds, so contamination rises.
mon = pd.read_csv("monsoon_samples.csv")
mon["unsafe"] = (mon["verdict"] == "Unsafe").astype(int)
mon = pd.get_dummies(mon, columns=["source"], dtype=int)
mon_pred = (model.predict_proba(mon[features])[:, 1] >= threshold).astype(int)
print("Monsoon samples                    :", len(mon))
print("Share unsafe, before the monsoon   :", round(y.mean() * 100, 1), "%")
print("Share unsafe, after the monsoon    :", round(mon["unsafe"].mean() * 100, 1), "%")
print("Recall on the test samples (0.3)   :", round(recall_score(y_test, flag), 3))
print("Recall on the monsoon samples      :", round(recall_score(mon["unsafe"], mon_pred), 3))
print("Precision on the monsoon samples   :", round(precision_score(mon["unsafe"], mon_pred), 3))
if recall_score(mon["unsafe"], mon_pred) < recall_score(y_test, flag) - 0.05:
    print("Feedback: recall has fallen after the monsoon - add monsoon samples and retrain.")
else:
    print("Feedback: recall holds after the monsoon - the model can stay in use.")