# Neural Net Wrapper
import torch 
import torch.nn as nn 
from typing import Callable, Union 


class NeuralConcept(nn.Module): 
    """
    Acts as a bridge between a Logic Predicate and a Neural Network.

    Example: 
        predicate = Predicate("Digit")
        net = MyCNN()
        concept = NeuralConcept(predicate, net, output_index=3)
        # This means: Digit(x) is True if MyCNN(x)'s 3rd output is high. 
    """
    def __init__(self, name: str, module: nn.Module, output_index: int = None): 
        super().__init__()
        self.name = name 
        self.module = module 
        self.output_index = output_index 

    def forward(self, x: torch.Tensor) -> torch.Tensor: 
        """
        Runs the neural net and extracts probability for this concept.
        """
        # Get raw output (Logic or Probs)
        output = self.module(x)

        # If specific index is needed (e.g., Class 3 is MNIST)
        if self.output_index is not None:
            # Ensure output is (Batch, Classes)
            if output.dim() > 1:
                return output[:, self.output_index]
            else: 
                return output[self.output_index]
        
        # If binary classifier (Single Output)
        return output.squeeze()
    
    def __repr__(self):
        return f"NeuralConcept({self.name})"
    

        