import torch
import torch.nn as nn
import math

# 维度：68 points * 2 coordinates (x, y) = 136
NUM_LANDMARKS = 136 
BATCH_SIZE = 32768 
SHAPE = (BATCH_SIZE, NUM_LANDMARKS)

W_VAL = 10.0
EPS_VAL = 2.0

class WingLoss(nn.Module):
    def __init__(self, w=10.0, epsilon=2.0, reduction='mean'):
        super(WingLoss, self).__init__()
        self.w = w
        self.epsilon = epsilon
        self.reduction = reduction
        # Constant C: w - w * ln(1 + w/epsilon)
        self.C = self.w - self.w * math.log(1 + self.w / self.epsilon)

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        delta = (target - pred).abs()
        # Case 1: Small errors (|x| < w) -> ln form
        loss_small = self.w * torch.log(1 + delta / self.epsilon)
        # Case 2: Large errors (|x| >= w) -> linear form
        loss_large = delta - self.C
        # Combine
        loss = torch.where(delta < self.w, loss_small, loss_large)

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, w=10.0, epsilon=2.0, reduction='mean'):
        super(Model, self).__init__()
        self.loss_fn = WingLoss(w=w, epsilon=epsilon, reduction=reduction)
    
    def forward(self, pred, target):
        return self.loss_fn(pred, target)

def get_inputs():
    
    pred = torch.randn(SHAPE, dtype=torch.float32)
    target = torch.randn(SHAPE, dtype=torch.float32)
    
    return [pred.contiguous(), target.contiguous()]

def get_init_inputs():
    return [W_VAL, EPS_VAL, 'none']