import torch.nn as nn
import torch

class CNNLSTM(nn.Module):
    def __init__(
        self,
        n_features,
        pred_len,
        cnn_channels=64,
        lstm_hidden=128,
        lstm_layers=2,
        dropout=0.2,
    ):
        super().__init__()

        # Conv1d 输入：[batch, n_features, input_len]
        self.cnn = nn.Sequential(
            nn.Conv1d(
                in_channels=n_features,
                out_channels=cnn_channels,
                kernel_size=3,
                padding=1,
            ),
            nn.BatchNorm1d(cnn_channels),
            nn.ReLU(),
            nn.Dropout(dropout),

            nn.Conv1d(
                in_channels=cnn_channels,
                out_channels=cnn_channels,
                kernel_size=3,
                padding=1,
            ),
            nn.BatchNorm1d(cnn_channels),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

        # LSTM 输入：[batch, input_len, cnn_channels]
        self.lstm = nn.LSTM(
            input_size=cnn_channels,
            hidden_size=lstm_hidden,
            num_layers=lstm_layers,
            batch_first=True,
            dropout=dropout if lstm_layers > 1 else 0.0,
        )

        self.head = nn.Sequential(
            nn.Linear(lstm_hidden, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, pred_len),
        )

    def forward(self, x):
        # x: [batch, input_len, n_features]
        x = x.permute(0, 2, 1)       # -> [batch, n_features, input_len]
        x = self.cnn(x)              # -> [batch, cnn_channels, input_len]
        x = x.permute(0, 2, 1)       # -> [batch, input_len, cnn_channels]

        lstm_out, _ = self.lstm(x)   # -> [batch, input_len, lstm_hidden]
        last_hidden = lstm_out[:, -1, :]

        y_pred = self.head(last_hidden)
        return y_pred

if __name__ == "__main__":
    model = CNNLSTM(n_features=10, pred_len=288, cnn_channels=64, lstm_hidden=32, lstm_layers=2, dropout=0.2)

    # 构造随机输入数据: [batch, input_len, n_features]
    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)