Update src/train_model.py

src/train_model.py

@@ -1,380 +1,325 @@
 """
-CYCLONE INTENSITY PREDICTION - XGBOOST MODEL
-Trains an XGBoost model to predict cyclone wind speed 24 hours ahead
-Includes hyperparameter tuning, feature importance, and comprehensive evaluation
+train_model.py
+==============
+Trains FuzzyNeuralNetwork models for all four disaster types.
+
+Usage:
+    python train_model.py                    # Train all
+    python train_model.py --disaster flood   # Train one
+    python train_model.py --disaster flood --epochs 300
+
+Synthetic Data Strategy:
+  Since real labeled training data is rarely available in a single format,
+  this script generates physically-motivated synthetic datasets.
+  
+  Each dataset is constructed so that the ground-truth risk label follows
+  the domain logic (e.g., high rainfall + low elevation + poor drainage → flood risk).
+  
+  When you have real data:
+    Replace the generate_*_data() functions with your own data loaders.
+    The rest of the training pipeline stays identical.
 """
 
-import pandas as pd
+import torch
 import numpy as np
-import xgboost as xgb
-from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
-from sklearn.model_selection import RandomizedSearchCV, KFold
-import matplotlib.pyplot as plt
-import seaborn as sns
-import joblib
 import os
-from datetime import datetime
+import argparse
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import roc_auc_score, mean_absolute_error
 
