import math
import torch
import torch.nn as nn


class PositionalEncoding(nn.Module):
    """
    正弦位置编码
    输入:  [batch, seq_len, d_model]
    输出:  [batch, seq_len, d_model]
    """
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)

        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float32)
            * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)

        # 如果 d_model 是奇数，这里要避免维度不匹配
        if d_model % 2 == 0:
            pe[:, 1::2] = torch.cos(position * div_term)
        else:
            pe[:, 1::2] = torch.cos(position * div_term[:-1])

        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer("pe", pe)

    def forward(self, x):
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len, :]
        return self.dropout(x)


class TransformerEncoderForecaster(nn.Module):
    """
    Transformer Encoder 时间序列预测模型

    输入:
        x: [batch, input_len, n_features]

    输出:
        y_pred: [batch, pred_len]
    """
    def __init__(
        self,
        n_features,
        pred_len,
        d_model=128,
        n_heads=8,
        num_layers=3,
        dim_feedforward=256,
        dropout=0.1,
        pooling="last",   # 可选: "last" 或 "mean"
    ):
        super().__init__()

        self.pred_len = pred_len
        self.pooling = pooling

        # 将原始多变量特征映射到 Transformer 的隐藏维度
        self.input_projection = nn.Linear(n_features, d_model)

        # 位置编码
        self.pos_encoder = PositionalEncoding(
            d_model=d_model,
            max_len=5000,
            dropout=dropout,
        )

        # Transformer Encoder 层
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )

        self.encoder = nn.TransformerEncoder(
            encoder_layer=encoder_layer,
            num_layers=num_layers,
        )

        # 输出层：将编码后的历史信息映射到未来 pred_len 个预测值
        self.head = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(128, pred_len),
        )

    def forward(self, x):
        # x: [batch, input_len, n_features]

        x = self.input_projection(x)      # [batch, input_len, d_model]
        x = self.pos_encoder(x)           # [batch, input_len, d_model]
        x = self.encoder(x)               # [batch, input_len, d_model]

        if self.pooling == "mean":
            x = x.mean(dim=1)             # [batch, d_model]
        else:
            x = x[:, -1, :]               # [batch, d_model]

        y_pred = self.head(x)             # [batch, pred_len]

        return y_pred

if __name__ == "__main__":
    model = TransformerEncoderForecaster(n_features=10, pred_len=288, d_model=128, n_heads=8, num_layers=3, dim_feedforward=256, dropout=0.1, pooling="last")
    print(model)
    # 构造随机输入数据: [batch, input_len, n_features]
    import torch
    batch_size = 4
    input_len = 288
    n_features = 10
    x = torch.randn(batch_size, input_len, n_features)
    # 前向传播
    y_pred = model(x)
    print("输出 shape:", y_pred.shape)