import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Instance Normalization followed by a Minkowski Transform.
    This version uses standard PyTorch operations for a fair baseline.
    """
    def __init__(self, num_features=64, p=2, eps=1e-5):
        super(Model, self).__init__()
        self.p = p
        if p <= 0:
            raise ValueError("p must be positive")
        
        # InstanceNorm2d operates on [N, C, H, W]
        self.instance_norm = nn.InstanceNorm2d(num_features, eps=eps, affine=False, track_running_stats=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Compute the InstanceNorm of x, then compute the Minkowski transform of the result.
        Input: [N, C, H, W]
        Output: [N, C]
        """
        # Input validation
        if x.dim() != 4:
            raise ValueError(f"Input tensor must be 4D, got {x.dim()}D")
        
        # Step 1: Apply Instance Normalization
        # This will now work correctly as H*W > 1
        normed_x = self.instance_norm(x) # Shape: [N, C, H, W]
        
        # Step 2: Compute Minkowski transform per channel on the normalized features
        N, C, H, W = normed_x.shape
        normed_x_flat = normed_x.view(N, C, H * W)
        abs_normed_flat = torch.abs(normed_x_flat)
        
        if self.p == 1:
            # L1 norm
            minkowski_vals = torch.sum(abs_normed_flat, dim=2) # Shape: [N, C]
        elif self.p == 2:
            # L2 norm
            minkowski_vals = torch.sqrt(torch.sum(abs_normed_flat ** 2, dim=2)) # Shape: [N, C]
        else:
            # General Lp norm
            minkowski_vals = torch.pow(torch.sum(torch.pow(abs_normed_flat, self.p), dim=2), 1.0/self.p) # Shape: [N, C]
        
        return minkowski_vals

# 参数配置
batch_size = 32
num_features = 64
height = 128
width = 128

def get_inputs():
    x = torch.randn(batch_size, num_features, height, width)
    return [x]

def get_init_inputs():
    return [num_features, 2] # num_features and p value
