"""
cmapss_model.py
===============
CNN-BiLSTM-Attention model with MC Dropout for RUL prediction.
Used for both C-MAPSS and AGTF30 turbine blade datasets.

Author: Mohammed Bello Sani
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class CNNBiLSTMAttention(nn.Module):
    """
    CNN + Bidirectional LSTM + Self-Attention + MC Dropout
    """
    def __init__(self, hp):
        super().__init__()
        self.hp = hp
        n_features = hp['n_features']
        cnn_filters = hp['cnn_filters']
        kernel_size = hp.get('cnn_kernel', 3)
        lstm_hidden = hp['lstm_hidden']
        lstm_layers = hp['lstm_layers']
        fc_hidden = hp.get('fc_hidden', 64)
        dropout = hp.get('mc_dropout', 0.5)

        # CNN feature extractor
        cnn_layers = []
        in_channels = n_features
        for i, filters in enumerate(cnn_filters):
            cnn_layers.append(nn.Conv1d(in_channels, filters, kernel_size, padding='same'))
            cnn_layers.append(nn.BatchNorm1d(filters))
            cnn_layers.append(nn.GELU())
            cnn_layers.append(nn.Dropout(dropout * 0.5))  # lighter dropout in CNN
            in_channels = filters
        self.cnn = nn.Sequential(*cnn_layers)
        self.cnn_out_dim = cnn_filters[-1]

        # BiLSTM
        self.lstm = nn.LSTM(
            input_size=self.cnn_out_dim,
            hidden_size=lstm_hidden,
            num_layers=lstm_layers,
            batch_first=True,
            bidirectional=True,
            dropout=hp.get('lstm_dropout', 0.3) if lstm_layers > 1 else 0
        )
        lstm_out_dim = lstm_hidden * 2  # bidirectional

        # Self-Attention
        self.attn_heads = hp.get('attn_heads', 4)
        self.attn = nn.MultiheadAttention(
            embed_dim=lstm_out_dim,
            num_heads=self.attn_heads,
            dropout=dropout * 0.3,
            batch_first=True
        )
        self.attn_norm = nn.LayerNorm(lstm_out_dim)

        # Fully connected regressor
        self.fc = nn.Sequential(
            nn.Linear(lstm_out_dim, fc_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(fc_hidden, 1)
        )

        # MC Dropout layers are already included; we keep them active during inference
        self.dropout_rate = dropout

    def forward(self, x):
        # x shape: (batch, seq_len, features)
        # Permute for CNN: (batch, features, seq_len)
        x_cnn = x.permute(0, 2, 1)
        x_cnn = self.cnn(x_cnn)            # (batch, filters, seq_len)
        x_cnn = x_cnn.permute(0, 2, 1)     # (batch, seq_len, filters)

        # BiLSTM
        lstm_out, _ = self.lstm(x_cnn)     # (batch, seq_len, lstm_hidden*2)

        # Self-Attention
        attn_out, _ = self.attn(lstm_out, lstm_out, lstm_out)
        attn_out = self.attn_norm(attn_out + lstm_out)  # residual

        # Global average pooling over time dimension
        pooled = attn_out.mean(dim=1)       # (batch, lstm_out_dim)

        # Final prediction
        out = self.fc(pooled).squeeze(-1)   # (batch,)
        return out

    def predict_with_uncertainty(self, x, n_samples=100):
        """
        Monte Carlo Dropout inference.
        Returns mean and standard deviation.
        """
        self.train()  # keep dropout active
        preds = []
        with torch.no_grad():
            for _ in range(n_samples):
                preds.append(self.forward(x).cpu().numpy())
        preds = np.stack(preds, axis=0)
        mean = preds.mean(axis=0)
        std = preds.std(axis=0)
        return mean, std


# Compatibility: if the old code expects the model to be loaded via torch.load,
# we need to ensure the class is defined in this module.
if __name__ == '__main__':
    print("CNNBiLSTMAttention model definition ready.")