ml_engine.py

"""
Real ML/RL Engine for Signal Engine
Implements ensemble ML, reinforcement learning, and genetic algorithms
"""

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score
import joblib
import os
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional
import asyncio
import aiohttp
from gateio_client import GateIOClient

class MLEngine:
    """Real ML engine with historical data training and self-learning"""
    
    def __init__(self, model_dir: str = "models"):
        self.model_dir = model_dir
        os.makedirs(model_dir, exist_ok=True)
        
        # Ensemble models
        self.models = {
            'rf': RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42),
            'gb': GradientBoostingClassifier(n_estimators=100, max_depth=5, random_state=42),
            'svm': SVC(probability=True, random_state=42),
            'lr': LogisticRegression(random_state=42)
        }
        
        self.scaler = StandardScaler()
        self.is_trained = False
        self.feature_names = []
        
    async def fetch_historical_data(self, symbol: str, interval: str = '1h', limit: int = 1000) -> pd.DataFrame:
        """Fetch historical candle data from Gate.io"""
        async with GateIOClient() as client:
            candles = await client.get_candles(symbol, interval=interval, limit=limit)
            
            df = pd.DataFrame(candles, columns=['timestamp', 'volume', 'close', 'high', 'low', 'open'])
            df['timestamp'] = pd.to_datetime(df['timestamp'], unit='s')
            df.set_index('timestamp', inplace=True)
            
            return df.sort_index()
    
    def extract_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Extract technical features for ML"""
        features = pd.DataFrame(index=df.index)
        
        # Price features
        features['close'] = df['close']
        features['high'] = df['high']
        features['low'] = df['low']
        features['volume'] = df['volume']
        
        # Returns
        features['return_1h'] = df['close'].pct_change(1)
        features['return_4h'] = df['close'].pct_change(4)
        features['return_24h'] = df['close'].pct_change(24)
        
        # Volatility
        features['volatility_24h'] = df['close'].pct_change().rolling(24).std()
        
        # Moving averages
        features['sma_7'] = df['close'].rolling(7).mean()
        features['sma_24'] = df['close'].rolling(24).mean()
        features['ema_12'] = df['close'].ewm(span=12).mean()
        
        # RSI
        delta = df['close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(14).mean()
        rs = gain / loss
        features['rsi'] = 100 - (100 / (1 + rs))
        
        # MACD
        ema_12 = df['close'].ewm(span=12).mean()
        ema_26 = df['close'].ewm(span=26).mean()
        features['macd'] = ema_12 - ema_26
        features['macd_signal'] = features['macd'].ewm(span=9).mean()
        
        # Bollinger Bands
        sma_20 = df['close'].rolling(20).mean()
        std_20 = df['close'].rolling(20).std()
        features['bb_upper'] = sma_20 + (std_20 * 2)
        features['bb_lower'] = sma_20 - (std_20 * 2)
        features['bb_width'] = (features['bb_upper'] - features['bb_lower']) / sma_20
        
        # Volume features
        features['volume_sma_24'] = df['volume'].rolling(24).mean()
        features['volume_ratio'] = df['volume'] / features['volume_sma_24']
        
        # Price momentum
        features['momentum_12'] = df['close'] - df['close'].shift(12)
        features['momentum_24'] = df['close'] - df['close'].shift(24)
        
        # Drop NaN values
        features = features.dropna()
        
        self.feature_names = features.columns.tolist()
        return features
    
    def create_labels(self, df: pd.DataFrame, lookahead: int = 4) -> pd.Series:
        """Create binary labels: 1 if price goes up in lookahead, 0 otherwise"""
        future_returns = df['close'].shift(-lookahead) / df['close'] - 1
        labels = (future_returns > 0).astype(int)
        return labels
    
    async def train(self, symbol: str = "BTC_USDT") -> Dict:
        """Train models on historical data"""
        print(f"Fetching historical data for {symbol}...")
        df = await self.fetch_historical_data(symbol)
        
        print("Extracting features...")
        features = self.extract_features(df)
        labels = self.create_labels(df)
        
        # Align features and labels
        aligned_data = pd.concat([features, labels], axis=1).dropna()
        X = aligned_data[self.feature_names]
        y = aligned_data.iloc[:, -1]
        
        # Scale features
        X_scaled = self.scaler.fit_transform(X)
        
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, test_size=0.2, random_state=42, shuffle=False
        )
        
        print(f"Training on {len(X_train)} samples, testing on {len(X_test)} samples...")
        
        # Train each model
        results = {}
        for name, model in self.models.items():
            print(f"Training {name}...")
            model.fit(X_train, y_train)
            
            # Evaluate
            y_pred = model.predict(X_test)
            accuracy = accuracy_score(y_test, y_pred)
            precision = precision_score(y_test, y_pred, average='binary')
            recall = recall_score(y_test, y_pred, average='binary')
            
            results[name] = {
                'accuracy': accuracy,
                'precision': precision,
                'recall': recall
            }
            
            print(f"{name}: Accuracy={accuracy:.3f}, Precision={precision:.3f}, Recall={recall:.3f}")
        
        # Save models
        self.save_models()
        self.is_trained = True
        
        return results
    
    def predict(self, features: Dict) -> Dict:
        """Generate prediction using ensemble"""
        if not self.is_trained:
            raise ValueError("Models not trained. Call train() first.")
        
        # Convert to DataFrame and scale
        X = pd.DataFrame([features])[self.feature_names]
        X_scaled = self.scaler.transform(X)
        
        # Get predictions from all models
        predictions = {}
        for name, model in self.models.items():
            pred_proba = model.predict_proba(X_scaled)[0]
            predictions[name] = {
                'probability_up': pred_proba[1],
                'probability_down': pred_proba[0]
            }
        
        # Ensemble prediction (weighted average)
        weights = {'rf': 0.3, 'gb': 0.3, 'svm': 0.2, 'lr': 0.2}
        ensemble_prob_up = sum(p['probability_up'] * weights[name] for name, p in predictions.items())
        
        direction = 'LONG' if ensemble_prob_up > 0.5 else 'SHORT'
        confidence = max(ensemble_prob_up, 1 - ensemble_prob_up)
        
        return {
            'direction': direction,
            'confidence': confidence,
            'probability_up': ensemble_prob_up,
            'probability_down': 1 - ensemble_prob_up,
            'individual_predictions': predictions
        }
    
    def save_models(self):
        """Save trained models to disk"""
        for name, model in self.models.items():
            joblib.dump(model, f"{self.model_dir}/{name}.pkl")
        joblib.dump(self.scaler, f"{self.model_dir}/scaler.pkl")
        joblib.dump(self.feature_names, f"{self.model_dir}/feature_names.pkl")
        print(f"Models saved to {self.model_dir}")
    
    def load_models(self):
        """Load trained models from disk"""
        for name in self.models.keys():
            self.models[name] = joblib.load(f"{self.model_dir}/{name}.pkl")
        self.scaler = joblib.load(f"{self.model_dir}/scaler.pkl")
        self.feature_names = joblib.load(f"{self.model_dir}/feature_names.pkl")
        self.is_trained = True
        print(f"Models loaded from {self.model_dir}")


class ReinforcementLearningAgent:
    """Simple RL agent for signal generation using Q-learning"""
    
    def __init__(self, state_size: int, action_size: int = 3):
        self.state_size = state_size
        self.action_size = action_size  # LONG, SHORT, HOLD
        self.q_table = np.zeros((state_size, action_size))
        self.learning_rate = 0.1
        self.discount_factor = 0.95
        self.epsilon = 0.1
        self.epsilon_decay = 0.995
        self.epsilon_min = 0.01
        
    def get_state(self, features: Dict) -> int:
        """Convert features to discrete state"""
        # Simplified: use RSI and moving average crossover
        rsi = features.get('rsi', 50)
        sma_7 = features.get('sma_7', 0)
        sma_24 = features.get('sma_24', 0)
        
        # Discretize into 10 states
        state = 0
        if rsi > 70: state += 4
        elif rsi < 30: state += 1
        
        if sma_7 > sma_24: state += 2
        
        return min(state, 9)
    
    def choose_action(self, state: int) -> int:
        """Choose action using epsilon-greedy policy"""
        if np.random.random() < self.epsilon:
            return np.random.choice(self.action_size)
        return np.argmax(self.q_table[state])
    
    def learn(self, state: int, action: int, reward: float, next_state: int):
        """Update Q-table using Q-learning"""
        best_next_action = np.argmax(self.q_table[next_state])
        td_target = reward + self.discount_factor * self.q_table[next_state][best_next_action]
        td_error = td_target - self.q_table[state][action]
        self.q_table[state][action] += self.learning_rate * td_error
        
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay
    
    def get_signal(self, features: Dict) -> str:
        """Get trading signal from RL agent"""
        state = self.get_state(features)
        action = self.choose_action(state)
        
        actions = ['LONG', 'SHORT', 'HOLD']
        return actions[action]


class GeneticAlgorithmOptimizer:
    """Genetic algorithm for hyperparameter optimization"""
    
    def __init__(self, population_size: int = 20, generations: int = 50):
        self.population_size = population_size
        self.generations = generations
        self.mutation_rate = 0.1
        self.crossover_rate = 0.7
        
    def create_individual(self) -> Dict:
        """Create random hyperparameters"""
        return {
            'n_estimators': np.random.randint(50, 200),
            'max_depth': np.random.randint(3, 15),
            'learning_rate': np.random.uniform(0.01, 0.3),
            'min_samples_split': np.random.randint(2, 10)
        }
    
    def initialize_population(self) -> List[Dict]:
        """Initialize random population"""
        return [self.create_individual() for _ in range(self.population_size)]
    
    def fitness(self, individual: Dict, X_train, y_train, X_test, y_test) -> float:
        """Evaluate fitness of individual (accuracy)"""
        model = GradientBoostingClassifier(
            n_estimators=individual['n_estimators'],
            max_depth=individual['max_depth'],
            learning_rate=individual['learning_rate'],
            min_samples_split=individual['min_samples_split'],
            random_state=42
        )
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        return accuracy_score(y_test, y_pred)
    
    def crossover(self, parent1: Dict, parent2: Dict) -> Tuple[Dict, Dict]:
        """Crossover two parents"""
        child1, child2 = parent1.copy(), parent2.copy()
        
        if np.random.random() < self.crossover_rate:
            # Swap random parameters
            for key in child1.keys():
                if np.random.random() < 0.5:
                    child1[key], child2[key] = child2[key], child1[key]
        
        return child1, child2
    
    def mutate(self, individual: Dict) -> Dict:
        """Mutate individual"""
        if np.random.random() < self.mutation_rate:
            key = np.random.choice(list(individual.keys()))
            if key == 'n_estimators':
                individual[key] = np.random.randint(50, 200)
            elif key == 'max_depth':
                individual[key] = np.random.randint(3, 15)
            elif key == 'learning_rate':
                individual[key] = np.random.uniform(0.01, 0.3)
            elif key == 'min_samples_split':
                individual[key] = np.random.randint(2, 10)
        return individual
    
    def select_parents(self, population: List[Dict], fitness_scores: List[float]) -> List[Dict]:
        """Select parents using tournament selection"""
        selected = []
        for _ in range(2):
            tournament = np.random.choice(len(population), size=3, replace=False)
            best_idx = tournament[np.argmax([fitness_scores[i] for i in tournament])]
            selected.append(population[best_idx])
        return selected
    
    def optimize(self, X_train, y_train, X_test, y_test) -> Dict:
        """Run genetic algorithm optimization"""
        population = self.initialize_population()
        best_individual = None
        best_fitness = 0
        
        for generation in range(self.generations):
            # Evaluate fitness
            fitness_scores = [self_fitness(ind, X_train, y_train, X_test, y_test) for ind in population]
            
            # Track best
            current_best_idx = np.argmax(fitness_scores)
            if fitness_scores[current_best_idx] > best_fitness:
                best_fitness = fitness_scores[current_best_idx]
                best_individual = population[current_best_idx]
            
            print(f"Generation {generation}: Best fitness = {best_fitness:.4f}")
            
            # Create new population
            new_population = []
            while len(new_population) < self.population_size:
                parents = self.select_parents(population, fitness_scores)
                child1, child2 = self.crossover(parents[0], parents[1])
                child1 = self.mutate(child1)
                child2 = self.mutate(child2)
                new_population.extend([child1, child2])
            
            population = new_population[:self.population_size]
        
        return best_individual


class SelfLearningSystem:
    """Self-learning system with online updates and performance tracking"""
    
    def __init__(self, ml_engine: MLEngine):
        self.ml_engine = ml_engine
        self.performance_history = []
        self.retrain_threshold = 0.05  # Retrain if accuracy drops by 5%
        self.last_accuracy = 0.0
        
    async def update(self, symbol: str = "BTC_USDT"):
        """Online learning: fetch new data and update if needed"""
        # Fetch recent predictions and their outcomes
        # This would require tracking predictions and their actual results
        
        # For now, just retrain periodically
        results = await self.ml_engine.train(symbol)
        current_accuracy = results['rf']['accuracy']
        
        self.performance_history.append({
            'timestamp': datetime.utcnow(),
            'accuracy': current_accuracy
        })
        
        # Check if retraining is needed
        if self.last_accuracy > 0 and (self.last_accuracy - current_accuracy) > self.retrain_threshold:
            print("Performance drop detected, retraining...")
            await self.ml_engine.train(symbol)
        
        self.last_accuracy = current_accuracy
        
    def get_performance_metrics(self) -> Dict:
        """Get performance metrics"""
        if not self.performance_history:
            return {}
        
        accuracies = [p['accuracy'] for p in self.performance_history]
        return {
            'current_accuracy': accuracies[-1],
            'average_accuracy': np.mean(accuracies),
            'best_accuracy': np.max(accuracies),
            'worst_accuracy': np.min(accuracies),
            'total_updates': len(self.performance_history)
        }