#!/usr/bin/env python3
"""
Deep Learning Engine mit PyTorch LSTM für Lotto-Vorhersagen
============================================================

PyTorch-basierte Implementierung für Python 3.14+ Kompatibilität.
TensorFlow unterstützt Python 3.14 noch nicht, daher verwenden wir PyTorch.

Features:
- LSTM-Netzwerk für zeitliche Sequenzen
- Sequence-to-Probability Mapping
- Feature Engineering Integration
- Model Persistence & Caching
- Hybrid mit RandomForest
"""

import numpy as np
import pandas as pd
import os
import json
import pickle
from datetime import datetime
from typing import Dict, List, Tuple, Optional
import warnings
warnings.filterwarnings('ignore')

# PyTorch Imports
try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import Dataset, DataLoader
    PYTORCH_AVAILABLE = True
except ImportError:
    PYTORCH_AVAILABLE = False
    print("⚠️  PyTorch not available. Install with: pip install torch")


class LSTMModel(nn.Module):
    """PyTorch LSTM Model for Lotto prediction."""

    def __init__(
        self,
        num_features: int,
        num_numbers: int,
        sequence_length: int,
        lstm_units_1: int = 128,
        lstm_units_2: int = 64,
        dense_units: int = 128,
        dropout_rate: float = 0.3
    ):
        super(LSTMModel, self).__init__()

        self.lstm1 = nn.LSTM(
            input_size=num_features,
            hidden_size=lstm_units_1,
            batch_first=True,
            bidirectional=True
        )
        self.dropout1 = nn.Dropout(dropout_rate)
        self.batch_norm1 = nn.BatchNorm1d(lstm_units_1 * 2)

        self.lstm2 = nn.LSTM(
            input_size=lstm_units_1 * 2,
            hidden_size=lstm_units_2,
            batch_first=True,
            bidirectional=True
        )
        self.dropout2 = nn.Dropout(dropout_rate)
        self.batch_norm2 = nn.BatchNorm1d(lstm_units_2 * 2)

        self.fc1 = nn.Linear(lstm_units_2 * 2, dense_units)
        self.batch_norm3 = nn.BatchNorm1d(dense_units)
        self.dropout3 = nn.Dropout(dropout_rate)

        self.fc2 = nn.Linear(dense_units, num_numbers)

    def forward(self, x):
        # LSTM 1
        x, _ = self.lstm1(x)
        x = self.dropout1(x)
        # Take last output
        x = x[:, -1, :]
        x = self.batch_norm1(x)

        # LSTM 2 needs 3D input
        x = x.unsqueeze(1)
        x, _ = self.lstm2(x)
        x = x[:, -1, :]
        x = self.dropout2(x)
        x = self.batch_norm2(x)

        # Dense layers
        x = torch.relu(self.fc1(x))
        x = self.batch_norm3(x)
        x = self.dropout3(x)

        # Output
        x = torch.sigmoid(self.fc2(x))

        return x


class LottoDataset(Dataset):
    """PyTorch Dataset for Lotto sequences."""

    def __init__(self, X, y):
        self.X = torch.FloatTensor(X)
        self.y = torch.FloatTensor(y)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