-# Configuration
-INPUT_FILE = "data/preprocessed.csv"
-MODEL_OUTPUT = "models/xgboost_cyclone_model.pkl"
-RESULTS_DIR = "results/"
-
-# Create directories
-os.makedirs("models", exist_ok=True)
-os.makedirs(RESULTS_DIR, exist_ok=True)
-
-print("=" * 80)
-print("CYCLONE INTENSITY PREDICTION - XGBOOST TRAINING")
-print("=" * 80)
-
-# ============================================================================
-# STEP 1: LOAD PREPROCESSED DATA
-# ============================================================================
-print("\n📂 STEP 1: Loading preprocessed data...")
-print("-" * 80)
-
-df = pd.read_csv(INPUT_FILE)
-print(f"✅ Loaded dataset: {df.shape[0]} rows × {df.shape[1]} columns")
-
-# ============================================================================
-# STEP 2: PREPARE FEATURES AND TARGET
-# ============================================================================
-print("\n🎯 STEP 2: Preparing features and target...")
-print("-" * 80)
-
-# Define columns to exclude from features
-exclude_cols = [
-    'DATE_TIME', 'TARGET_24H', 'SPLIT', 'BASIN', 'TECH',
-    'CYCLONE_NUMBER', 'STORM_TYPE',  # Already one-hot encoded
-    'YEAR'  # Keep MONTH, HOUR, DOY but not YEAR (prevents overfitting)
-]
-
-# Get feature columns
-feature_cols = [col for col in df.columns if col not in exclude_cols]
-target_col = 'TARGET_24H'
-
-print(f"📊 Total features: {len(feature_cols)}")
-print(f"🎯 Target variable: {target_col}")
-
-# Split by pre-defined split column
-train_df = df[df['SPLIT'] == 'train'].copy()
-test_df = df[df['SPLIT'] == 'test'].copy()
-
-X_train = train_df[feature_cols]
-y_train = train_df[target_col]
-X_test = test_df[feature_cols]
-y_test = test_df[target_col]
-
-print(f"\n✅ Training set: {X_train.shape[0]} samples")
-print(f"✅ Test set: {X_test.shape[0]} samples")
-print(f"   Test set percentage: {len(test_df)/len(df)*100:.1f}%")
-
-# Check for any remaining NaN values
-if X_train.isnull().any().any():
-    print("\n⚠️  Warning: NaN values found in training features")
-    print(X_train.isnull().sum()[X_train.isnull().sum() > 0])
-    print("   Filling NaN with 0...")
-    X_train = X_train.fillna(0)
-    X_test = X_test.fillna(0)
-
-# ============================================================================
-# STEP 3: BASELINE MODEL (DEFAULT PARAMETERS)
-# ============================================================================
-print("\n🤖 STEP 3: Training baseline XGBoost model...")
-print("-" * 80)
-
-baseline_model = xgb.XGBRegressor(
-    n_estimators=100,
-    max_depth=6,
-    learning_rate=0.1,
-    random_state=42,
-    n_jobs=-1
+from src.fuzzy_neural_network import FuzzyNeuralNetwork, FNNTrainer, save_model
+from src.disaster_predictors import (
+    FLOOD_FEATURES, CYCLONE_FEATURES, LANDSLIDE_FEATURES, EARTHQUAKE_FEATURES
 )
 
-baseline_model.fit(X_train, y_train)
-y_pred_baseline = baseline_model.predict(X_test)
-
-# Baseline metrics
-mae_baseline = mean_absolute_error(y_test, y_pred_baseline)
-rmse_baseline = np.sqrt(mean_squared_error(y_test, y_pred_baseline))
-r2_baseline = r2_score(y_test, y_pred_baseline)
+MODEL_DIR = "models"
+SEED = 42
+np.random.seed(SEED)
+torch.manual_seed(SEED)
 
-print(f"\n✅ BASELINE MODEL RESULTS:")
-print(f"   MAE:  {mae_baseline:.2f} knots")
-print(f"   RMSE: {rmse_baseline:.2f} knots")
-print(f"   R²:   {r2_baseline:.4f}")
 
 # ============================================================================
-# STEP 4: HYPERPARAMETER TUNING
+# SYNTHETIC DATA GENERATORS
 # ============================================================================
-print("\n⚙️  STEP 4: Hyperparameter tuning with RandomizedSearchCV...")
-print("-" * 80)
-print("   This may take several minutes...")
-
-# Define hyperparameter search space
-param_distributions = {
-    'n_estimators': [100, 200, 300, 500],
-    'max_depth': [3, 5, 7, 9, 11],
-    'learning_rate': [0.01, 0.05, 0.1, 0.2],
-    'subsample': [0.6, 0.8, 1.0],
-    'colsample_bytree': [0.6, 0.8, 1.0],
-    'min_child_weight': [1, 3, 5],
-    'gamma': [0, 0.1, 0.2, 0.3],
-    'reg_alpha': [0, 0.1, 0.5, 1.0],  # L1 regularization
-    'reg_lambda': [0.5, 1.0, 2.0]     # L2 regularization
+# Each function returns (X: np.ndarray, y: np.ndarray)
+# X shape: (n_samples, n_features) — already normalized to [0, 1]
+# y shape: (n_samples,) — continuous risk score in [0, 1]
+
+def generate_flood_data(n: int = 5000):
+    """
+    Flood risk is driven by:
+      - High rainfall
+      - Low elevation
+      - High soil saturation
+      - Low drainage capacity
+      - Close proximity to rivers
+    """
+    rng = np.random.default_rng(SEED)
+
+    rainfall_norm      = rng.beta(2, 5, n)          # Skewed: most days low rainfall
+    elevation_norm     = rng.beta(3, 2, n)           # Skewed: most areas higher ground
+    slope_norm         = rng.beta(2, 5, n)
+    soil_sat_norm      = rng.beta(2, 3, n)
+    dist_river_norm    = rng.beta(2, 2, n)
+    drainage_norm      = rng.beta(3, 2, n)           # Most areas have decent drainage
+    hist_flood_norm    = rng.beta(1.5, 3, n)
+    pop_density_norm   = rng.beta(2, 2, n)
+
+    X = np.column_stack([
+        rainfall_norm, elevation_norm, slope_norm, soil_sat_norm,
+        dist_river_norm, drainage_norm, hist_flood_norm, pop_density_norm
+    ])
+
+    # Domain-informed risk formula
+    risk = (
+        0.35 * rainfall_norm +
+        0.25 * (1 - elevation_norm) +         # Low elevation → higher risk
+        0.15 * soil_sat_norm +
+        0.10 * (1 - drainage_norm) +           # Poor drainage → higher risk
+        0.08 * (1 - dist_river_norm) +         # Close to river → higher risk
+        0.07 * hist_flood_norm
+    )
+
+    # Add noise and clip
+    risk += rng.normal(0, 0.05, n)
+    y = np.clip(risk, 0.0, 1.0).astype(np.float32)
+
+    return X.astype(np.float32), y
+
+
+def generate_cyclone_data(n: int = 3000):
+    rng = np.random.default_rng(SEED + 1)
+
+    wind_norm       = rng.beta(2, 5, n)
+    pressure_norm   = rng.beta(3, 2, n)       # Higher value = lower pressure = worse
+    sst_norm        = rng.beta(3, 3, n)
+    curvature_norm  = rng.beta(2, 3, n)
+    dist_coast_norm = rng.beta(2, 2, n)
+    surge_norm      = rng.beta(2, 4, n)
+    moisture_norm   = rng.beta(3, 3, n)
+    shear_norm      = rng.beta(2, 3, n)       # High shear weakens cyclones
+
+    X = np.column_stack([
+        wind_norm, pressure_norm, sst_norm, curvature_norm,
+        dist_coast_norm, surge_norm, moisture_norm, shear_norm
+    ])
+
+    risk = (
+        0.30 * wind_norm +
+        0.25 * (1 - pressure_norm) +           # Low pressure = higher intensity
+        0.15 * sst_norm +                      # Warm water feeds cyclones
+        0.10 * surge_norm +
+        0.10 * (1 - dist_coast_norm) +
+        0.05 * moisture_norm +
+        0.05 * (1 - shear_norm)                # Low shear = stronger cyclone
+    )
+
+    risk += rng.normal(0, 0.05, n)
+    y = np.clip(risk, 0.0, 1.0).astype(np.float32)
+
+    return X.astype(np.float32), y
+
+
+def generate_landslide_data(n: int = 4000):
+    rng = np.random.default_rng(SEED + 2)
+
+    slope_norm      = rng.beta(2, 3, n)
+    rainfall_norm   = rng.beta(2, 5, n)
+    soil_norm       = rng.beta(3, 2, n)        # Higher = more stable soil
+    veg_norm        = rng.beta(3, 2, n)        # Higher = more vegetation = more stable
+    seismic_norm    = rng.beta(1.5, 4, n)
+    fault_norm      = rng.beta(2, 2, n)        # Higher = farther from fault
+    aspect_norm     = rng.beta(2, 2, n)
+    hist_norm       = rng.beta(1.5, 4, n)
+
+    X = np.column_stack([
+        slope_norm, rainfall_norm, soil_norm, veg_norm,
+        seismic_norm, fault_norm, aspect_norm, hist_norm
+    ])
+
+    risk = (
+        0.30 * slope_norm +
+        0.25 * rainfall_norm +
+        0.15 * (1 - soil_norm) +               # Unstable soil → higher risk
+        0.10 * (1 - veg_norm) +                # No vegetation → higher risk
+        0.10 * seismic_norm +
+        0.05 * (1 - fault_norm) +              # Close to fault → higher risk
+        0.05 * hist_norm
+    )
+
+    risk += rng.normal(0, 0.05, n)
+    y = np.clip(risk, 0.0, 1.0).astype(np.float32)
+
+    return X.astype(np.float32), y
+
+
+def generate_earthquake_data(n: int = 3000):
+    rng = np.random.default_rng(SEED + 3)
+
+    hist_seism_norm  = rng.beta(2, 4, n)
+    fault_norm       = rng.beta(2, 2, n)       # Higher = farther from fault
+    liquef_norm      = rng.beta(2, 4, n)
+    depth_norm       = rng.beta(3, 2, n)       # Higher = deeper = less damage
+    stress_norm      = rng.beta(2, 3, n)
+    vuln_norm        = rng.beta(2, 3, n)
+    pop_norm         = rng.beta(2, 2, n)
+    amp_norm         = rng.beta(2, 3, n)
+
+    X = np.column_stack([
+        hist_seism_norm, fault_norm, liquef_norm, depth_norm,
+        stress_norm, vuln_norm, pop_norm, amp_norm
+    ])
+
+    risk = (
+        0.25 * hist_seism_norm +
+        0.20 * (1 - fault_norm) +              # Close to fault = more risk
+        0.15 * liquef_norm +
+        0.10 * (1 - depth_norm) +              # Shallow = more damage
+        0.10 * stress_norm +
+        0.10 * vuln_norm +
+        0.05 * pop_norm +
+        0.05 * amp_norm
+    )
+
+    risk += rng.normal(0, 0.05, n)
+    y = np.clip(risk, 0.0, 1.0).astype(np.float32)
+
+    return X.astype(np.float32), y
+
+
+DATA_GENERATORS = {
+    "flood":      (generate_flood_data,      FLOOD_FEATURES),
+    "cyclone":    (generate_cyclone_data,     CYCLONE_FEATURES),
+    "landslide":  (generate_landslide_data,   LANDSLIDE_FEATURES),
+    "earthquake": (generate_earthquake_data,  EARTHQUAKE_FEATURES),
 }
 
-# Create base model
-xgb_model = xgb.XGBRegressor(random_state=42, n_jobs=-1)
-
-# RandomizedSearchCV with cross-validation
-random_search = RandomizedSearchCV(
-    xgb_model,
-    param_distributions=param_distributions,
-    n_iter=50,  # Try 50 random combinations
-    cv=5,       # 5-fold cross-validation
-    scoring='neg_mean_absolute_error',
-    n_jobs=-1,
-    random_state=42,
-    verbose=1
-)
-
-random_search.fit(X_train, y_train)
-
-print(f"\n✅ Best hyperparameters found:")
-for param, value in random_search.best_params_.items():
-    print(f"   {param}: {value}")
-
-# Get best model
-best_model = random_search.best_estimator_
 
 # ============================================================================
-# STEP 5: EVALUATE BEST MODEL
+# TRAINING PIPELINE
 # ============================================================================
-print("\n📊 STEP 5: Evaluating best model...")
-print("-" * 80)
-
-# Predictions
-y_pred_train = best_model.predict(X_train)
-y_pred_test = best_model.predict(X_test)
-
-# Training metrics
-mae_train = mean_absolute_error(y_train, y_pred_train)
-rmse_train = np.sqrt(mean_squared_error(y_train, y_pred_train))
-r2_train = r2_score(y_train, y_pred_train)
-
-# Test metrics
-mae_test = mean_absolute_error(y_test, y_pred_test)
-rmse_test = np.sqrt(mean_squared_error(y_test, y_pred_test))
-r2_test = r2_score(y_test, y_pred_test)
-
-print(f"\n✅ OPTIMIZED MODEL RESULTS:")
-print(f"\n   TRAINING SET:")
-print(f"      MAE:  {mae_train:.2f} knots")
-print(f"      RMSE: {rmse_train:.2f} knots")
-print(f"      R²:   {r2_train:.4f}")
-print(f"\n   TEST SET:")
-print(f"      MAE:  {mae_test:.2f} knots")
-print(f"      RMSE: {rmse_test:.2f} knots")
-print(f"      R²:   {r2_test:.4f}")
-
-# Improvement over baseline
-print(f"\n   📈 IMPROVEMENT OVER BASELINE:")
-print(f"      MAE improvement:  {mae_baseline - mae_test:.2f} knots ({(mae_baseline-mae_test)/mae_baseline*100:.1f}%)")
-print(f"      RMSE improvement: {rmse_baseline - rmse_test:.2f} knots ({(rmse_baseline-rmse_test)/rmse_baseline*100:.1f}%)")
 
-# ============================================================================
-# STEP 6: FEATURE IMPORTANCE
-# ============================================================================
-print("\n📊 STEP 6: Analyzing feature importance...")
-print("-" * 80)
-
-# Get feature importance
-feature_importance = pd.DataFrame({
-    'feature': feature_cols,
-    'importance': best_model.feature_importances_
-}).sort_values('importance', ascending=False)
-
-print(f"\n✅ TOP 20 MOST IMPORTANT FEATURES:")
-print(feature_importance.head(20).to_string(index=False))
-
-# Save full feature importance
-feature_importance.to_csv(f"{RESULTS_DIR}feature_importance.csv", index=False)
-
-# Plot feature importance
-plt.figure(figsize=(12, 8))
-top_features = feature_importance.head(20)
-plt.barh(range(len(top_features)), top_features['importance'])
-plt.yticks(range(len(top_features)), top_features['feature'])
-plt.xlabel('Importance Score')
-plt.title('Top 20 Feature Importance - XGBoost Model')
-plt.gca().invert_yaxis()
-plt.tight_layout()
-plt.savefig(f"{RESULTS_DIR}feature_importance.png", dpi=300, bbox_inches='tight')
-print(f"\n✅ Saved: {RESULTS_DIR}feature_importance.png")
-plt.close()
-
-# ============================================================================
-# STEP 7: PREDICTION ANALYSIS
-# ============================================================================
-print("\n📈 STEP 7: Analyzing predictions...")
-print("-" * 80)
-
-# Create results dataframe
-results_df = pd.DataFrame({
-    'actual': y_test,
-    'predicted': y_pred_test,
-    'error': y_test - y_pred_test,
-    'abs_error': np.abs(y_test - y_pred_test)
-})
-
-# Error statistics by intensity ranges
-print(f"\n✅ ERROR ANALYSIS BY INTENSITY RANGE:")
-intensity_bins = [0, 34, 64, 96, 200]
-intensity_labels = ['Tropical Depression (<34kt)', 
-                    'Tropical Storm (34-63kt)',
-                    'Hurricane/Cyclone (64-95kt)',
-                    'Major Hurricane (>95kt)']
-
-results_df['intensity_category'] = pd.cut(results_df['actual'], 
-                                           bins=intensity_bins, 
-                                           labels=intensity_labels)
-
-for category in intensity_labels:
-    category_data = results_df[results_df['intensity_category'] == category]
-    if len(category_data) > 0:
-        print(f"\n   {category}:")
-        print(f"      Samples: {len(category_data)}")
-        print(f"      MAE: {category_data['abs_error'].mean():.2f} knots")
-        print(f"      Max Error: {category_data['abs_error'].max():.2f} knots")
+def evaluate_model(model: FuzzyNeuralNetwork, X: torch.Tensor, y: torch.Tensor) -> dict:
+    model.eval()
+    with torch.no_grad():
+        preds = model(X).numpy()
+    y_np = y.numpy()
+
+    # Binarize at 0.5 for AUC
+    try:
+        auc = roc_auc_score((y_np > 0.5).astype(int), preds)
+    except Exception:
+        auc = float('nan')
+
+    mae = mean_absolute_error(y_np, preds)
+
+    return {
+        "MAE": round(float(mae), 4),
+        "AUC-ROC": round(float(auc), 4),
+        "Mean Prediction": round(float(preds.mean()), 4),
+        "Std Prediction": round(float(preds.std()), 4),
+    }
+
+
+def train_disaster_model(disaster_type: str, epochs: int = 200, n_samples: int = None):
+    print(f"\n{'='*60}")
+    print(f"  Training FNN for: {disaster_type.upper()}")
+    print(f"{'='*60}")
+
+    generator_fn, feature_names = DATA_GENERATORS[disaster_type]
+    n = n_samples or {"flood": 5000, "cyclone": 3000, "landslide": 4000, "earthquake": 3000}[disaster_type]
+
+    print(f"Generating {n} synthetic samples...")
+    X, y = generator_fn(n)
+
+    # Train/val/test split
+    X_trainval, X_test, y_trainval, y_test = train_test_split(X, y, test_size=0.15, random_state=SEED)
+    X_train, X_val, y_train, y_val = train_test_split(X_trainval, y_trainval, test_size=0.15, random_state=SEED)
+
+    print(f"  Train: {len(X_train)} | Val: {len(X_val)} | Test: {len(X_test)}")
+
+    # Tensors
+    X_train_t = torch.tensor(X_train)
+    y_train_t = torch.tensor(y_train)
+    X_val_t   = torch.tensor(X_val)
+    y_val_t   = torch.tensor(y_val)
+    X_test_t  = torch.tensor(X_test)
+    y_test_t  = torch.tensor(y_test)
+
+    # Model
+    n_features = len(feature_names)
+    model = FuzzyNeuralNetwork(
+        n_features=n_features,
+        n_terms=3,
+        hidden_dims=[64, 32],
+        dropout=0.2
+    )
+
+    print(f"  Model: FNN with {n_features} inputs, 3 fuzzy terms, 64→32 deep head")
+    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"  Trainable parameters: {total_params:,}")
+
+    # Train
+    trainer = FNNTrainer(model, lr=1e-3, weight_decay=1e-4)
+    history = trainer.fit(
+        X_train_t, y_train_t,
+        X_val_t,   y_val_t,
+        epochs=epochs, batch_size=64, patience=25
+    )
+
+    # Evaluate
+    print("\n  Test set evaluation:")
+    metrics = evaluate_model(model, X_test_t, y_test_t)
+    for k, v in metrics.items():
+        print(f"    {k}: {v}")
+
+    # Save
+    os.makedirs(MODEL_DIR, exist_ok=True)
+    model_path = os.path.join(MODEL_DIR, f"fnn_{disaster_type}_model.pt")
+    save_model(model, model_path, feature_names)
+
+    # Save feature names as text too
+    feat_path = os.path.join(MODEL_DIR, "feature_names", f"{disaster_type}_features.txt")
+    os.makedirs(os.path.dirname(feat_path), exist_ok=True)
+    with open(feat_path, "w") as f:
+        f.write("\n".join(feature_names))
+
+    print(f"\n  Model saved to: {model_path}")
+    return metrics
+
+
+def train_all(epochs: int = 200):
+    results = {}
+    for disaster_type in DATA_GENERATORS:
+        metrics = train_disaster_model(disaster_type, epochs=epochs)
+        results[disaster_type] = metrics
+
+    print("\n" + "="*60)
+    print("  TRAINING SUMMARY")
+    print("="*60)
+    for dt, metrics in results.items():
+        print(f"  {dt.upper():12s} | MAE: {metrics['MAE']:.4f} | AUC: {metrics['AUC-ROC']:.4f}")
+    print("="*60)
 
-# ============================================================================
-# STEP 8: VISUALIZATION
-# ============================================================================
-print("\n📊 STEP 8: Creating visualizations...")
-print("-" * 80)
-
-# 1. Actual vs Predicted scatter plot
-plt.figure(figsize=(10, 8))
-plt.scatter(y_test, y_pred_test, alpha=0.5, s=30)
-plt.plot([y_test.min(), y_test.max()], 
-         [y_test.min(), y_test.max()], 
-         'r--', lw=2, label='Perfect Prediction')
-plt.xlabel('Actual Wind Speed (knots)', fontsize=12)
-plt.ylabel('Predicted Wind Speed (knots)', fontsize=12)
-plt.title(f'XGBoost: Actual vs Predicted\nMAE: {mae_test:.2f} kt, RMSE: {rmse_test:.2f} kt, R²: {r2_test:.3f}', 
-          fontsize=14)
-plt.legend()
-plt.grid(True, alpha=0.3)
-plt.tight_layout()
-plt.savefig(f"{RESULTS_DIR}actual_vs_predicted.png", dpi=300, bbox_inches='tight')
-print(f"✅ Saved: {RESULTS_DIR}actual_vs_predicted.png")
-plt.close()
-
-# 2. Error distribution
-plt.figure(figsize=(12, 5))
-
-plt.subplot(1, 2, 1)
-plt.hist(results_df['error'], bins=50, edgecolor='black', alpha=0.7)
-plt.xlabel('Prediction Error (knots)', fontsize=11)
-plt.ylabel('Frequency', fontsize=11)
-plt.title('Error Distribution', fontsize=12)
-plt.axvline(0, color='red', linestyle='--', linewidth=2, label='Zero Error')
-plt.legend()
-plt.grid(True, alpha=0.3)
-
-plt.subplot(1, 2, 2)
-plt.boxplot([results_df[results_df['intensity_category'] == cat]['abs_error'].dropna() 
-             for cat in intensity_labels],
-            labels=['TD', 'TS', 'Hurricane', 'Major'])
-plt.ylabel('Absolute Error (knots)', fontsize=11)
-plt.xlabel('Storm Intensity Category', fontsize=11)
-plt.title('Error by Storm Intensity', fontsize=12)
-plt.xticks(rotation=45, ha='right')
-plt.grid(True, alpha=0.3, axis='y')
-
-plt.tight_layout()
-plt.savefig(f"{RESULTS_DIR}error_analysis.png", dpi=300, bbox_inches='tight')
-print(f"✅ Saved: {RESULTS_DIR}error_analysis.png")
-plt.close()
-
-# 3. Learning curve (training history)
-print(f"✅ Saved: {RESULTS_DIR}learning_curve.png")
 
 # ============================================================================
-# STEP 9: SAVE MODEL
+# ENTRY POINT
 # ============================================================================
-print("\n💾 STEP 9: Saving trained model...")
-print("-" * 80)
-
-# Save the model
-joblib.dump(best_model, MODEL_OUTPUT)
-print(f"✅ Model saved: {MODEL_OUTPUT}")
-
-# Save feature names
-feature_names_file = "models/feature_names.txt"
-with open(feature_names_file, 'w') as f:
-    for feature in feature_cols:
-        f.write(f"{feature}\n")
-print(f"✅ Feature names saved: {feature_names_file}")
-
-# Save model metadata
-metadata = {
-    'training_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
-    'n_features': len(feature_cols),
-    'n_train_samples': len(X_train),
-    'n_test_samples': len(X_test),
-    'test_mae': float(mae_test),
-    'test_rmse': float(rmse_test),
-    'test_r2': float(r2_test),
-    'best_params': random_search.best_params_
-}
 
-metadata_df = pd.DataFrame([metadata])
-metadata_df.to_csv(f"{RESULTS_DIR}model_metadata.csv", index=False)
-print(f"✅ Metadata saved: {RESULTS_DIR}model_metadata.csv")
-
-# ============================================================================
-# STEP 10: SUMMARY REPORT
-# ============================================================================
-print("\n" + "=" * 80)
-print("🎉 TRAINING COMPLETE!")
-print("=" * 80)
-
-print(f"""
-📊 FINAL MODEL PERFORMANCE:
-   • Test MAE:  {mae_test:.2f} knots
-   • Test RMSE: {rmse_test:.2f} knots
-   • Test R²:   {r2_test:.4f}
-   
-📁 OUTPUT FILES:
-   • Model: {MODEL_OUTPUT}
-   • Feature names: {feature_names_file}
-   • Feature importance: {RESULTS_DIR}feature_importance.csv
-   • Visualizations: {RESULTS_DIR}*.png
-   • Metadata: {RESULTS_DIR}model_metadata.csv
-
-🎯 MODEL INTERPRETATION:
-   • MAE of {mae_test:.2f} kt means on average, predictions are off by ~{mae_test:.0f} knots
-   • For a 100kt cyclone, expect ±{mae_test:.0f} kt error range
-   • R² of {r2_test:.4f} means model explains {r2_test*100:.1f}% of variance
-
-🚀 NEXT STEPS:
-   1. Use this model for real-time predictions
-   2. Monitor performance on new cyclones
-   3. Retrain with more recent data periodically
-   4. Consider ensemble with other models (Random Forest, Neural Networks)
-
-💡 USAGE:
-   ```python
-   import joblib
-   model = joblib.load('{MODEL_OUTPUT}')
-   
-   # Make prediction for new cyclone data
-   prediction = model.predict(new_cyclone_features)
-   print(f"Predicted wind speed in 24h: {{prediction[0]:.1f}} knots")
-   ```
-""")
-
-print("=" * 80)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train FNN disaster models")
+    parser.add_argument(
+        "--disaster",
+        choices=list(DATA_GENERATORS.keys()) + ["all"],
+        default="all",
+        help="Which disaster model to train"
+    )
+    parser.add_argument("--epochs", type=int, default=200)
+    parser.add_argument("--samples", type=int, default=None)
+
+    args = parser.parse_args()
+
+    if args.disaster == "all":
+        train_all(epochs=args.epochs)
+    else:
+        train_disaster_model(args.disaster, epochs=args.epochs, n_samples=args.samples)