import torch

class Logic:
    """
    Pre-built differentiable logic gates.
    """
    @staticmethod
    def implies(condition_idx: int, consequence_idx: int):
        """ Rule: If A is high, B must be high. (A -> B) """
        def loss_fn(latents: torch.Tensor):
            a = torch.sigmoid(latents[:, condition_idx])
            b = torch.sigmoid(latents[:, consequence_idx])
            # Penalize when A=1 and B=0
            return (a * (1.0 - b)).mean()
        return loss_fn

    @staticmethod
    def exclusion(idx_a: int, idx_b: int):
        """ Rule: A and B cannot both be high. (NAND) """
        def loss_fn(latents: torch.Tensor):
            a = torch.sigmoid(latents[:, idx_a])
            b = torch.sigmoid(latents[:, idx_b])
            return (a * b).mean()
        return loss_fn

    @staticmethod
    def max_value(idx: int, limit: float):
        """ Rule: Neuron[idx] value must be <= limit """
        def loss_fn(latents: torch.Tensor):
            val = latents[:, idx]
            return torch.relu(val - limit).mean()
        return loss_fn