class DeepLearningEngine:
    """
    PyTorch-based LSTM Deep Learning Engine for Lotto predictions.
    """

    def __init__(
        self,
        num_numbers: int = 49,
        sequence_length: int = 20,
        cache_dir: str = None,
        fast_mode: bool = True
    ):
        """
        Args:
            num_numbers: Maximum number (49 for Lotto, 50 for Eurojackpot)
            sequence_length: How many past draws to consider
            cache_dir: Directory for model persistence
            fast_mode: Use faster training (fewer epochs)
        """
        self.num_numbers = num_numbers
        self.sequence_length = sequence_length
        self.cache_dir = cache_dir
        self.fast_mode = fast_mode

        # Device
        self.device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

        # Model components
        self.model = None
        self.is_trained = False
        self.training_history = {}

        # Configuration
        self.config = {
            'lstm_units_1': 128,
            'lstm_units_2': 64,
            'dense_units': 128,
            'dropout_rate': 0.3,
            'learning_rate': 0.001,
            'batch_size': 32,
            'epochs': 30 if fast_mode else 100,
            'validation_split': 0.2
        }

        if cache_dir:
            os.makedirs(cache_dir, exist_ok=True)
            self.model_path = os.path.join(cache_dir, f'lstm_model_{num_numbers}.pth')
            self.config_path = os.path.join(cache_dir, f'lstm_config_{num_numbers}.json')
        else:
            self.model_path = None
            self.config_path = None

        print(f"🧠 Deep Learning Engine initialized (PyTorch)")
        print(f"   Numbers: 1-{num_numbers}")
        print(f"   Sequence Length: {sequence_length}")
        print(f"   Device: {self.device}")
        print(f"   Fast Mode: {fast_mode}")

    def _build_model(self, num_features: int) -> LSTMModel:
        """Builds LSTM architecture."""
        model = LSTMModel(
            num_features=num_features,
            num_numbers=self.num_numbers,
            sequence_length=self.sequence_length,
            lstm_units_1=self.config['lstm_units_1'],
            lstm_units_2=self.config['lstm_units_2'],
            dense_units=self.config['dense_units'],
            dropout_rate=self.config['dropout_rate']
        )
        return model.to(self.device)

    def _prepare_sequences(
        self,
        df: pd.DataFrame,
        features_df: pd.DataFrame
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Prepares sequences for LSTM training."""
        print(f"   Preparing sequences (length={self.sequence_length})...")

        # Ensure data is sorted by date
        if 'datum' in df.columns:
            df = df.sort_values('datum').reset_index(drop=True)

        # Extract number columns
        num_cols = [col for col in df.columns if col.startswith('Z')]

        # Normalize features to [0, 1]
        features_normalized = features_df.copy()
        for col in features_normalized.columns:
            min_val = features_normalized[col].min()
            max_val = features_normalized[col].max()
            if max_val > min_val:
                features_normalized[col] = (features_normalized[col] - min_val) / (max_val - min_val)
            else:
                features_normalized[col] = 0.5

        X_sequences = []
        y_targets = []

        # Create sequences
        for i in range(self.sequence_length, len(df)):
            # Get sequence of features
            sequence = features_normalized.iloc[i - self.sequence_length:i].values
            X_sequences.append(sequence)

            # Target: next draw as binary vector
            target = np.zeros(self.num_numbers)
            next_draw = df.iloc[i][num_cols].values
            for num in next_draw:
                if 1 <= num <= self.num_numbers:
                    target[int(num) - 1] = 1
            y_targets.append(target)

        X = np.array(X_sequences)
        y = np.array(y_targets)

        print(f"   ✅ Created {len(X)} sequences")
        print(f"   Shape: X={X.shape}, y={y.shape}")

        return X, y

    def train(
        self,
        df: pd.DataFrame,
        features_df: pd.DataFrame,
        force_retrain: bool = False
    ) -> bool:
        """Trains LSTM model on historical data."""
        if not PYTORCH_AVAILABLE:
            print("❌ PyTorch not available")
            return False

        # Check for cached model
        if not force_retrain and self.model_path and os.path.exists(self.model_path):
            print("   📦 Loading cached LSTM model...")
            try:
                checkpoint = torch.load(self.model_path, weights_only=False)
                num_features = checkpoint['num_features']

                self.model = self._build_model(num_features)
                self.model.load_state_dict(checkpoint['model_state_dict'])
                self.model.eval()
                self.is_trained = True

                # Load config
                if os.path.exists(self.config_path):
                    with open(self.config_path, 'r') as f:
                        self.training_history = json.load(f)

                print(f"   ✅ Loaded cached model")
                return True
            except Exception as e:
                print(f"   ⚠️  Failed to load cached model: {e}")
                print("   🔄 Training new model...")

        print(f"\n🧠 TRAINING DEEP LEARNING MODEL (PyTorch LSTM)")
        print("=" * 70)

        # Prepare data
        X, y = self._prepare_sequences(df, features_df)

        if len(X) < 100:
            print("   ⚠️  Not enough data for training (need >100 sequences)")
            return False

        # Split data
        split_idx = int(len(X) * (1 - self.config['validation_split']))
        X_train, X_val = X[:split_idx], X[split_idx:]
        y_train, y_val = y[:split_idx], y[split_idx:]

        # Create datasets
        train_dataset = LottoDataset(X_train, y_train)
        val_dataset = LottoDataset(X_val, y_val)

        train_loader = DataLoader(
            train_dataset,
            batch_size=self.config['batch_size'],
            shuffle=True
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=self.config['batch_size'],
            shuffle=False
        )

        # Build model
        print(f"   Building LSTM architecture...")
        num_features = X.shape[2]
        self.model = self._build_model(num_features)

        # Show summary
        total_params = sum(p.numel() for p in self.model.parameters())
        print(f"\n   📊 Model Summary:")
        print(f"   Total parameters: {total_params:,}")

        # Loss and optimizer
        criterion = nn.BCELoss()
        optimizer = optim.Adam(self.model.parameters(), lr=self.config['learning_rate'])
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', factor=0.5, patience=5
        )

        # Train
        print(f"\n   🚀 Training for {self.config['epochs']} epochs...")
        print(f"   Batch size: {self.config['batch_size']}")
        print(f"   Validation split: {self.config['validation_split']:.1%}")

        best_val_loss = float('inf')
        patience_counter = 0
        patience = 10

        try:
            for epoch in range(self.config['epochs']):
                # Training
                self.model.train()
                train_loss = 0.0
                for batch_X, batch_y in train_loader:
                    batch_X = batch_X.to(self.device)
                    batch_y = batch_y.to(self.device)

                    optimizer.zero_grad()
                    outputs = self.model(batch_X)
                    loss = criterion(outputs, batch_y)
                    loss.backward()
                    optimizer.step()

                    train_loss += loss.item()

                train_loss /= len(train_loader)

                # Validation
                self.model.eval()
                val_loss = 0.0
                with torch.no_grad():
                    for batch_X, batch_y in val_loader:
                        batch_X = batch_X.to(self.device)
                        batch_y = batch_y.to(self.device)

                        outputs = self.model(batch_X)
                        loss = criterion(outputs, batch_y)
                        val_loss += loss.item()

                val_loss /= len(val_loader)

                # Learning rate scheduling
                scheduler.step(val_loss)

                # Print progress every 5 epochs
                if (epoch + 1) % 5 == 0:
                    print(f"   Epoch {epoch+1}/{self.config['epochs']}: "
                          f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

                # Early stopping
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    patience_counter = 0

                    # Save best model
                    if self.model_path:
                        torch.save({
                            'model_state_dict': self.model.state_dict(),
                            'num_features': num_features,
                            'config': self.config
                        }, self.model_path)
                else:
                    patience_counter += 1
                    if patience_counter >= patience:
                        print(f"   Early stopping at epoch {epoch+1}")
                        break

            # Load best model
            if self.model_path and os.path.exists(self.model_path):
                checkpoint = torch.load(self.model_path, weights_only=False)
                self.model.load_state_dict(checkpoint['model_state_dict'])

            # Store training history
            self.training_history = {
                'trained_at': datetime.now().isoformat(),
                'num_samples': len(X),
                'num_features': num_features,
                'final_loss': float(train_loss),
                'final_val_loss': float(val_loss),
                'best_val_loss': float(best_val_loss),
                'epochs_trained': epoch + 1
            }

            self.is_trained = True

            # Save config
            if self.config_path:
                with open(self.config_path, 'w') as f:
                    json.dump(self.training_history, f, indent=2)

            # Print results
            print(f"\n   ✅ TRAINING COMPLETED")
            print(f"   Best Val Loss: {best_val_loss:.4f}")
            print(f"   Epochs: {self.training_history['epochs_trained']}")

            return True

        except Exception as e:
            print(f"\n   ❌ Training failed: {e}")
            import traceback
            traceback.print_exc()
            return False

    def predict(
        self,
        recent_df: pd.DataFrame,
        recent_features: pd.DataFrame
    ) -> Dict[int, float]:
        """Predicts probabilities for each number."""
        if not self.is_trained or self.model is None:
            print("⚠️  Model not trained, returning uniform distribution")
            return {i: 0.5 for i in range(1, self.num_numbers + 1)}

        if len(recent_df) < self.sequence_length:
            print(f"⚠️  Not enough recent data (need {self.sequence_length}, got {len(recent_df)})")
            return {i: 0.5 for i in range(1, self.num_numbers + 1)}

        # Prepare sequence
        recent_features_normalized = recent_features.copy()
        for col in recent_features_normalized.columns:
            min_val = recent_features_normalized[col].min()
            max_val = recent_features_normalized[col].max()
            if max_val > min_val:
                recent_features_normalized[col] = (recent_features_normalized[col] - min_val) / (max_val - min_val)
            else:
                recent_features_normalized[col] = 0.5

        sequence = recent_features_normalized.iloc[-self.sequence_length:].values
        X = torch.FloatTensor(sequence).unsqueeze(0).to(self.device)

        # Predict
        self.model.eval()
        with torch.no_grad():
            predictions = self.model(X)[0].cpu().numpy()

        # Convert to dictionary
        result = {i + 1: float(predictions[i]) for i in range(self.num_numbers)}

        return result

    def get_model_info(self) -> Dict:
        """Returns model information."""
        return {
            'is_trained': self.is_trained,
            'pytorch_available': PYTORCH_AVAILABLE,
            'num_numbers': self.num_numbers,
            'sequence_length': self.sequence_length,
            'device': str(self.device),
            'config': self.config,
            'training_history': self.training_history,
            'model_exists': self.model is not None
        }


class HybridDeepLearningPredictor:
    """Kombiniert RandomForest + LSTM für robustere Vorhersagen."""

    def __init__(
        self,
        dl_engine: DeepLearningEngine,
        rf_weight: float = 0.4,
        dl_weight: float = 0.6
    ):
        """
        Args:
            dl_engine: Deep Learning Engine
            rf_weight: Weight for RandomForest predictions
            dl_weight: Weight for Deep Learning predictions
        """
        self.dl_engine = dl_engine
        self.rf_weight = rf_weight
        self.dl_weight = dl_weight

        print(f"🔀 Hybrid Predictor: RF={rf_weight:.1%} + DL={dl_weight:.1%}")

    def predict(
        self,
        rf_predictions: Dict[int, float],
        recent_df: pd.DataFrame,
        recent_features: pd.DataFrame
    ) -> Dict[int, float]:
        """Combines RandomForest and Deep Learning predictions."""
        # Get DL predictions
        dl_predictions = self.dl_engine.predict(recent_df, recent_features)

        # Combine
        combined = {}
        for num in range(1, self.dl_engine.num_numbers + 1):
            rf_score = rf_predictions.get(num, 0.5)
            dl_score = dl_predictions.get(num, 0.5)

            combined[num] = (
                self.rf_weight * rf_score +
                self.dl_weight * dl_score
            )

        # Normalize to [0, 1]
        min_score = min(combined.values())
        max_score = max(combined.values())
        if max_score > min_score:
            combined = {
                num: (score - min_score) / (max_score - min_score)
                for num, score in combined.items()
            }

        return combined
