"""
Train ML models on the telecom dataset and save them as .joblib files.
Run: python train_models.py
"""
import sys
import os
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

import pandas as pd
import numpy as np
import joblib
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.impute import SimpleImputer

print("=" * 60)
print("QoS Buddy — Model Training")
print("=" * 60)

# ── Load dataset ──
DATA_PATHS = [
    Path(__file__).parent / "data" / "processed" / "network_kpis_full.csv",
    Path(__file__).parent / "unified_telecom_dataset_FILLED.csv",
    Path(__file__).parent.parent / "data" / "unified_telecom_dataset_FILLED.csv",
]

df = None
for p in DATA_PATHS:
    if p.exists():
        print(f"\nLoading dataset: {p.name}")
        df = pd.read_csv(p)
        print(f"  Rows: {len(df):,}  |  Columns: {len(df.columns)}")
        break

if df is None:
    print("ERROR: No dataset found. Run generate_data.py first.")
    sys.exit(1)

# ── Feature columns ──
FEATURE_COLS = [
    "RSRP", "SINR", "PRB_UL", "PRB_DL",
    "throughput_DL", "throughput_UL",
    "latency_ms", "MOS",
    "call_drop_rate", "handover_success_rate",
    "BLER_UL", "BLER_DL", "CQI", "RLC_retransmissions",
    "hour",
]

# Keep only columns that exist
FEATURE_COLS = [c for c in FEATURE_COLS if c in df.columns]
print(f"\nFeature columns ({len(FEATURE_COLS)}): {FEATURE_COLS}")

# ── Prepare features ──
X_raw = df[FEATURE_COLS].copy()
imputer = SimpleImputer(strategy="median")
X = imputer.fit_transform(X_raw)

# ── Fit unified scaler ──
print("\n[1/4] Fitting unified scaler...")
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
scaler.feature_names_in_ = np.array(FEATURE_COLS)
print(f"  Scaler fitted on {X_scaled.shape[0]:,} samples, {X_scaled.shape[1]} features")

# ── Train BO2 SLA Risk model ──
print("\n[2/4] Training BO2 SLA Risk model (Random Forest)...")
if "sla_breach_risk" in df.columns:
    y_sla = pd.to_numeric(df["sla_breach_risk"], errors="coerce").fillna(0).values
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y_sla, test_size=0.2, random_state=42)
    
    rf_sla = RandomForestRegressor(n_estimators=100, max_depth=12, random_state=42, n_jobs=-1)
    rf_sla.fit(X_train, y_train)
    
    y_pred = rf_sla.predict(X_test)
    r2 = r2_score(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    
    cv_scores = cross_val_score(rf_sla, X_scaled[:5000], y_sla[:5000], cv=5, scoring="r2")
    
    print(f"  R²: {r2:.4f}  |  MAE: {mae:.4f}  |  RMSE: {rmse:.4f}")
    print(f"  CV R² scores: {cv_scores.round(4)}")
    print(f"  CV R² mean: {cv_scores.mean():.4f} ± {cv_scores.std():.4f}")
    
    # Store metrics in model
    rf_sla.metrics_ = {
        "r2": float(r2),
        "mae": float(mae),
        "rmse": float(rmse),
        "cv_r2_mean": float(cv_scores.mean()),
        "cv_r2_std": float(cv_scores.std()),
        "cv_scores": cv_scores.tolist(),
    }
    rf_sla.feature_names_in_ = np.array(FEATURE_COLS)
    rf_sla.training_stats_ = {
        "feature_importances": dict(zip(FEATURE_COLS, rf_sla.feature_importances_.tolist())),
        "cv_scores": cv_scores.tolist(),
    }
else:
    print("  WARNING: sla_breach_risk column not found, creating dummy model")
    rf_sla = RandomForestRegressor(n_estimators=10, random_state=42)
    rf_sla.fit(X_scaled[:100], np.random.rand(100))

# ── Train BO3 CAPEX Score model ──
print("\n[3/4] Training BO3 CAPEX Score model (Random Forest)...")
if "capex_score" in df.columns:
    y_capex = pd.to_numeric(df["capex_score"], errors="coerce").fillna(0).values
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y_capex, test_size=0.2, random_state=42)
    
    rf_capex = RandomForestRegressor(n_estimators=100, max_depth=12, random_state=42, n_jobs=-1)
    rf_capex.fit(X_train, y_train)
    
    y_pred = rf_capex.predict(X_test)
    r2_c = r2_score(y_test, y_pred)
    mae_c = mean_absolute_error(y_test, y_pred)
    rmse_c = np.sqrt(mean_squared_error(y_test, y_pred))
    
    print(f"  R²: {r2_c:.4f}  |  MAE: {mae_c:.4f}  |  RMSE: {rmse_c:.4f}")
    rf_capex.feature_names_in_ = np.array(FEATURE_COLS)
else:
    print("  WARNING: capex_score column not found, creating dummy model")
    rf_capex = RandomForestRegressor(n_estimators=10, random_state=42)
    rf_capex.fit(X_scaled[:100], np.random.rand(100))

# ── Save models ──
print("\n[4/4] Saving models...")
models_dir = Path(__file__).parent.parent / "models"
models_dir.mkdir(parents=True, exist_ok=True)

joblib.dump(scaler,    models_dir / "unified_scaler.joblib")
joblib.dump(rf_sla,    models_dir / "bo2_sla_risk.joblib")
joblib.dump(rf_capex,  models_dir / "bo3_capex_score.joblib")

print(f"  Saved to: {models_dir}")
print(f"  - unified_scaler.joblib")
print(f"  - bo2_sla_risk.joblib")
print(f"  - bo3_capex_score.joblib")

# ── Save runtime metadata ──
import json
meta = {
    "trained": True,
    "target": "sla_breach_risk",
    "features": FEATURE_COLS,
    "metrics": rf_sla.metrics_ if hasattr(rf_sla, "metrics_") else {},
    "training_stats": rf_sla.training_stats_ if hasattr(rf_sla, "training_stats_") else {},
    "fit_at": pd.Timestamp.now().isoformat(),
    "version": 73,
    "is_v7_champion": True,
    "dataset": "network_kpis_full.csv",
    "n_samples": len(df),
}
meta_path = models_dir / "ml_runtime_metadata.json"
with open(meta_path, "w") as f:
    json.dump(meta, f, indent=2)
print(f"  - ml_runtime_metadata.json")

print("\n" + "=" * 60)
print("✅ Training complete!")
print(f"   SLA Risk Model  → R²={rf_sla.metrics_['r2']:.4f}" if hasattr(rf_sla, "metrics_") else "")
print(f"   CAPEX Model     → R²={r2_c:.4f}" if "r2_c" in dir() else "")
print("=" * 60)