# Copyright 2025, MetaQuotes Ltd.
# https://www.mql5.com/en/users/johnhlomohang/

import torch.nn as nn
import torch

class EntropyModel(nn.Module):
    """Enhanced model with 8 input features"""
    def __init__(self, dropout_rate=0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(8, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            
            nn.Linear(32, 16),
            nn.BatchNorm1d(16),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            
            nn.Linear(16, 8),
            nn.BatchNorm1d(8),
            nn.ReLU(),
            
            nn.Linear(8, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        return self.net(x)
    
    def predict_with_uncertainty(self, x, n_samples=10):
        """Monte Carlo dropout for prediction uncertainty"""
        self.train()  # Enable dropout
        predictions = []
        with torch.no_grad():
            for _ in range(n_samples):
                pred = self.net(x)
                predictions.append(pred.cpu().numpy())
        
        predictions = np.array(predictions)
        mean_pred = np.mean(predictions, axis=0)
        std_pred = np.std(predictions, axis=0)
        
        self.eval()  # Back to eval mode
        return mean_pred, std_pred