import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt

# 1. Load your newly minted dataset
df = pd.read_csv("final_kano_dataset.csv")

# 2. Drop any rows with blank data (leftover from our rolling average calculations)
df = df.dropna()

# 3. Define the Features (Inputs) and the Target (Output)
# We tell the AI to look at rainfall, temp, and humidity to predict the river flow
features = ['PRECTOTCORR', 'T2M_MAX', 'T2M_MIN', 'RH2M', '3Day_Rain_Sum']
X = df[features]
y = df['Streamflow_m3s']

# 4. Split data: 80% for training the AI, 20% for testing it
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 5. Build and Train the Random Forest AI
print("Training the AI model... (This might take a few seconds)")
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 6. Make Predictions on the test data
predictions = model.predict(X_test)

# 7. Calculate accuracy metrics for the Chapter 4 Results section
r2 = r2_score(y_test, predictions)
rmse = np.sqrt(mean_squared_error(y_test, predictions))

print("\n--- AI Model Results ---")
print(f"R-squared Score: {r2:.2f} (Closer to 1.0 is better)")
print(f"Root Mean Square Error: {rmse:.2f} m³/s (Average prediction error)")

# 8. Create a visual graph for the project report
plt.figure(figsize=(10, 5))
# We plot just the first 100 days of the test set so the graph is easy to read
plt.plot(y_test.values[:100], label="Actual Streamflow", color='blue', alpha=0.7)
plt.plot(predictions[:100], label="AI Predicted Flow", color='red', alpha=0.7, linestyle='dashed')
plt.title("Kano River: Actual vs. AI Predicted Streamflow (Sample Test Period)")
plt.xlabel("Days")
plt.ylabel("Streamflow (m³/s)")
plt.legend()

# Save the graph as an image file in your folder
plt.savefig("AI_Prediction_Graph.png", bbox_inches='tight')
print("\nSuccess! A graph named 'AI_Prediction_Graph.png' has been saved in your folder.")