import torch
import torch.nn as nn
import torch.nn.functional as F
from data.get_adj import get_gso

class STSSLModel(nn.Module):
    def __init__(self, args):
        super(STSSLModel, self).__init__()
        
        # 设置默认参数
        if 'd_model' not in args:
            args['d_model'] = 64
        if 'd_output' not in args:
            args['d_output'] = args['output_dim']
        if 'input_length' not in args:
            args['input_length'] = args['n_his']
        if 'dropout' not in args:
            args['dropout'] = 0.1
        if 'nmb_prototype' not in args:
            args['nmb_prototype'] = 10
        if 'batch_size' not in args:
            args['batch_size'] = 64
        if 'shm_temp' not in args:
            args['shm_temp'] = 0.1
        if 'yita' not in args:
            args['yita'] = 0.5
        if 'percent' not in args:
            args['percent'] = 0.1
        if 'device' not in args:
            args['device'] = 'cpu'
        if 'gso_type' not in args:
            args['gso_type'] = 'sym_norm_lap'
        if 'graph_conv_type' not in args:
            args['graph_conv_type'] = 'cheb_graph_conv'
            
        # 保存参数
        self.args = args
        self.num_nodes = args['num_nodes']
        self.input_dim = args['input_dim']
        self.output_dim = args['output_dim']
        self.horizon = args['horizon']
        self.d_model = args['d_model']
        
        # 获取邻接矩阵
        self.gso = get_gso(args)
        
        # 时间嵌入
        self.T_i_D_emb = nn.Parameter(torch.empty(288, args['d_model']))
        self.D_i_W_emb = nn.Parameter(torch.empty(7, args['d_model']))
        
        # 节点嵌入
        self.node_emb_u = nn.Parameter(torch.randn(self.num_nodes, args['d_model']))
        self.node_emb_d = nn.Parameter(torch.randn(self.num_nodes, args['d_model']))
        
        # 编码器 - 使用1个输入通道
        self.encoder = STEncoder(
            Kt=3, Ks=3, 
            input_dim=1,  # 只使用第一个通道
            hidden_dim=args['d_model'], 
            input_length=args['input_length'], 
            num_nodes=args['num_nodes'], 
            droprate=args['dropout']
        )
        
        # 预测头
        self.predictor = nn.Linear(args['d_model'], args['output_dim'])
        
        # 初始化参数
        self.reset_parameters()
        
    def reset_parameters(self):
        nn.init.xavier_uniform_(self.node_emb_u)
        nn.init.xavier_uniform_(self.node_emb_d)
        nn.init.xavier_uniform_(self.T_i_D_emb)
        nn.init.xavier_uniform_(self.D_i_W_emb)
        
    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)
        
        # 编码
        encoded = self.encoder(x, self.gso)
        
        # 预测
        # 取最后一个时间步的输出进行预测
        last_hidden = encoded[:, -1, :, :]  # (batch_size, num_nodes, d_model)
        
        # 预测未来horizon个时间步
        predictions = []
        for t in range(self.horizon):
            pred = self.predictor(last_hidden)  # (batch_size, num_nodes, output_dim)
            predictions.append(pred)
        
        # 堆叠预测结果
        output = torch.stack(predictions, dim=1)  # (batch_size, horizon, num_nodes, output_dim)
        
        return output


class STEncoder(nn.Module):
    def __init__(self, Kt, Ks, input_dim, hidden_dim, input_length, num_nodes, droprate):
        super(STEncoder, self).__init__()
        self.num_nodes = num_nodes
        self.input_length = input_length
        
        # 简化的时空编码器 - 使用1个输入通道
        self.conv1 = nn.Conv2d(input_dim, hidden_dim//2, kernel_size=(Kt, Ks), padding=(Kt//2, Ks//2))
        self.conv2 = nn.Conv2d(hidden_dim//2, hidden_dim, kernel_size=(Kt, Ks), padding=(Kt//2, Ks//2))
        self.dropout = nn.Dropout(droprate)
        
    def forward(self, x, graph):
        # x: (batch_size, seq_len, num_nodes, features)
        batch_size, seq_len, num_nodes, features = x.shape
        
        # 调整维度
        x = x.permute(0, 3, 1, 2)  # (batch_size, features, seq_len, num_nodes)
        
        # 卷积操作
        x = F.relu(self.conv1(x))
        x = self.dropout(x)
        x = F.relu(self.conv2(x))
        
        # 调整回原维度
        x = x.permute(0, 2, 3, 1)  # (batch_size, seq_len, num_nodes, features)
        
        return x