"""
Recurrent Neural Network Model
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: RNN/LSTM implementation for sequence modeling
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class SimpleRNN(nn.Module):
    """
    Simple Recurrent Neural Network using LSTM
    
    An RNN architecture using LSTM cells for sequence modeling tasks.
    Demonstrates PyTorch's RNN modules and sequence processing.
    """
    
    def __init__(self, input_size, hidden_size=64, num_layers=2, num_classes=10, dropout=0.2):
        """
        Initialize the RNN
        
        Args:
            input_size: Number of features in input sequence
            hidden_size: Number of hidden units in LSTM
            num_layers: Number of LSTM layers
            num_classes: Number of output classes
            dropout: Dropout probability
        """
        super(SimpleRNN, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # LSTM layer
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=False
        )
        
        # Fully connected layer
        self.fc = nn.Linear(hidden_size, num_classes)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        """
        Forward pass through the RNN
        
        Args:
            x: Input tensor of shape (batch_size, seq_len, input_size)
            
        Returns:
            Output tensor of shape (batch_size, num_classes)
        """
        # LSTM forward pass
        # Output shape: (batch_size, seq_len, hidden_size)
        lstm_out, (hidden, cell) = self.lstm(x)
        
        # Use the last output from the sequence
        # Take the last time step's output
        last_output = lstm_out[:, -1, :]
        
        # Apply dropout
        last_output = self.dropout(last_output)
        
        # Fully connected layer
        output = self.fc(last_output)
        
        return output
    
    def predict(self, x):
        """
        Make predictions on input data
        
        Args:
            x: Input tensor
            
        Returns:
            Predicted class indices
        """
        self.eval()
        with torch.no_grad():
            outputs = self.forward(x)
            _, predicted = torch.max(outputs.data, 1)
        return predicted


class SimpleGRU(nn.Module):
    """
    Simple Gated Recurrent Unit Network
    
    An alternative RNN architecture using GRU cells.
    """
    
    def __init__(self, input_size, hidden_size=64, num_layers=2, num_classes=10, dropout=0.2):
        """
        Initialize the GRU
        
        Args:
            input_size: Number of features in input sequence
            hidden_size: Number of hidden units in GRU
            num_layers: Number of GRU layers
            num_classes: Number of output classes
            dropout: Dropout probability
        """
        super(SimpleGRU, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        # GRU layer
        self.gru = nn.GRU(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=False
        )
        
        # Fully connected layer
        self.fc = nn.Linear(hidden_size, num_classes)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        """
        Forward pass through the GRU
        
        Args:
            x: Input tensor of shape (batch_size, seq_len, input_size)
            
        Returns:
            Output tensor of shape (batch_size, num_classes)
        """
        # GRU forward pass
        gru_out, hidden = self.gru(x)
        
        # Use the last output from the sequence
        last_output = gru_out[:, -1, :]
        
        # Apply dropout
        last_output = self.dropout(last_output)
        
        # Fully connected layer
        output = self.fc(last_output)
        
        return output