import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Bray-Curtis Distance implementation.
    Computes the Bray-Curtis distance between two sets of vectors.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Compute Bray-Curtis distance between x and y.

        Args:
            x (torch.Tensor): First set of vectors [batch_size, feature_dim]
            y (torch.Tensor): Second set of vectors [batch_size, feature_dim]

        Returns:
            torch.Tensor: Bray-Curtis distances [batch_size]
        """
        # Input validation
        if x.shape != y.shape:
            raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
        
        if x.dim() != 2:
            raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
        
        # Compute Bray-Curtis distance: Σ|x_i - y_i| / Σ(|x_i| + |y_i|)
        # Step 1: Compute absolute differences
        diff = torch.abs(x - y)
        
        # Step 2: Compute numerator (sum of absolute differences)
        numerator = torch.sum(diff, dim=1)
        
        # Step 3: Compute denominator (sum of absolute values)
        denominator = torch.sum(torch.abs(x) + torch.abs(y), dim=1)
        
        # Step 4: Handle division by zero
        # When denominator is 0 (both vectors are all zeros), distance is 0
        distance = torch.where(denominator > 0, numerator / denominator, torch.zeros_like(numerator))
        
        return distance

batch_size = 512
feature_dim = 512

def get_inputs():
    # Generate two sets of positive vectors (ecological data is typically non-negative)
    x = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
    y = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
    return [x, y]

def get_init_inputs():
    return []  # No special initialization inputs needed
