"""ECA (Efficient Channel Attention) — 残差初始化版本，注册到 ultralytics

关键改进：使用可学习缩放因子 alpha (初始=0)，使得
  output = x * (1 + alpha * (sigmoid(y) - 1))
  当 alpha=0 时 output=x（恒等映射），不破坏预训练特征。
"""
import math
import torch
import torch.nn as nn


class ECA(nn.Module):
    def __init__(self, c1=None, c2=None, gamma=2, b=1):
        super().__init__()
        self.gamma = gamma
        self.b = b
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.sigmoid = nn.Sigmoid()
        self.alpha = nn.Parameter(torch.zeros(1))
        self._conv = None
        if c1 is not None:
            self._init_conv(c1)

    def _init_conv(self, channels):
        t = int(abs((math.log2(channels) + self.b) / self.gamma))
        k = t if t % 2 else t + 1
        self._conv = nn.Conv1d(1, 1, kernel_size=k, padding=k // 2, bias=False)

    def forward(self, x):
        if self._conv is None:
            self._init_conv(x.shape[1])
            self._conv = self._conv.to(x.device)
        y = self.avg_pool(x)
        y = self._conv(y.squeeze(-1).transpose(-1, -2))
        y = y.transpose(-1, -2).unsqueeze(-1)
        return x * (1.0 + self.alpha * (self.sigmoid(y) - 1.0))


def register_eca():
    import ultralytics.nn.tasks as tasks
    if "ECA" not in vars(tasks):
        tasks.ECA = ECA