import torch
import torch.nn as nn
from model.MegaCRN.MegaCRN import MegaCRN

class MegaCRNModel(nn.Module):
    def __init__(self, args):
        super(MegaCRNModel, self).__init__()
        
        # 设置默认参数
        if 'rnn_units' not in args:
            args['rnn_units'] = 64
        if 'num_layers' not in args:
            args['num_layers'] = 1
        if 'cheb_k' not in args:
            args['cheb_k'] = 3
        if 'ycov_dim' not in args:
            args['ycov_dim'] = 1
        if 'mem_num' not in args:
            args['mem_num'] = 20
        if 'mem_dim' not in args:
            args['mem_dim'] = 64
        if 'cl_decay_steps' not in args:
            args['cl_decay_steps'] = 2000
        if 'use_curriculum_learning' not in args:
            args['use_curriculum_learning'] = True
        if 'horizon' not in args:
            args['horizon'] = 12
            
        # 创建MegaCRN模型
        self.model = MegaCRN(
            num_nodes=args['num_nodes'],
            input_dim=1,  # 固定为1，因为我们只使用第一个通道
            output_dim=args['output_dim'],
            horizon=args['horizon'],
            rnn_units=args['rnn_units'],
            num_layers=args['num_layers'],
            cheb_k=args['cheb_k'],
            ycov_dim=args['ycov_dim'],
            mem_num=args['mem_num'],
            mem_dim=args['mem_dim'],
            cl_decay_steps=args['cl_decay_steps'],
            use_curriculum_learning=args['use_curriculum_learning']
        )
        
        self.args = args
        self.batches_seen = 0  # 添加batches_seen计数器
        
    def forward(self, x):
        # x shape: (batch_size, seq_len, num_nodes, features)
        # 按照DDGCRN的模式，只使用第一个通道
        x = x[..., 0].unsqueeze(-1)  # (batch_size, seq_len, num_nodes, 1)
        
        # 创建y_cov (这里使用零张量，实际使用时可能需要根据具体需求调整)
        y_cov = torch.zeros(x.shape[0], self.args['horizon'], x.shape[2], self.args['ycov_dim'], device=x.device)
        
        # 创建labels (这里使用零张量，实际使用时可能需要根据具体需求调整)
        labels = torch.zeros(x.shape[0], self.args['horizon'], x.shape[2], self.args['output_dim'], device=x.device)
        
        # 前向传播
        output, h_att, query, pos, neg = self.model(x, y_cov, labels=labels, batches_seen=self.batches_seen)
        
        # 更新batches_seen
        self.batches_seen += 1
        
        return output