# Path: difflogic/nn/neural_logic/modules/_utils.py
# def exclude_mask(inputs, cnt=2, dim=1):
# def mask_value(inputs, mask, value):


import torch
import torch.nn as nn
from jacinle.utils.enum import JacEnum
from ._utils import meshgrid, meshgrid_exclude_self
class InputTransformMethod(JacEnum):
  CONCAT = 'concat'
  DIFF = 'diff'
  CMP = 'cmp'


class InputTransform(nn.Module):
  """Transform the unary predicates to binary predicates by operations."""

  def __init__(self, method, exclude_self=True):
    super().__init__()
    self.method = InputTransformMethod.from_string(method)
    self.exclude_self = exclude_self

  def forward(self, inputs):
    assert inputs.dim() == 3

    x, y = meshgrid(inputs, dim=1)

    if self.method is InputTransformMethod.CONCAT:
      combined = torch.cat((x, y), dim=3)
    elif self.method is InputTransformMethod.DIFF:
      combined = x - y
    elif self.method is InputTransformMethod.CMP:
      combined = torch.cat([x < y, x == y, x > y], dim=3)
    else:
      raise ValueError('Unknown input transform method: {}.'.format(
          self.method))

    if self.exclude_self:

# Path: difflogic/nn/neural_logic/modules/_utils.py
# def exclude_mask(inputs, cnt=2, dim=1):
# def mask_value(inputs, mask, value):


# Path: difflogic/nn/neural_logic/modules/input_transform.py
import torch
import torch.nn as nn
from jacinle.utils.enum import JacEnum
from ._utils import meshgrid, meshgrid_exclude_self

class InputTransformMethod(JacEnum):
  CONCAT = 'concat'
  DIFF = 'diff'
  CMP = 'cmp'


class InputTransform(nn.Module):
  """Transform the unary predicates to binary predicates by operations."""

  def __init__(self, method, exclude_self=True):
    super().__init__()
    self.method = InputTransformMethod.from_string(method)
    self.exclude_self = exclude_self

  def forward(self, inputs):
    assert inputs.dim() == 3

    x, y = meshgrid(inputs, dim=1)

    if self.method is InputTransformMethod.CONCAT:
      combined = torch.cat((x, y), dim=3)
    elif self.method is InputTransformMethod.DIFF:
      combined = x - y
    elif self.method is InputTransformMethod.CMP:
      combined = torch.cat([x < y, x == y, x > y], dim=3)
    else:
      raise ValueError('Unknown input transform method: {}.'.format(
          self.method))

    if self.exclude_self:

      combined = meshgrid_exclude_self(combined, dim=1)