SelfTrainService/trainer.py

"""
ML Model Trainer for Schedule Optimization
Handles model training and retraining with multiple models and ensemble
"""
import os
import pickle
import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, Dict, Tuple
import numpy as np

from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import xgboost as xgb
import catboost as cb
import lightgbm as lgb

from .config import CONFIG
from .data_store import ScheduleDataStore
from .feature_extractor import FeatureExtractor


class ModelTrainer:
    """Train and manage ML models for schedule optimization"""
    
    def __init__(self, model_dir: Optional[str] = None):
        self.model_dir = Path(model_dir or CONFIG.MODEL_DIR)
        self.model_dir.mkdir(parents=True, exist_ok=True)
        
        self.data_store = ScheduleDataStore()
        self.feature_extractor = FeatureExtractor()
        
        self.models = {}  # Dictionary of trained models
        self.model_scores = {}  # Performance scores for each model
        self.ensemble_weights = {}  # Weights for ensemble
        self.best_model_name = None
        self.last_trained = None
        self.training_history = []
    
    def _get_model(self, model_name: str):
        """Get model instance by name"""
        if model_name == "gradient_boosting":
            return GradientBoostingRegressor(
                n_estimators=CONFIG.EPOCHS,
                learning_rate=CONFIG.LEARNING_RATE,
                random_state=42
            )
        
        elif model_name == "random_forest":
            return RandomForestRegressor(
                n_estimators=CONFIG.EPOCHS,
                random_state=42,
                n_jobs=-1
            )
        
        elif model_name == "xgboost":
            return xgb.XGBRegressor(
                n_estimators=CONFIG.EPOCHS,
                learning_rate=CONFIG.LEARNING_RATE,
                random_state=42,
                verbosity=0
            )
        
        elif model_name == "lightgbm":
            return lgb.LGBMRegressor(
                n_estimators=CONFIG.EPOCHS,
                learning_rate=CONFIG.LEARNING_RATE,
                random_state=42,
                verbose=-1
            )
        
        elif model_name == "catboost":
            return cb.CatBoostRegressor(
                iterations=CONFIG.EPOCHS,
                learning_rate=CONFIG.LEARNING_RATE,
                random_state=42,
                verbose=False
            )
        
        return None
    
    def should_retrain(self) -> bool:
        """Check if model should be retrained"""
        if not self.last_trained:
            # Never trained
            return True
        
        # Check time since last training
        hours_since_training = (
            datetime.now() - self.last_trained
        ).total_seconds() / 3600
        
        if hours_since_training >= CONFIG.RETRAIN_INTERVAL_HOURS:
            # Check if enough new data
            new_schedules = self.data_store.get_schedules_since(self.last_trained)
            if len(new_schedules) >= CONFIG.MIN_SCHEDULES_FOR_RETRAIN:
                return True
        
        return False
    
    def train(self, force: bool = False) -> Dict:
        """Train or retrain all models"""
        
        if not force and not self.should_retrain():
            return {
                "success": False,
                "reason": "Retraining not needed yet"
            }
        
        # Load data
        schedules = self.data_store.load_schedules()
        
        if len(schedules) < CONFIG.MIN_SCHEDULES_FOR_TRAINING:
            return {
                "success": False,
                "reason": f"Not enough data. Need {CONFIG.MIN_SCHEDULES_FOR_TRAINING}, have {len(schedules)}"
            }
        
        # Prepare dataset
        X, y = self.feature_extractor.prepare_dataset(schedules)
        
        if len(X) == 0:
            return {
                "success": False,
                "error": "No valid features extracted"
            }
        
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=CONFIG.TRAIN_TEST_SPLIT, random_state=42
        )
        
        # Train all models
        self.models = {}
        self.model_scores = {}
        all_metrics = {}
        
        for model_name in CONFIG.MODEL_TYPES:
            print(f"Training {model_name}...")
            model = self._get_model(model_name)
            
            if model is None:
                print(f"Skipping {model_name} - not available")
                continue
            
            # Train model
            model.fit(X_train, y_train)
            
            # Evaluate
            train_pred = model.predict(X_train)
            test_pred = model.predict(X_test)

            train_r2 = r2_score(y_train, train_pred) # type: ignore
            test_r2 = r2_score(y_test, test_pred)   # type: ignore
            test_rmse = np.sqrt(mean_squared_error(y_test, test_pred))  # type: ignore
            
            self.models[model_name] = model
            self.model_scores[model_name] = test_r2
            
            all_metrics[model_name] = {
                "train_r2": train_r2,
                "test_r2": test_r2,
                "train_rmse": np.sqrt(mean_squared_error(y_train, train_pred)), # type: ignore
                "test_rmse": test_rmse
            }
            
            print(f"  {model_name}: R² = {test_r2:.4f}, RMSE = {test_rmse:.4f}")
        
        # Compute ensemble weights based on performance
        if CONFIG.USE_ENSEMBLE and len(self.models) > 1:
            total_score = sum(self.model_scores.values())
            self.ensemble_weights = {
                name: score / total_score 
                for name, score in self.model_scores.items()
            }
        else:
            self.ensemble_weights = {}
        
        # Find best model
        if self.model_scores:
            self.best_model_name = max(self.model_scores.items(), key=lambda x: x[1])[0]
        
        # Save model
        self.last_trained = datetime.now()
        self.save_model()
        
        # Record training history
        history_entry = {
            "timestamp": self.last_trained.isoformat(),
            "metrics": all_metrics,
            "best_model": self.best_model_name,
            "ensemble_weights": self.ensemble_weights,
            "config": {
                "models_trained": list(self.models.keys()),
                "version": CONFIG.MODEL_VERSION
            }
        }
        self.training_history.append(history_entry)
        self._save_history()
        
        return {
            "success": True,
            "models_trained": list(self.models.keys()),
            "best_model": self.best_model_name,
            "metrics": all_metrics,
            "ensemble_weights": self.ensemble_weights,
            "samples_used": len(X),
            "timestamp": self.last_trained.isoformat()
        }
    
    def predict(self, features: Dict[str, float], use_ensemble: bool = True) -> Tuple[float, float]:
        """Predict schedule quality and confidence"""
        if not self.models:
            self.load_model()
        
        if not self.models:
            return 0.0, 0.0
        
        # Convert features to vector
        feature_vector = np.array([
            [features.get(f, 0.0) for f in CONFIG.FEATURES]
        ])
        
        if use_ensemble and CONFIG.USE_ENSEMBLE and self.ensemble_weights:
            # Ensemble prediction
            prediction = 0.0
            for model_name, weight in self.ensemble_weights.items():
                if model_name in self.models:
                    pred = self.models[model_name].predict(feature_vector)[0]
                    prediction += weight * pred
            
            # Confidence based on ensemble agreement
            predictions = [
                self.models[name].predict(feature_vector)[0]
                for name in self.models.keys()
            ]
            std_dev = np.std(predictions)
            confidence = max(0.5, min(1.0, 1.0 - (std_dev / 50)))  # Higher agreement = higher confidence
        else:
            # Use best single model
            best_model = self.models.get(self.best_model_name)
            if best_model is None:
                best_model = list(self.models.values())[0]
            
            prediction = best_model.predict(feature_vector)[0]
            confidence = min(1.0, 0.8 + (prediction / 100) * 0.2)
        
        return float(prediction), float(confidence)
    
    def save_model(self):
        """Save all models to disk"""
        if not self.models:
            return
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        model_path = self.model_dir / f"models_{timestamp}.pkl"
        latest_path = self.model_dir / "models_latest.pkl"
        
        model_data = {
            "models": self.models,
            "ensemble_weights": self.ensemble_weights,
            "best_model_name": self.best_model_name,
            "last_trained": self.last_trained,
            "config": {
                "version": CONFIG.MODEL_VERSION,
                "features": CONFIG.FEATURES,
                "models_trained": list(self.models.keys())
            }
        }
        
        with open(model_path, 'wb') as f:
            pickle.dump(model_data, f)
        
        with open(latest_path, 'wb') as f:
            pickle.dump(model_data, f)
    
    def load_model(self) -> bool:
        """Load models from disk"""
        latest_path = self.model_dir / "models_latest.pkl"
        
        if not latest_path.exists():
            return False
        
        try:
            with open(latest_path, 'rb') as f:
                model_data = pickle.load(f)
            
            self.models = model_data["models"]
            self.ensemble_weights = model_data.get("ensemble_weights", {})
            self.best_model_name = model_data.get("best_model_name")
            self.last_trained = model_data.get("last_trained")
            return True
        except Exception as e:
            print(f"Error loading models: {e}")
            return False
    
    def _save_history(self):
        """Save training history"""
        history_path = self.model_dir / "training_history.json"
        with open(history_path, 'w') as f:
            json.dump(self.training_history, f, indent=2, default=str)
    
    def get_model_info(self) -> Dict:
        """Get information about current models"""
        if not self.models:
            self.load_model()
        
        return {
            "models_loaded": list(self.models.keys()) if self.models else [],
            "best_model": self.best_model_name,
            "ensemble_enabled": CONFIG.USE_ENSEMBLE,
            "ensemble_weights": self.ensemble_weights,
            "last_trained": self.last_trained.isoformat() if self.last_trained else None,
            "should_retrain": self.should_retrain(),
            "schedules_available": self.data_store.count_schedules(),
            "training_runs": len(self.training_history)
        }