import numpy as np

try:
    from river import compose, linear_model, preprocessing, metrics, optim
    from river.ensemble import AdaptiveRandomForestRegressor
    from river.tree import HoeffdingTreeRegressor
    HAS_RIVER = True
except ImportError:
    HAS_RIVER = False

class OnlineLearningTracker:
    """True online learning system (1-sample streaming updates) using River."""
    
    def __init__(self, feature_dimension: int):
        self.feature_dimension = feature_dimension
        self.models = {}
        self.metrics = {}
        
        if HAS_RIVER:
            # Linear online model
            self.models['online_linear'] = compose.Pipeline(
                preprocessing.StandardScaler(),
                linear_model.LinearRegression(optimizer=optim.SGD(0.01))
            )
            
            # Adaptive Random Forest (Concept drift handler)
            self.models['online_arf'] = AdaptiveRandomForestRegressor(seed=42)
            
            # Hoeffding Tree for streaming fast updates
            self.models['online_ht'] = HoeffdingTreeRegressor()
            
            # Track MAE for each
            self.metrics['online_linear'] = metrics.MAE()
            self.metrics['online_arf'] = metrics.MAE()
            self.metrics['online_ht'] = metrics.MAE()
            
    def update(self, features: np.ndarray, target: float):
        """Update the online models step-by-step."""
        if not HAS_RIVER: return
        
        # Convert numpy array to dictionary for River
        x_dict = {f"f_{i}": val for i, val in enumerate(features[:self.feature_dimension])}
        
        for name, model in self.models.items():
            try:
                # 1. Update metric BEFORE training
                y_pred = model.predict_one(x_dict)
                self.metrics[name].update(target, y_pred)
                
                # 2. Train model with actual target
                model.learn_one(x_dict, target)
            except Exception:
                pass
                
    def predict(self, features: np.ndarray) -> dict:
        """Returns predictions and their streaming MAE weights."""
        if not HAS_RIVER: return {}
        
        preds = {}
        x_dict = {f"f_{i}": val for i, val in enumerate(features[:self.feature_dimension])}
        
        for name, model in self.models.items():
            try:
                pred = model.predict_one(x_dict)
                # Map streaming error -> score (e.g. 100 * (1 - error/50))
                error = self.metrics[name].get()
                score = max(0, min(100, 100 * (1 - error / 50)))
                preds[name] = {"prediction": np.clip(pred, 0, 100), "score": score}
            except Exception:
                pass
                
        return preds
