from typing import Any

# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
know_dag = torch.from_numpy(np.loadtxt('../Frcsub_DATA/Qwen_dag.txt')).float()

class GAE(nn.Module):
    def __init__(self,hidden_dim, know_dim):
        super(GAE,self).__init__()
        self.hidden_dim = hidden_dim
        self.know_dim = know_dim
        
        self.encoder = nn.ModuleList([
            nn.Sequential(
            nn.Linear(1, self.hidden_dim),
            nn.BatchNorm1d(self.hidden_dim),
            nn.Sigmoid(),
            nn.Linear(self.hidden_dim, self.know_dim),
            nn.BatchNorm1d(self.know_dim),
            nn.Sigmoid()
        ) for _ in range(self.know_dim)
        ]
        )

        self.decoder = nn.ModuleList([
            nn.Sequential(
            nn.Linear(1, self.hidden_dim),
            nn.BatchNorm1d(self.hidden_dim),
            nn.Sigmoid(),
            nn.Linear(self.hidden_dim, 1),
            nn.Sigmoid() 
            ) for _ in range(self.know_dim)
        ]
        )
    
    def forward(self,stu_emb, know_dag):  # stu_emb (32, know_dim) 
        encoded_feature = []
        for i in range(self.know_dim):
            x_i = stu_emb[:,i].unsqueeze(1)  # x_i (32, 1)
            h_i = self.encoder[i](x_i)  # h_i  (32, know_dim)
            encoded_feature.append(h_i.unsqueeze(1))  # (32, 1 ,know_dim)

        H = torch.cat(encoded_feature,dim = 1)  # H (32, know_dim, know_dim)
        H_ = H * know_dag.unsqueeze(0) 
        H_ = H_.sum(dim = 1)  # H_ (32, 8)
        
        decoded_output = []
        for i in range(self.know_dim):
            H_i = H_[:,i].unsqueeze(dim = 1)  # (32, 1)
            x_hat_i = self.decoder[i](H_i)
            decoded_output.append(x_hat_i)

        X_hat = torch.cat(decoded_output, dim = 1)
        return X_hat

class NeuralCD(nn.Module):
    def __init__(self, student_num, exer_num, know_num, know_dag) -> None:
        super(NeuralCD,self).__init__()
        self.student_num = student_num
        self.exer_num = exer_num
        self.know_num = know_num

        self.student_emb = nn.Embedding(self.student_num, self.know_num)
        self.exer_diff = nn.Embedding(self.exer_num, self.know_num)
        self.exer_disc = nn.Embedding(self.exer_num, 1)
        
        self.gae = GAE(32,8)
        self.linear1 = nn.Sequential(
            nn.Linear(self.know_num, 512),
            nn.BatchNorm1d(512),
        )
        self.dropout1 = nn.Dropout(p = 0.5)
        self.linear2 = nn.Sequential(
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
        )
        self.dropout2 = nn.Dropout(p = 0.5)
        self.linear3 = nn.Linear(256, 1)
        # self.know_dag = nn.Parameter(torch.empty(know_num, know_num))
        ''' 初始化参数 '''
        self._initialize_()
        self.know_dag = nn.Parameter(know_dag.float(), requires_grad=True)
        
    def _initialize_(self):
        for name, param in self.named_parameters():
            if param.dim() == 1:
                continue
            if 'weight' in name or name == "know_dag":
                nn.init.xavier_normal_(param)
    
    def forward(self, stu_id, exer_id, K_R):
        stu_pre = self.student_emb(stu_id)
        stu_emb = torch.sigmoid(stu_pre)
        exer_diff = torch.sigmoid(self.exer_diff(exer_id))
        exer_disc = torch.sigmoid(self.exer_disc(exer_id))
        
        stu_emb_re = self.gae(stu_emb,self.know_dag)

        x = K_R * (stu_emb - exer_diff) * exer_disc
        x = self.dropout1(torch.sigmoid(self.linear1(x)))
        x = self.dropout2(torch.sigmoid(self.linear2(x)))
        x = torch.sigmoid(self.linear3(x))
        
        return x, stu_emb, stu_emb_re
    
    def apply_clipper(self):
        clipper = NoneNegClipper()
        self.linear1.apply(clipper)
        self.linear2.apply(clipper)
        self.linear3.apply(clipper)
        # self.linear4.apply(clipper)

    def get_stu_knowledge_status(self,stu_id):
        stu_emb = torch.sigmoid(self.student_emb(stu_id))
        stu_emb_re = self.gae(stu_emb,self.know_dag)
        return stu_emb.data, stu_emb_re.data

代码加了BatchNorm1d后就报错
ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 32])