# ---------------------------------------------------------------------------
# flood_warning.py
#
# Capstone project: a next-day flood warning for a river town, from today's
# rain, upstream rain, river level, soil wetness and dam releases.
#
# CBSE Class XII, Artificial Intelligence (843), Part C - Capstone Project.
# SDG 13 (Climate Action) and SDG 11 (Sustainable Cities and Communities).
#
# Data Science Methodology, 10 steps. CLASSIFICATION of a RARE event: floods
# happen on a few days a year, so accuracy is useless and the project is
# judged on recall, precision and F1. Two models - Logistic Regression and a
# Random Forest, both told to weight the rare class - are tested on the two
# latest monsoons, and checked with leave-one-year-out cross validation.
#
# Libraries: pandas, numpy, matplotlib, scikit-learn, joblib.
# ---------------------------------------------------------------------------

# %% Step 1 - Business Understanding
print("=" * 64)
print("A FLOOD EARLY-WARNING MODEL")
print("=" * 64)
print("Problem : the river crosses its danger level with little warning")
print("Goal    : warn a day ahead, so people and cattle can move in time")
print()

# %% Step 2 - Analytic Approach
print("Analytic approach : predictive - classification of a rare event")
print()

# %% Step 3 - Data Requirements
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import LeaveOneGroupOut
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score, accuracy_score
import joblib

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

# %% Step 4 - Data Collection
data = pd.read_csv("river_gauge.csv", parse_dates=["date"])
print("Days collected :", len(data))
print("Monsoons       :", data["date"].dt.year.min(), "to", data["date"].dt.year.max())
print(data.head(6))
print()

# %% Step 5 - Data Understanding
print("--- Step 5: Data Understanding ---")
print("Missing values:")
print(data.isnull().sum().to_string())
print()
floods = data["flood_next_day"].sum()
print("Days followed by a flood :", floods, "of", len(data))
print("Share of flood days      :", round(data["flood_next_day"].mean() * 100, 1), "%")
print("Accuracy of a model that NEVER warns :", round((1 - data["flood_next_day"].mean()) * 100, 1), "%")
print()

# Chart 1: one monsoon, with the danger level.
one = data[data["date"].dt.year == 2022]
plt.figure(figsize=(10, 4))
plt.plot(one["date"], one["river_level_m"], color="#4e79a7")
plt.axhline(DANGER, color="#e15759", linestyle="--", label="Danger level (104 m)")
plt.title("River level through the 2022 monsoon")
plt.ylabel("River level (m)")
plt.legend()
plt.grid(True, linestyle="--", alpha=0.3)
plt.tight_layout()
plt.savefig("chart1_season_2022.png")
plt.close()

# Chart 2: how rare floods are.
plt.figure(figsize=(6, 4))
plt.bar(["No flood next day", "Flood next day"], data["flood_next_day"].value_counts().sort_index(),
        color=["#9c9c9c", "#e15759"])
plt.title("An unbalanced dataset")
plt.ylabel("Days")
plt.tight_layout()
plt.savefig("chart2_balance.png")
plt.close()

# Chart 3: upstream rain against today's river level.
calm = data[data["flood_next_day"] == 0]
flood = data[data["flood_next_day"] == 1]
plt.figure(figsize=(7, 5))
plt.scatter(calm["upstream_rain_mm"], calm["river_level_m"], s=10, alpha=0.4, color="#9c9c9c", label="No flood next day")
plt.scatter(flood["upstream_rain_mm"], flood["river_level_m"], s=30, color="#e15759", label="Flood next day")
plt.axhline(DANGER, color="#e15759", linestyle="--", linewidth=0.8)
plt.title("Upstream rain and today's river level")
plt.xlabel("Upstream rain today (mm)")
plt.ylabel("River level today (m)")
plt.legend()
plt.grid(True, linestyle="--", alpha=0.3)
plt.tight_layout()
plt.savefig("chart3_upstream.png")
plt.close()

