src/models/lstm_branch.py

"""
lstm_branch.py — Bidirectional LSTM temporal encoder for weather + AQI data.

Processes 7-day sequences of weather and air quality features into
a compact temporal representation. Uses a 2-layer bidirectional LSTM
to capture both forward and backward temporal dependencies.
"""

import logging

import torch
import torch.nn as nn

from src.training.config import (
    TIMESERIES_FEATURES, TIMESERIES_WINDOW,
    LSTM_HIDDEN_SIZE, LSTM_NUM_LAYERS, LSTM_BIDIRECTIONAL, LSTM_FEATURE_DIM,
)

logger = logging.getLogger(__name__)


class LSTMBranch(nn.Module):
    """
    Bidirectional LSTM encoder for temporal weather/AQI sequences.

    Architecture:
        Input:  (batch, 7, 6) — 7 days × 6 features
        Output: (batch, 256) — temporal feature vector

    Features: [temperature, humidity, wind_speed, wind_direction,
               precipitation, PM2.5]
    """

    def __init__(
        self,
        input_size: int = TIMESERIES_FEATURES,
        hidden_size: int = LSTM_HIDDEN_SIZE,
        num_layers: int = LSTM_NUM_LAYERS,
        bidirectional: bool = LSTM_BIDIRECTIONAL,
        dropout: float = 0.2,
    ):
        """
        Args:
            input_size: Number of features per timestep (6).
            hidden_size: LSTM hidden state size per direction (128).
            num_layers: Number of stacked LSTM layers (2).
            bidirectional: Whether to use bidirectional LSTM.
            dropout: Dropout between LSTM layers.
        """
        super().__init__()

        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.num_directions = 2 if bidirectional else 1

        # Input projection layer for better feature learning
        self.input_proj = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.LayerNorm(hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

        # Bidirectional LSTM
        self.lstm = nn.LSTM(
            input_size=hidden_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=bidirectional,
            dropout=dropout if num_layers > 1 else 0,
        )

        # Attention mechanism to weight timesteps
        self.attention = nn.Sequential(
            nn.Linear(hidden_size * self.num_directions, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, 1),
        )

        # Output feature dimension
        self.out_features = hidden_size * self.num_directions  # 256

        # Final projection
        self.output_proj = nn.Sequential(
            nn.Linear(self.out_features, self.out_features),
            nn.LayerNorm(self.out_features),
            nn.ReLU(),
        )

        logger.info(
            f"LSTMBranch initialized: {input_size}→{self.out_features}d, "
            f"{num_layers} layers, {'bi' if bidirectional else 'uni'}directional"
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through LSTM encoder.

        Args:
            x: Input tensor of shape (batch, seq_len, features).

        Returns:
            Temporal feature vector of shape (batch, 256).
        """
        batch_size = x.size(0)

        # Project input features
        x = self.input_proj(x)  # (batch, seq_len, hidden_size)

        # LSTM encoding
        lstm_out, (h_n, c_n) = self.lstm(x)
        # lstm_out: (batch, seq_len, hidden_size * num_directions)

        # Attention-weighted pooling over timesteps
        attn_weights = self.attention(lstm_out)  # (batch, seq_len, 1)
        attn_weights = torch.softmax(attn_weights, dim=1)
        context = torch.sum(lstm_out * attn_weights, dim=1)  # (batch, 256)

        # Final projection
        output = self.output_proj(context)  # (batch, 256)

        return output

    def get_attention_weights(self, x: torch.Tensor) -> torch.Tensor:
        """
        Get attention weights for interpretability.

        Returns:
            Attention weights of shape (batch, seq_len).
        """
        x = self.input_proj(x)
        lstm_out, _ = self.lstm(x)
        attn_weights = self.attention(lstm_out).squeeze(-1)
        attn_weights = torch.softmax(attn_weights, dim=1)
        return attn_weights


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    model = LSTMBranch()

    # Test forward pass
    dummy = torch.randn(4, 7, 6)  # 4 samples, 7 days, 6 features
    output = model(dummy)
    print(f"Input: {dummy.shape} → Output: {output.shape}")

    # Test attention weights
    attn = model.get_attention_weights(dummy)
    print(f"Attention weights: {attn.shape}")
    print(f"Attention values: {attn[0].detach().numpy()}")

    # Parameter count
    total = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total:,}")