import torch
import torch.nn.functional as F 
from abc import ABC, abstractmethod


class TNorm(ABC): 
    """Abstract Base Class for Logic Operations"""

    @abstractmethod 
    def and_op(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        pass 

    @abstractmethod 
    def or_op(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        pass

    @abstractmethod
    def not_op(self, x: torch.Tensor) -> torch.Tensor:
        pass

    @abstractmethod
    def implies_op(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        pass


class ProductTNorm(TNorm): 
    """
    Product Logic (Probabilistic).
    Best for gradients because it's smooth everywhere.
    """
    def and_op(self, x, y): 
        return x * y 
    
    def or_op(self, x, y):
        return x + y - (x * y)
    
    def not_op(self, x): 
        return 1.0 - x
    
    def implies_op(self, x, y): 
        return 1.0 - x + (x * y)
    

class GodelTNorm(TNorm):
    """
    Godel Logic (Min/Mix). 
    Good for 'fuzzy' crisp decisions but has weak gradients (vanishing gradients).
    """
    def and_op(self, x, y):
        return torch.min(x, y)

    def or_op(self, x, y):
        return torch.max(x, y)
    
    def not_op(self, x): 
        return 1.0 - x
    
    def implies_op(self, x, y):
        # Kleene-Dienes: max(1 - A, B)
        return torch.max(1.0 - x, y)
    

class LukasiewiczTNorm(TNorm):
    """
    Bounded Logic. Good for arithmetic constraints.
    """
    def and_op(self, x, y): 
        return torch.max(torch.zeros_like(x), x + y - 1.0)
    
    def or_op(self, x, y): 
        return torch.min(torch.ones_like(x), x + y)
    
    def not_op(self, x): 
        return 1.0 - x
    
    def implies_op(self, x, y): 
        return torch.min(torch.ones_like(x), 1.0 - x + y)
    

def get_logic_engine(logic_type: str) -> TNorm:
    if logic_type == "product": 
        return ProductTNorm()
    elif logic_type == "godel": 
        return GodelTNorm()
    elif logic_type == "lukasiewicz": 
        return LukasiewiczTNorm()
    else: 
        raise ValueError(f"Unknown logic type: {logic_type}")
    