# %% Step 6 - Data Preparation
print("--- Step 6: Data Preparation ---")
before = len(data)
data = data.drop_duplicates(subset="date").sort_values("date").reset_index(drop=True)
print("Duplicate days removed      :", before - len(data))
data["rain_mm"] = pd.to_numeric(data["rain_mm"])
blank = data["rain_mm"].isnull().sum()
# A blank local rain reading is filled with that day's upstream rain, the
# nearest gauge - better than a seasonal average on a stormy day.
data["rain_mm"] = data["rain_mm"].fillna(data["upstream_rain_mm"])
print("Blank rain readings filled  :", blank, "(from the upstream gauge)")

# Feature engineering: rain over the last three days, and the gap between
# today's level and the danger mark.
data["year"] = data["date"].dt.year
data["rain_3day_mm"] = data.groupby("year")["rain_mm"].transform(lambda r: r.rolling(3, min_periods=1).sum())
data["upstream_3day_mm"] = data.groupby("year")["upstream_rain_mm"].transform(lambda r: r.rolling(3, min_periods=1).sum())
data["metres_below_danger"] = DANGER - data["river_level_m"]
features = ["rain_mm", "upstream_rain_mm", "rain_3day_mm", "upstream_3day_mm", "river_level_m",
            "rise_24h_m", "metres_below_danger", "soil_wetness", "dam_release_cumecs"]
print("Features used               :", len(features))
print()

# Chart 4: three-day upstream rain for flood and calm days.
plt.figure(figsize=(6.5, 4.3))
plt.boxplot([data[data["flood_next_day"] == 0]["upstream_3day_mm"], data[data["flood_next_day"] == 1]["upstream_3day_mm"]],
            labels=["No flood next day", "Flood next day"])
plt.title("Upstream rain over the last three days")
plt.ylabel("mm")
plt.grid(True, axis="y", linestyle="--", alpha=0.3)
plt.tight_layout()
plt.savefig("chart4_upstream_3day.png")
plt.close()

# %% Step 7 - Modelling
print("--- Step 7: Modelling ---")
# A warning system is used on future monsoons, so the test is the two
# latest monsoons, 2023 and 2024 - never mixed into training.
train = data[data["year"] <= 2022]
test = data[data["year"] >= 2023]
X_train, y_train = train[features], train["flood_next_day"]
X_test, y_test = test[features], test["flood_next_day"]
print("Training days (2015-2022) :", len(train), " flood days :", y_train.sum())
print("Testing days (2023-2024)  :", len(test), " flood days :", y_test.sum())

# class_weight="balanced" tells each model that a missed flood day counts
# as much as all the calm days together - otherwise it learns to say "no".
logistic = make_pipeline(StandardScaler(), LogisticRegression(class_weight="balanced", max_iter=2000))
forest = RandomForestClassifier(n_estimators=300, min_samples_leaf=3, class_weight="balanced", random_state=0)
logistic.fit(X_train, y_train)
forest.fit(X_train, y_train)
print()

# %% Step 8 - Evaluation
print("--- Step 8: Evaluation ---")
results = {}
for name, model in [["Logistic Regression", logistic], ["Random Forest", forest]]:
    pred = model.predict(X_test)
    tn, fp, fn, tp = confusion_matrix(y_test, pred).ravel()
    results[name] = pred
    print(name)
    print("  " + name + " TP, FN, FP, TN :", tp, fn, fp, tn)
    print("  " + name + " accuracy  :", round(accuracy_score(y_test, pred), 3))
    print("  " + name + " precision :", round(precision_score(y_test, pred, zero_division=0), 3))
    print("  " + name + " recall    :", round(recall_score(y_test, pred), 3))
    print("  " + name + " F1        :", round(f1_score(y_test, pred), 3))
print()

