"""Neural + tree hybrid ensemble predictor.

Stacks gradient-boosted trees with a Transformer/BiLSTM network and a Bayesian
network (Monte-Carlo dropout) for uncertainty-aware signal filtering.
"""
import numpy as np
import pandas as pd
from sklearn.preprocessing import RobustScaler
from sklearn.model_selection import TimeSeriesSplit
import xgboost as xgb
import lightgbm as lgb
from catboost import CatBoostClassifier
from hmmlearn import hmm
import talib
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader


class TimeSeriesDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

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

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

class TransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.3):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + self.dropout(attn_out))
        ff_out = self.ff(x)
        x = self.norm2(x + self.dropout(ff_out))
        return x

class NeuralEnsemblePredictor:
    def __init__(self, forecast_period=12, confidence_threshold=0.6, device='cuda' if torch.cuda.is_available() else 'cpu'):
        self.forecast_period = forecast_period
        self.confidence_threshold = confidence_threshold
        self.device = device
        self.scaler = RobustScaler()
        self.regime_model = hmm.GaussianHMM(n_components=3, random_state=42)

        # Tree-based models with pre-tuned parameters
        self.models = {
            'xgboost': xgb.XGBClassifier(
                colsample_bytree=0.9821911945239713,
                learning_rate=0.07305593001592295,
                max_depth=9,
                min_child_weight=5,
                subsample=0.7524259059189972,
                random_state=42,
                eval_metric='logloss'
            ),
            'lightgbm': lgb.LGBMClassifier(
                feature_fraction=0.9249583953429453,
                learning_rate=0.025468440525690465,
                max_depth=7,
                min_child_samples=41,
                subsample=0.8092209312217534,
                random_state=44
            ),
            'catboost': CatBoostClassifier(
                colsample_bylevel=0.7762414086064304,
                depth=4,
                learning_rate=0.08067885793496102,
                subsample=0.858134636983408,
                random_state=45,
                verbose=0
            )
        }

        # PyTorch neural networks
        self.nn_model = self._build_neural_network().to(self.device)
        self.bnn_model = self._build_bayesian_network().to(self.device)

        # Non-linear meta-model
        self.meta_model = lgb.LGBMClassifier(
            n_estimators=100, max_depth=3, learning_rate=0.05, random_state=46
        )

    def _build_neural_network(self):
        """Transformer-based neural network"""
        class Net(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = nn.ModuleList([
                    TransformerBlock(d_model=6, nhead=2, dropout=0.3)
                    for _ in range(2)
                ])
                self.lstm = nn.LSTM(6, 32, batch_first=True, bidirectional=True)
                self.attn = nn.Linear(64, 1)  # Attention over 64 from bidirectional LSTM
                self.fc = nn.Sequential(
                    nn.Linear(64, 64),
                    nn.ReLU(),
                    nn.Dropout(0.4),
                    nn.Linear(64, 2)
                )

            def forward(self, x):
                for t in self.transformer:
                    x = t(x)
                x, _ = self.lstm(x)  # Shape: (batch, 20, 64)
                attn_weights = torch.softmax(self.attn(x), dim=1)  # Shape: (batch, 20, 1)
                x = (x * attn_weights).sum(dim=1)  # Shape: (batch, 64)
                x = self.fc(x)
                return x

        return Net()

    def _build_bayesian_network(self):
        """Bayesian neural network with Monte Carlo dropout"""
        class BNN(nn.Module):
            def __init__(self):
                super().__init__()
                self.flatten = nn.Flatten()
                self.fc1 = nn.Linear(20 * 6, 64)
                self.fc2 = nn.Linear(64, 32)
                self.fc3 = nn.Linear(32, 2)
                self.dropout = nn.Dropout(0.3)
                self.relu = nn.ReLU()

            def forward(self, x, training=False):
                x = self.flatten(x)
                x = self.relu(self.fc1(x))
                x = self.dropout(x) if training else x
                x = self.relu(self.fc2(x))
                x = self.dropout(x) if training else x
                x = self.fc3(x)
                return x

        return BNN()

    def detect_market_regime(self, data):
        returns = np.log(data['close'] / data['close'].shift(1))
        volatility = returns.rolling(window=20).std()
        combined = pd.DataFrame({'returns': returns, 'volatility': volatility}).dropna()
        self.regime_model.fit(combined.values)
        regimes = self.regime_model.predict(combined.values)
        regime_series = pd.Series(index=data.index, dtype='float64')
        regime_series.iloc[len(data)-len(regimes):] = regimes
        return regime_series

    def create_advanced_features(self, df):
        data = df.copy()
        data['market_regime'] = self.detect_market_regime(data)
        for period in [21, 55]:
            data[f'ema_{period}'] = talib.EMA(data['close'], timeperiod=period)
            data[f'trend_{period}'] = (data[f'ema_{period}'] - data[f'ema_{period}'].shift(period)) / data[f'ema_{period}'].shift(period)
        data['atr_ratio'] = talib.ATR(data['high'], data['low'], data['close'], 14) / data['close']
        data['rsi'] = talib.RSI(data['close'], 14)
        data['volume_ma'] = talib.EMA(data['volume'], timeperiod=20)
        data['volume_ratio'] = data['volume'] / data['volume_ma']

        returns = data['close'].shift(-self.forecast_period) / data['close'] - 1
        data['target'] = np.where(returns > 0.005, 1, np.where(returns < -0.005, 0, None))
        return data.dropna()

    def prepare_features(self, data, for_nn=False):
        feature_columns = [
            'market_regime', 'atr_ratio', 'rsi', 'volume_ratio',
            'trend_21', 'trend_55'
        ]
        X = data[feature_columns]
        y = data['target'].astype(int)

        if for_nn:
            X_3d = np.array([X.iloc[i-20:i].values for i in range(20, len(X))])
            y_3d = y.iloc[20:].values
            return X_3d, y_3d
        return X, y

    def train_nn(self, X_nn, y_nn, model, epochs=50, batch_size=32):
        dataset = TimeSeriesDataset(X_nn, y_nn)
        train_size = int(0.8 * len(dataset))
        val_size = len(dataset) - train_size
        train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=batch_size)

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=0.0001)

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

        for epoch in range(epochs):
            model.train()
            train_loss = 0
            for X_batch, y_batch in train_loader:
                X_batch, y_batch = X_batch.to(self.device), y_batch.to(self.device)
                optimizer.zero_grad()
                outputs = model(X_batch)
                loss = criterion(outputs, y_batch)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()

            model.eval()
            val_loss = 0
            with torch.no_grad():
                for X_batch, y_batch in val_loader:
                    X_batch, y_batch = X_batch.to(self.device), y_batch.to(self.device)
                    outputs = model(X_batch)
                    val_loss += criterion(outputs, y_batch).item()

            train_loss /= len(train_loader)
            val_loss /= len(val_loader)
            print(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                patience_counter = 0
                torch.save(model.state_dict(), 'nn_best.pth')
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    print("Early stopping")
                    break

        model.load_state_dict(torch.load('nn_best.pth'))
        return model

    def fit(self, train_data):
        processed_data = self.create_advanced_features(train_data)
        X, y = self.prepare_features(processed_data)
        X_nn, y_nn = self.prepare_features(processed_data, for_nn=True)
        X_scaled = self.scaler.fit_transform(X)
        X_scaled = pd.DataFrame(X_scaled, columns=X.columns)

        # Train tree-based models (no tuning)
        tscv = TimeSeriesSplit(n_splits=5)
        oof_preds = np.zeros((len(X_scaled), len(self.models) + 2))

        for fold, (train_idx, val_idx) in enumerate(tscv.split(X_scaled)):
            X_train, X_val = X_scaled.iloc[train_idx], X_scaled.iloc[val_idx]
            y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

            for i, (name, model) in enumerate(self.models.items()):
                if name == 'xgboost':
                    model.fit(X_train, y_train, eval_set=[(X_val, y_val)],verbose=False)
                elif name == 'lightgbm':
                    model.fit(X_train, y_train, eval_set=[(X_val, y_val)])
                elif name == 'catboost':
                    model.fit(X_train, y_train, eval_set=(X_val, y_val))
                oof_preds[val_idx, i] = model.predict_proba(X_val)[:, 1]

        # Train neural network
        nn_idx_shift = len(X_scaled) - len(X_nn)
        self.nn_model = self.train_nn(X_nn, y_nn, self.nn_model)
        self.nn_model.eval()
        with torch.no_grad():
            nn_preds = torch.softmax(self.nn_model(torch.tensor(X_nn, dtype=torch.float32).to(self.device)), dim=1)[:, 1].cpu().numpy()
        oof_preds[nn_idx_shift:, len(self.models)] = nn_preds

        # Train Bayesian NN with Monte Carlo dropout
        self.bnn_model = self.train_nn(X_nn, y_nn, self.bnn_model)
        self.bnn_model.train()  # Enable dropout for MC estimation
        mc_preds = []
        with torch.no_grad():
            for _ in range(10):  # 10 Monte Carlo samples
                preds = torch.softmax(self.bnn_model(torch.tensor(X_nn, dtype=torch.float32).to(self.device), training=True), dim=1)[:, 1].cpu().numpy()
                mc_preds.append(preds)
        bnn_mean = np.mean(mc_preds, axis=0)
        bnn_std = np.std(mc_preds, axis=0)
        oof_preds[nn_idx_shift:, len(self.models) + 1] = bnn_mean

        # Train meta-model
        self.meta_model.fit(oof_preds, y, sample_weight=np.exp(np.linspace(-1, 0, len(y))))

        # Final training of tree models
        for name, model in self.models.items():
            if name in ['xgboost']:
                model.fit(X_scaled, y, eval_set=[(X_scaled, y)], verbose=False)
            elif name in ['lightgbm']:
                model.fit(X_scaled, y, eval_set=[(X_scaled, y)])
            else:
                model.fit(X_scaled, y, eval_set=(X_scaled, y))

        return processed_data

    def predict(self, data):
        processed_data = self.create_advanced_features(data)
        X, _ = self.prepare_features(processed_data)
        X_nn, _ = self.prepare_features(processed_data, for_nn=True)
        X_scaled = self.scaler.transform(X)
        X_scaled = pd.DataFrame(X_scaled, columns=X.columns)

        base_preds = np.zeros((len(X_scaled), len(self.models) + 2))
        for i, (name, model) in enumerate(self.models.items()):
            base_preds[:, i] = model.predict_proba(X_scaled)[:, 1]

        nn_idx_shift = len(X_scaled) - len(X_nn)
        self.nn_model.eval()
        with torch.no_grad():
            nn_preds = torch.softmax(self.nn_model(torch.tensor(X_nn, dtype=torch.float32).to(self.device)), dim=1)[:, 1].cpu().numpy()
        base_preds[nn_idx_shift:, len(self.models)] = nn_preds

        # Bayesian NN predictions with uncertainty
        self.bnn_model.train()  # Enable dropout
        mc_preds = []
        with torch.no_grad():
            for _ in range(10):
                preds = torch.softmax(self.bnn_model(torch.tensor(X_nn, dtype=torch.float32).to(self.device), training=True), dim=1)[:, 1].cpu().numpy()
                mc_preds.append(preds)
        bnn_mean = np.mean(mc_preds, axis=0)
        bnn_std = np.std(mc_preds, axis=0)
        base_preds[nn_idx_shift:, len(self.models) + 1] = bnn_mean

        # Meta-model prediction
        meta_proba = self.meta_model.predict_proba(base_preds)

        signals = pd.Series(0, index=processed_data.index)
        valid_indices = processed_data.index[20:]  # Skip first 20 due to lookback
        meta_proba_valid = meta_proba[nn_idx_shift:]  # Align with NN predictions

        long_mask = (meta_proba_valid[:, 1] > self.confidence_threshold) & (bnn_std < 0.2)
        short_mask = (meta_proba_valid[:, 0] > self.confidence_threshold) & (bnn_std < 0.2)
        signals.loc[valid_indices[long_mask]] = 1
        signals.loc[valid_indices[short_mask]] = -1

        return signals