"""
Basic Neural Network Model
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Basic feedforward neural network implementation
"""

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


class BasicNeuralNetwork(nn.Module):
    """
    Basic Feedforward Neural Network
    
    A simple multi-layer perceptron with configurable hidden layers.
    Demonstrates PyTorch's dynamic computation graph and automatic differentiation.
    """
    
    def __init__(self, input_size, hidden_size=64, output_size=10, num_layers=2, dropout=0.2):
        """
        Initialize the neural network
        
        Args:
            input_size: Number of input features
            hidden_size: Number of neurons in hidden layers
            output_size: Number of output classes
            num_layers: Number of hidden layers
            dropout: Dropout probability for regularization
        """
        super(BasicNeuralNetwork, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        
        # Create layers dynamically
        layers = []
        
        # Input layer
        layers.append(nn.Linear(input_size, hidden_size))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(dropout))
        
        # Hidden layers
        for _ in range(num_layers - 1):
            layers.append(nn.Linear(hidden_size, hidden_size))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
        
        # Output layer
        layers.append(nn.Linear(hidden_size, output_size))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, x):
        """
        Forward pass through the network
        
        Args:
            x: Input tensor of shape (batch_size, input_size)
            
        Returns:
            Output tensor of shape (batch_size, output_size)
        """
        # Flatten input if needed
        if x.dim() > 2:
            x = x.view(x.size(0), -1)
        
        return self.network(x)
    
    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