# Leave-one-year-out cross validation: each monsoon in turn is held out and
# predicted by a model trained on all the other monsoons. Days from the same
# season never sit on both sides of the split.
groups = data["year"]
logo = LeaveOneGroupOut()
pooled = {"Logistic Regression": np.zeros(len(data), dtype=int), "Random Forest": np.zeros(len(data), dtype=int)}
year_recall = {"Logistic Regression": [], "Random Forest": []}
years = []
year_floods = []
for tr, te in logo.split(data[features], data["flood_next_day"], groups):
    years.append(int(groups.iloc[te[0]]))
    yt = data["flood_next_day"].iloc[te]
    year_floods.append(int(yt.sum()))
    for name in pooled:
        if name == "Logistic Regression":
            m = make_pipeline(StandardScaler(), LogisticRegression(class_weight="balanced", max_iter=2000))
        else:
            m = RandomForestClassifier(n_estimators=300, min_samples_leaf=3, class_weight="balanced", random_state=0)
        m.fit(data[features].iloc[tr], data["flood_next_day"].iloc[tr])
        p = m.predict(data[features].iloc[te])
        pooled[name][te] = p
        # Recall is undefined for a monsoon with no floods: nothing to catch.
        year_recall[name].append(recall_score(yt, p) if yt.sum() > 0 else np.nan)
print("Leave-one-year-out recall for each held-out monsoon:")
print("  Year  Floods  Logistic  Forest")
for i in range(len(years)):
    if year_floods[i] == 0:
        print("  ", years[i], "    0     no floods to catch")
    else:
        print("  ", years[i], str(year_floods[i]).rjust(5), "   ", round(year_recall["Logistic Regression"][i], 2),
              "    ", round(year_recall["Random Forest"][i], 2))
# Pooled: every day's held-out prediction, from the model that never saw its monsoon.
y_all = data["flood_next_day"]
for name in pooled:
    print(name, "- pooled recall    :", round(recall_score(y_all, pooled[name]), 3))
    print(name, "- pooled precision :", round(precision_score(y_all, pooled[name]), 3))
    print(name, "- pooled F1        :", round(f1_score(y_all, pooled[name]), 3))
print()

# Chart 5: confusion matrices on 2023-2024.
plt.figure(figsize=(10, 4.3))
for k, name in enumerate(["Logistic Regression", "Random Forest"]):
    cm = confusion_matrix(y_test, results[name], labels=[1, 0])
    plt.subplot(1, 2, k + 1)
    plt.imshow(cm, cmap="Reds")
    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=12,
                     color="white" if cm[r][c] > cm.max() / 2 else "black")
    plt.xticks([0, 1], ["Warned", "Not warned"])
    plt.yticks([0, 1], ["Flood", "No flood"])
    plt.title(name)
plt.tight_layout()
plt.savefig("chart5_confusion.png")
plt.close()

# Chart 6: recall in each held-out year (years with no floods left blank).
x = np.arange(len(years))
plt.figure(figsize=(9, 4))
plt.bar(x - 0.2, np.nan_to_num(year_recall["Logistic Regression"]), width=0.4, color="#4e79a7", label="Logistic Regression")
plt.bar(x + 0.2, np.nan_to_num(year_recall["Random Forest"]), width=0.4, color="#59a14f", label="Random Forest")
for i in range(len(years)):
    if year_floods[i] == 0:
        plt.text(x[i], 0.04, "no\nfloods", ha="center", fontsize=8)
plt.xticks(x, [str(y) + "\n(" + str(f) + ")" for y, f in zip(years, year_floods)])
plt.ylim(0, 1.05)
plt.ylabel("Recall (floods warned of)")
plt.xlabel("Held-out monsoon (flood days)")
plt.title("Leave-one-year-out cross validation")
plt.legend()
plt.tight_layout()
plt.savefig("chart6_by_year.png")
plt.close()

# The threshold: Logistic Regression's probability, and how many floods
# are caught - and false alarms raised - at each level.
prob = logistic.predict_proba(X_test)[:, 1]
thresholds = np.round(np.arange(0.1, 0.91, 0.1), 1)
recalls = []
alarms = []
for t in thresholds:
    warn = prob >= t
    recalls.append(recall_score(y_test, warn))
    alarms.append(int((warn & (y_test == 0)).sum()))
chosen = 0.5
print("Threshold used                     :", chosen)
print("False alarms in two monsoons at 0.5 :", alarms[list(thresholds).index(chosen)])
print()

# Chart 7: floods caught and false alarms at each threshold.
fig, ax1 = plt.subplots(figsize=(8, 4.3))
ax1.plot(thresholds, recalls, marker="o", color="#e15759", label="Floods warned of (recall)")
ax1.set_ylabel("Recall")
ax1.set_ylim(0, 1.05)
ax1.set_xlabel("Probability needed to issue a warning")
ax2 = ax1.twinx()
ax2.bar(thresholds, alarms, width=0.05, alpha=0.35, color="#4e79a7", label="False alarms")
ax2.set_ylabel("False alarms in 2023-2024")
ax1.set_title("Warnings caught against false alarms")
fig.legend(loc="upper right", bbox_to_anchor=(0.88, 0.88))
fig.tight_layout()
plt.savefig("chart7_threshold.png")
plt.close()

# %% Step 9 - Deployment
print("--- Step 9: Deployment ---")
joblib.dump(logistic, "flood_model.pkl")
model = joblib.load("flood_model.pkl")
today = pd.DataFrame({
    "rain_mm": [12.0, 85.0], "upstream_rain_mm": [20.0, 140.0], "rain_3day_mm": [30.0, 190.0],
    "upstream_3day_mm": [45.0, 260.0], "river_level_m": [98.4, 102.9], "rise_24h_m": [0.1, 1.4],
    "metres_below_danger": [5.6, 1.1], "soil_wetness": [0.45, 0.95], "dam_release_cumecs": [0, 900],
})
p_today = model.predict_proba(today)[:, 1]
for i, label in enumerate(["An ordinary monsoon day", "After three days of heavy rain"]):
    action = "ISSUE FLOOD WARNING for tomorrow" if p_today[i] >= chosen else "no warning"
    print(label.ljust(32), "probability", round(p_today[i], 3), "->", action)
print()

# %% Step 10 - Feedback
print("--- Step 10: Feedback ---")
# The 2025 monsoon arrives. Each evening the model issues (or does not
# issue) a warning for the next day; afterwards the warnings are checked.
s25 = pd.read_csv("monsoon_2025.csv", parse_dates=["date"])
s25["year"] = 2025
s25["rain_3day_mm"] = s25["rain_mm"].rolling(3, min_periods=1).sum()
s25["upstream_3day_mm"] = s25["upstream_rain_mm"].rolling(3, min_periods=1).sum()
s25["metres_below_danger"] = DANGER - s25["river_level_m"]
warn25 = model.predict_proba(s25[features])[:, 1] >= chosen
flood25 = s25["flood_next_day"] == 1
print("Flood days in the 2025 monsoon      :", int(flood25.sum()))
print("Of these, warned the evening before :", int((warn25 & flood25).sum()))
print("False alarms in 2025                :", int((warn25 & ~flood25).sum()))
print("Recall in 2025                      :", round(recall_score(flood25, warn25), 3))

# Chart 8: the 2025 monsoon, warnings and floods.
plt.figure(figsize=(10, 4))
plt.plot(s25["date"], s25["river_level_m"], color="#4e79a7", label="River level")
plt.axhline(DANGER, color="#e15759", linestyle="--", label="Danger level")
warned_days = s25["date"][warn25]
plt.scatter(warned_days, [DANGER + 1.2] * len(warned_days), marker="v", color="#f28e2b", s=40, label="Warning issued for next day")
plt.title("Feedback: the 2025 monsoon")
plt.ylabel("River level (m)")
plt.legend(loc="upper left")
plt.grid(True, linestyle="--", alpha=0.3)
plt.tight_layout()
plt.savefig("chart8_monsoon_2025.png")
plt.close()
