import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from polar import PolarCode, get_frozen
from models import g_Full, f_Full, weights_init
from utils import min_sum_log_sum_exp, min_sum_log_sum_exp_4, countSetBits, log_sum_exp, STEQuantize
from collections import defaultdict
import os



###��������ģ�ͼܹ��Ż���
#������Ƚ���������ܹ�����Transformer��ResNetģ�飩
#����ע�������������õز�׽����֮��Ĺ�ϵ
#ʵ�ָ�������ȿ����üܹ�
class TransformerLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = nn.ReLU()

    def forward(self, src):
        # src��״: (sequence_length, batch_size, d_model)
        src2 = self.self_attn(src, src, src)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src
####ADD down
    
class DeepPolar(PolarCode):
    def __init__(self, args, device, N, K, ell = 2, infty = 1000., depth_map : defaultdict = None):
        # ��ʼ��DeepPolar��
        # args: ����ģ�Ͳ��������ö���
        # device: �����豸����CPU��GPU��
        # N: Polar��ĳ���
        # K: ��Ϣλ������
        # ell: ÿ����ӿ��С
        # infty: �����ֵ�����ڱ�ʾ���ɿ����ŵ�
        # depth_map: ÿ����ӿ��С���ֵ�
        args.use_transformer = True  # �Ƿ�ʹ��Transformer
        args.nhead = 4              # Transformerͷ��
        args.dim_feedforward = 2048 # FFNά��
        args.transformer_dropout = 0.1 # Dropout��
    
    # ����Transformer���ʼ������ʼ��Transformer����޸�
        self.transformer_layers = nn.ModuleDict()
        if hasattr(args, 'use_transformer') and args.use_transformer:
            for d in range(1, self.n_ell+1):
                self.transformer_layers[str(d)] = TransformerLayer(
                    d_model=self.depth_map[d],
                    nhead=getattr(args, 'nhead', 4),  # Ĭ��4��ͷ
                    dim_feedforward=getattr(args, 'dim_feedforward', 2048),
                    dropout=getattr(args, 'transformer_dropout', 0.1)
                ).to(device)
        # �������������ʼ���޸ģ���Ҫ��DeepPolar��__init__���������Ӷ�������ĳ�ʼ����
        if args.use_transformer:
            self.transformer_layers = nn.ModuleDict()
            for d in range(1, self.n_ell+1):
                self.transformer_layers[str(d)] = TransformerLayer(
                    d_model=self.depth_map[d],
                    nhead=args.nhead,
                    dim_feedforward=args.dim_feedforward,
                    dropout=args.dropout_p
                )
        
        if args.use_attention:
            self.attention_layers = nn.ModuleDict()
            for d in range(1, self.n_ell+1):
                self.attention_layers[str(d)] = nn.MultiheadAttention(
                    embed_dim=self.depth_map[d],
                    num_heads=args.nhead,
                    dropout=args.dropout_p
                )
        # ��ȡ����λ����Ϣλ��
        Fr = get_frozen(N, K, self.args.rate_profile)
        # ���ø���PolarCode�ĳ�ʼ������
        super().__init__(n = int(np.log2(N)), K = K, Fr=Fr, infty = infty)
        self.N = N  # Polar��ĳ���

        # ����ṩ��depth_map����ʹ������ȷ��ÿһ����ӿ��С
        if depth_map is not None:
            # ���depth_map��ֵ�ĳ˻��Ƿ����N
            assert np.prod(list(depth_map.values())) == N
            # ���depth_map�ļ��Ƿ��1��ʼ������ĳ����
            assert min(list(depth_map.keys())) == 1
            assert max(list(depth_map.keys())) <= int(np.log2(N))
            self.ell = None  # ��ʹ�ù̶���ellֵ
            self.n_ell = len(depth_map.keys())  # ��Ȳ���
            assert max(list(depth_map.keys())) == self.n_ell  # ��������Ƿ������Ȳ���
            self.depth_map = depth_map  # �洢depth_map
        else:
            # ���û���ṩdepth_map����ʹ�ù̶���ellֵ
            self.ell = ell
            self.n_ell = int(np.log(N)/np.log(self.ell))  # ������Ȳ���
            self.depth_map = defaultdict(int)  # ����һ��Ĭ���ֵ�
            for d in range(1, self.n_ell+1):
                self.depth_map[d] = self.ell  # Ϊÿһ�������ӿ��С
            assert np.prod(list(self.depth_map.values())) == N  # ����ӿ��С�ĳ˻��Ƿ����N

        self.device = device  # �洢�����豸
        self.fnet_dict = None  # ���ڴ洢��������
        self.gnet_dict = None  # ���ڴ洢��������
        self.infty = infty  # �洢�����ֵ

    @staticmethod
    def get_onehot(actions):
        # ������ת��Ϊone-hot����
        # actions: ��������
        # ����: one-hot���������

        inds = (0.5 + 0.5*actions).long()  # ������ת��Ϊ����
        if len(actions.shape) == 2:
            # ������������Ƕ�ά�ģ�����one-hot����
            return torch.eye(2, device = inds.device)[inds].reshape(actions.shape[0], -1)
        elif len(actions.shape) == 3:
            # ���������������ά�ģ�����one-hot����
            return torch.eye(2, device = inds.device)[inds].reshape(actions.shape[0], actions.shape[1], -1)

    def define_kernel_nns(self, ell, unfrozen = None, fnet = 'KO', gnet = 'KO', shared = False):
        # �������ͱ�������
        # ell: ÿ����ӿ��С
        # unfrozen: δ�����λλ��
        # fnet: ������������
        # gnet: ������������
        # shared: �Ƿ����������

        if 'KO' in fnet:
            self.fnet_dict = {}  # ��ʼ�����������ֵ�
        else:
            self.fnet_dict = None  # �����ʹ�ý������磬����ΪNone
        self.shared = shared  # �洢�Ƿ�������

        if 'KO' in gnet:
            self.gnet_dict = {}  # ��ʼ�����������ֵ�
        else:
            self.gnet_dict = None  # �����ʹ�ñ������磬����ΪNone

        dec_hidden_size = self.args.dec_hidden_size  # ������������ز��С
        enc_hidden_size = self.args.enc_hidden_size  # ������������ز��С
        depth = 1  # ��ǰ���
        assert len(unfrozen) > 0, "No unfrozen bits!"  # ����Ƿ���δ�����λ

        self.fnet_dict[depth] = {}  # ��ʼ����ǰ��ȵĽ�������
        if fnet == 'KO_parallel' or fnet == 'KO_last_parallel':
            bit_position = 0  # ��ǰλλ��
            self.fnet_dict[depth][bit_position] = {}  # ��ʼ����ǰλλ�õĽ�������
            # �����С
            input_size = ell
            # �����С
            output_size = ell
            # ������������
            self.fnet_dict[depth][bit_position] = f_Full(input_size, dec_hidden_size, output_size, activation = self.args.dec_activation, dropout_p = self.args.dropout_p, depth = self.args.f_depth, use_norm = self.args.use_norm).to(self.device)
        elif 'KO' in fnet:
            if shared:
                self.fnet_dict[depth] = {}  # ��ʼ����ǰ��ȵĽ�������
                for current_position in range(ell):
                    # �������������Ľ�������
                    self.fnet_dict[depth][current_position] = f_Full(ell + current_position, dec_hidden_size, 1, activation = self.args.dec_activation, dropout_p = self.args.dropout_p, depth = self.args.f_depth, use_norm = self.args.use_norm).to(self.device)
            else:
                bit_position = 0  # ��ǰλλ��
                for current_position in unfrozen:
                    if not self.fnet_dict[depth].get(bit_position):
                        self.fnet_dict[depth][bit_position] = {}  # ��ʼ����ǰλλ�õĽ�������
                    input_size = ell + (int(self.args.onehot)+1)*current_position  # �����С
                    # ������������
                    self.fnet_dict[depth][bit_position][current_position] = f_Full(input_size, dec_hidden_size, 1, activation = self.args.dec_activation, dropout_p = self.args.dropout_p, depth = self.args.f_depth, use_norm = self.args.use_norm).to(self.device)

        if 'KO' in gnet:
            self.gnet_dict[depth] = {}  # ��ʼ����ǰ��ȵı�������
            if shared:
                if gnet == 'KO':
                    # �������������ı�������
                    self.gnet_dict[depth] = g_Full(ell, enc_hidden_size, ell-1, depth = self.args.g_depth, skip_depth = self.args.g_skip_depth, skip_layer = self.args.g_skip_layer, ell = ell, activation = self.args.enc_activation, use_skip = self.args.skip).to(self.device)
            else:
                bit_position = 0  # ��ǰλλ��
                if gnet == 'KO':
                    # ������������
                    self.gnet_dict[depth][bit_position] = g_Full(ell, enc_hidden_size, ell-1, depth = self.args.g_depth, skip_depth = self.args.g_skip_depth, skip_layer = self.args.g_skip_layer, ell = ell, activation = self.args.enc_activation, use_skip = self.args.skip).to(self.device)

    def define_and_load_nns(self, ell, kernel_load_path = None, fnet = 'KO', gnet = 'KO', shared = True, dataparallel = False):
        # ���岢���ؽ���ͱ�������
        # ell: ÿ����ӿ��С
        # kernel_load_path: �������������·��
        # fnet: ������������
        # gnet: ������������
        # shared: �Ƿ����������
        # dataparallel: �Ƿ�ʹ�����ݲ���

        if 'KO' in fnet:
            self.fnet_dict = {}  # ��ʼ�����������ֵ�
        else:
            self.fnet_dict = None  # �����ʹ�ý������磬����ΪNone
        self.shared = shared  # �洢�Ƿ�������

        if 'KO' in gnet:
            self.gnet_dict = {}  # ��ʼ�����������ֵ�
        else:
            self.gnet_dict = None  # �����ʹ�ñ������磬����ΪNone

        dec_hidden_size = self.args.dec_hidden_size  # ������������ز��С
        enc_hidden_size = self.args.enc_hidden_size  # ������������ز��С

        for depth in range(self.n_ell, 0, -1):
            if depth in self.args.polar_depths:
                continue  # �����ǰ�����ָ��������б��У�����
            ell = self.depth_map[depth]  # ��ȡ��ǰ��ȵ��ӿ��С
            proj_size = np.prod([self.depth_map[d] for d in range(1, depth+1)])  # ���㵱ǰ��ȵ�ͶӰ��С

            if fnet == 'KO_last_parallel' and depth == 1:
                self.fnet_dict[depth] = {}  # ��ʼ����ǰ��ȵĽ�������
                for bit_position in range(self.N // proj_size):
                    proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)  # ���㵱ǰλλ�õ�ͶӰ
                    get_num_info_proj = lambda proj : sum([int(x in self.info_positions) for x in proj])  # ����ͶӰ�е���Ϣλ����
                    num_info_in_proj = get_num_info_proj(proj)  # ��ȡͶӰ�е���Ϣλ����
                    subproj_len = len(proj) // ell  # ������ͶӰ�ĳ���
                    subproj = [proj[i:i+subproj_len] for i in range(0, len(proj), subproj_len)]  # �ָ�ͶӰΪ��ͶӰ
                    num_info_in_subproj = [get_num_info_proj(x) for x in subproj]  # ����ÿ����ͶӰ�е���Ϣλ����
                    unfrozen = [i for i, x in enumerate(num_info_in_subproj) if x >= 1]  # ��ȡδ�������ͶӰλ��

                    input_size = ell  # �����С
                    output_size = ell  # �����С
                    # ������������
                    self.fnet_dict[depth][bit_position] = f_Full(input_size, dec_hidden_size, output_size, activation = self.args.dec_activation, dropout_p = self.args.dropout_p, depth = self.args.f_depth, use_norm = self.args.use_norm).to(self.device)

                    if len(unfrozen) > 0:
                        if kernel_load_path is not None:
                            try:
                                # ���Լ���Ԥѵ���Ľ����������
                                ckpt = torch.load(os.path.join(kernel_load_path + '_parallel', f'{ell}_{len(unfrozen)}.pt'))
                                ckpt_exists = True
                            except FileNotFoundError:
                                print(f"Parallel File not found for ell = {ell}, num_unfrozen = {len(unfrozen)}")
                                ckpt_exists = False
                        else:
                            ckpt_exists = False
                        if ckpt_exists:
                            # ���ؽ����������
                            f_ckpt = ckpt[0][1][0].state_dict()
                            self.fnet_dict[depth][bit_position].load_state_dict(f_ckpt)
                            if dataparallel:
                                # ʹ�����ݲ���
                                self.fnet_dict[depth][bit_position] = nn.DataParallel(self.fnet_dict[depth][bit_position])
            elif 'KO' in fnet:
                self.fnet_dict[depth] = {}  # ��ʼ����ǰ��ȵĽ�������
                if shared:
                    self.fnet_dict[depth] = {}  # ��ʼ����ǰ��ȵĽ�������
                    for current_position in range(ell):
                        # �������������Ľ�������
                        self.fnet_dict[depth][current_position] = f_Full(ell + current_position, dec_hidden_size, 1, activation = self.args.dec_activation, dropout_p = self.args.dropout_p, depth = self.args.f_depth, use_norm = self.args.use_norm).to(self.device)
                        if dataparallel:
                            # ʹ�����ݲ���
                            self.fnet_dict[depth][current_position] = nn.DataParallel(self.fnet_dict[depth][current_position])
                else:
                    for bit_position in range(self.N // proj_size):
                        proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)  # ���㵱ǰλλ�õ�ͶӰ
                        get_num_info_proj = lambda proj : sum([int(x in self.info_positions) for x in proj])  # ����ͶӰ�е���Ϣλ����
                        num_info_in_proj = get_num_info_proj(proj)  # ��ȡͶӰ�е���Ϣλ����
                        subproj_len = len(proj) // ell  # ������ͶӰ�ĳ���
                        subproj = [proj[i:i+subproj_len] for i in range(0, len(proj), subproj_len)]  # �ָ�ͶӰΪ��ͶӰ
                        num_info_in_subproj = [get_num_info_proj(x) for x in subproj]  # ����ÿ����ͶӰ�е���Ϣλ����
                        unfrozen = [i for i, x in enumerate(num_info_in_subproj) if x >= 1]  # ��ȡδ�������ͶӰλ��

                        if len(unfrozen) > 0:
                            if kernel_load_path is not None:
                                try:
                                    # ���Լ���Ԥѵ���Ľ����������
                                    ckpt = torch.load(os.path.join(kernel_load_path, f'{ell}_{len(unfrozen)}.pt'))
                                    ckpt_exists = True
                                except FileNotFoundError:
                                    print(f"File not found for ell = {ell}, num_unfrozen = {len(unfrozen)}")
                                    ckpt_exists = False
                            else:
                                ckpt_exists = False
                            for current_position in unfrozen:
                                if not self.fnet_dict[depth].get(bit_position):
                                    self.fnet_dict[depth][bit_position] = {}  # ��ʼ����ǰλλ�õĽ�������
                                input_size = ell + (int(self.args.onehot)+1)*current_position  # �����С
                                output_size = 1  # �����С
                                # ������������
                                self.fnet_dict[depth][bit_position][current_position] = f_Full(input_size, dec_hidden_size, output_size, activation = self.args.dec_activation, dropout_p = self.args.dropout_p, depth = self.args.f_depth, use_norm = self.args.use_norm).to(self.device)
                                if ckpt_exists:
                                    try:
                                        # ���ؽ����������
                                        f_ckpt = ckpt[0][1][0][current_position].state_dict()
                                    except:
                                        print(unfrozen)
                                    self.fnet_dict[depth][bit_position][current_position].load_state_dict(f_ckpt)
                                    if dataparallel:
                                        # ʹ�����ݲ���
                                        self.fnet_dict[depth][bit_position][current_position] = nn.DataParallel(self.fnet_dict[depth][bit_position][current_position])

            if 'KO' in gnet:
                self.gnet_dict[depth] = {}  # ��ʼ����ǰ��ȵı�������
                if shared:
                    if gnet == 'KO':
                        if not dataparallel:
                            # �������������ı�������
                            self.gnet_dict[depth] = g_Full(ell, enc_hidden_size, ell-1, depth = self.args.g_depth, skip_depth = self.args.g_skip_depth, skip_layer = self.args.g_skip_layer, ell = ell, use_skip = self.args.skip).to(self.device)
                        else:
                            # ʹ�����ݲ��д������������ı�������
                            self.gnet_dict[depth] = nn.DataParallel(g_Full(ell, enc_hidden_size, ell-1, depth = self.args.g_depth, skip_depth = self.args.g_skip_depth, skip_layer = self.args.g_skip_layer, ell = ell, use_skip = self.args.skip)).to(self.device)
                else:
                    for bit_position in range(self.N // proj_size):
                        proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)  # ���㵱ǰλλ�õ�ͶӰ
                        get_num_info_proj = lambda proj : sum([int(x in self.info_positions) for x in proj])  # ����ͶӰ�е���Ϣλ����
                        num_info_in_proj = get_num_info_proj(proj)  # ��ȡͶӰ�е���Ϣλ����
                        subproj_len = len(proj) // ell  # ������ͶӰ�ĳ���
                        subproj = [proj[i:i+subproj_len] for i in range(0, len(proj), subproj_len)]  # �ָ�ͶӰΪ��ͶӰ
                        num_info_in_subproj = [get_num_info_proj(x) for x in subproj]  # ����ÿ����ͶӰ�е���Ϣλ����
                        unfrozen = [i for i, x in enumerate(num_info_in_subproj) if x >= 1]

                        if num_info_in_proj > 0:
                            if gnet == 'KO':
                                self.gnet_dict[depth][bit_position] = g_Full(ell, enc_hidden_size, ell-1, depth = self.args.g_depth, skip_depth = self.args.g_skip_depth, skip_layer = self.args.g_skip_layer, ell = ell, activation = self.args.enc_activation, use_skip = self.args.skip).to(self.device)
                            if kernel_load_path is not None:
                                try:
                                    ckpt = torch.load(os.path.join(kernel_load_path, f'{ell}_{len(unfrozen)}.pt'))
                                    self.gnet_dict[depth][bit_position].load_state_dict(ckpt[1][1][0].state_dict())
                                except FileNotFoundError:
                                    print(f"File not found for ell = {ell}, num_unfrozen = {len(unfrozen)}")
                                    pass
                            if dataparallel:
                                self.gnet_dict[depth][bit_position] = nn.DataParallel(self.gnet_dict[depth][bit_position])


                        # print(f"g : {depth}, {bit_position}, {len(unfrozen)}")
        if kernel_load_path is not None:
            print("Loaded kernel from ", kernel_load_path)

    def load_nns(self, fnet_dict, gnet_dict = None, shared = False):
        self.fnet_dict = fnet_dict
        self.gnet_dict = gnet_dict

        for depth in fnet_dict.keys():
            if self.fnet_dict is not None:
                for bit_position in self.fnet_dict[depth].keys():
                    if not isinstance(self.fnet_dict[depth][bit_position], dict):#shared or self.args.decoder_type == 'KO_parallel' or self.args.decoder_type == 'KO_RNN':
                        self.fnet_dict[depth][bit_position].to(self.device)
                    else:
                        for current_position in self.fnet_dict[depth][bit_position].keys():
                            self.fnet_dict[depth][bit_position][current_position].to(self.device)
            if gnet_dict is not None:
                if shared:
                    self.gnet_dict[depth].to(self.device)
                else:
                    for bit_position in self.gnet_dict[depth].keys():
                        self.gnet_dict[depth][bit_position].to(self.device)
        print("NN weights loaded!")

    def load_partial_nns(self, fnet_dict, gnet_dict = None):

        for depth in fnet_dict.keys():
            if fnet_dict is not None:
                for bit_position in fnet_dict[depth].keys():
                    if isinstance(fnet_dict[depth][bit_position], dict):
                        for current_position in fnet_dict[depth][bit_position].keys():
                            self.fnet_dict[depth][bit_position][current_position] = fnet_dict[depth][bit_position][current_position].to(self.device)
                    else:
                        self.fnet_dict[depth][bit_position] = fnet_dict[depth][bit_position].to(self.device)

            if gnet_dict is not None:
                for bit_position in gnet_dict[depth].keys():
                    self.gnet_dict[depth][bit_position] = gnet_dict[depth][bit_position].to(self.device)
        print("NN weights loaded!")

    def kernel_encode(self, ell, gnet, msg_bits, info_positions, binary = False):
        input_shape = msg_bits.shape[-1]
        assert input_shape <= ell
        u = torch.ones(msg_bits.shape[0], self.N, dtype=torch.float).to(self.device)
        u[:, info_positions] = msg_bits
        output =torch.cat([gnet(u.unsqueeze(1)).squeeze(1), u[:, -1:]], 1)

        power_constrained_u = self.power_constraint(output)
        if binary:
            stequantize = STEQuantize.apply
            power_constrained_u = stequantize(power_constrained_u)
        return power_constrained_u

####�������Ľ��ı������ܹ����滻���е�kernel_encode������
# ���������ǿ���Ľ��������Ĺ���Լ������
#��ǿ�ı������������������ӵ�DeepPolar���У���Ϊ���б���������������򲹳䡣
    def enhanced_encode(self, msg_bits, binary=False):
        u = torch.ones(msg_bits.shape[0], self.N, dtype=torch.float).to(self.device)
        u[:, self.info_positions] = msg_bits
        
    # ...ǰ��Ĵ���...
        
        # Ӧ��Transformer��
        chunks = [u[:, i:i+proj_size//ell].unsqueeze(1) for i in 
                range(bit_position*proj_size, (bit_position+1)*proj_size, proj_size//ell)]
        concatenated = torch.cat(chunks, 1)  # shape: (batch_size, num_chunks, chunk_size)
        
        # ת��ΪTransformer��Ҫ����״������
        transformer_input = concatenated.permute(1, 0, 2)  # (num_chunks, batch_size, chunk_size)
        transformer_out = self.transformer_layers[d](transformer_input)
        output = transformer_out.permute(1, 0, 2)  # ת����(batch_size, num_chunks, chunk_size)
        
    # ...����Ĵ���...

        # ����λ�ñ���
        position = torch.arange(self.N).unsqueeze(0).to(self.device)
        pos_enc = torch.zeros_like(u)
        pos_enc[:, :] = position.float()
        u = u + pos_enc
        
        # �ֲ㴦��
        for d in range(1, self.n_ell+1):
            proj_size = np.prod([self.depth_map[dd] for dd in range(1, d+1)])
            ell = self.depth_map[d]
            
            for bit_position in range(self.N // proj_size):
                proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)
                num_info = sum([int(x in self.info_positions) for x in proj])
                
                if num_info > 0:
                    # ʹ��Transformer������Ϣλ
                    chunks = [u[:, i:i+proj_size//ell].unsqueeze(1) for i in 
                            range(bit_position*proj_size, (bit_position+1)*proj_size, proj_size//ell)]
                    concatenated = torch.cat(chunks, 1)
                    
                    # Ӧ��Transformer��
                    transformer_out = self.transformer_layers[d](concatenated.permute(1, 0, 2))
                    output = transformer_out.permute(1, 0, 2).reshape(u.shape[0], -1)
                    
                    u = torch.cat((u[:, :bit_position*proj_size], output, 
                                u[:, (bit_position+1)*proj_size:]), dim=1)
        
        # �Ľ��Ĺ���Լ��
        power_constrained_u = self.enhanced_power_constraint(u)
        
        if binary:
            stequantize = STEQuantize.apply
            power_constrained_u = stequantize(power_constrained_u)
        
        return power_constrained_u

    def enhanced_power_constraint(self, codewords):
        # ��������Ӧ���ʵ���
        norm = torch.norm(codewords, p=2, dim=1, keepdim=True)
        scale = torch.sqrt(self.N) / (norm + 1e-6)
        return codewords * scale
####ADD Down


#�Ľ��Ľ������ܹ����滻���е�deeppolar_decode������
#��ǿ�������Ĳ��д����������Ż���Ϣλ�����߼�
###��ǿ�Ľ�����������ͬ�����ӵ�DeepPolar���У����������н������������棬ͨ����������ʹ�����ֽ��뷽ʽ��
    def enhanced_decode(self, noisy_code):
        assert noisy_code.shape[1] == self.N
        
        # ��ʼ��
        decoded_llrs = self.infty*torch.ones(noisy_code.shape[0], self.N, device=noisy_code.device)
        partial_sums = torch.ones(noisy_code.shape[0], self.n_ell+1, self.N, device=noisy_code.device)
        
        # ��������Ԥ����
        noisy_code = self.noise_preprocess(noisy_code)
        
        # �ݹ����
        decoded_llrs, partial_sums = self.enhanced_decode_depth(
            noisy_code.unsqueeze(2), self.n_ell, 0, decoded_llrs, partial_sums)
        
        return decoded_llrs[:, self.info_positions], torch.sign(decoded_llrs[:, self.info_positions])

    def noise_preprocess(self, noisy_code):
        # Ӧ��С�����������Ԥ��������
        if self.args.use_wavelet:
            noisy_code = self.wavelet_denoise(noisy_code)
        return noisy_code

    def enhanced_decode_depth(self, llrs, depth, bit_position, decoded_llrs, partial_sums):
        half_index = np.prod([self.depth_map[d] for d in range(1, depth)]) if depth > 1 else 1
        ell = self.depth_map[depth]
        left_bit_position = self.depth_map[depth] * bit_position
        
        # ��Ϣλ���
        proj_size = np.prod([self.depth_map[d] for d in range(1, depth+1)])
        proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)
        num_info = sum([int(x in self.info_positions) for x in proj])
        
        # ʹ��ע����������ǿ����
        dec_chunks = [llrs[:, (j)*half_index:(j+1)*half_index].clone() for j in range(ell)]
        
        if num_info > 0 and depth in self.fnet_dict:
            # Ӧ��ע��������
            concatenated = torch.cat(dec_chunks, 2)
            if self.args.use_attention:
                concatenated = self.attention_layers[depth](concatenated)
            
            # ���н���
            if self.args.parallel_decode:
                Lu = self.fnet_dict[depth][bit_position](concatenated)
                decoded_llrs[:, left_bit_position + np.array(self.info_positions)] = Lu.squeeze(1)
            else:
                # ���н���
                for current_position in range(ell):
                    if current_position > 0:
                        prev_decoded = partial_sums[:, depth-1, (current_position-1)*half_index:current_position*half_index].unsqueeze(2)
                        dec_chunks.append(prev_decoded)
                    
                    if current_position in [p-left_bit_position for p in self.info_positions if left_bit_position <= p < left_bit_position+proj_size]:
                        concatenated = torch.cat(dec_chunks, 2)
                        Lu = self.fnet_dict[depth][bit_position][current_position](concatenated)
                        decoded_llrs[:, left_bit_position + current_position] = Lu.squeeze(2).squeeze(1)
        
        # �ݹ鴦��
        if depth > 1:
            for current_position in range(ell):
                bit_position_offset = left_bit_position + current_position
                if bit_position_offset in self.info_positions:
                    decoded_llrs, partial_sums = self.enhanced_decode_depth(
                        llrs[:, current_position*half_index:(current_position+1)*half_index], 
                        depth-1, bit_position_offset, decoded_llrs, partial_sums)
        
        return decoded_llrs, partial_sums
######ADD Down

    def deeppolar_encode(self, msg_bits, binary = False):
        u = torch.ones(msg_bits.shape[0], self.N, dtype=torch.float).to(self.device)
        u[:, self.info_positions] = msg_bits
        for d in range(1, self.n_ell+1):
            # num_bits = self.ell**(d-1)
            num_bits = np.prod([self.depth_map[dd] for dd in range(1, d)]) if d > 1 else 1
            # proj_size = self.ell**(d)
            proj_size = np.prod([self.depth_map[dd] for dd in range(1, d+1)])
            ell = self.depth_map[d]
            for bit_position, i in enumerate(np.arange(0, self.N, ell*num_bits)):

                # [u v] encoded to [(u xor v),v)]
                proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)
                get_num_info_proj = lambda proj : sum([int(x in self.info_positions) for x in proj])
                num_info_in_proj = get_num_info_proj(proj)

                subproj_len = len(proj) // ell
                subproj = [proj[i:i+subproj_len] for i in range(0, len(proj), subproj_len)]
                num_info_in_subproj = [get_num_info_proj(x) for x in subproj]
                num_nonzero_subproj = sum([int(x != 0) for x in num_info_in_subproj])
                
                if num_info_in_proj > 0:
                    info_bits_present = True          
                else:
                    info_bits_present = False         
                if d in self.args.polar_depths:
                    info_bits_present = False

                enc_chunks = []
                ell = self.depth_map[d]
                for j in range(ell):
                    chunk = u[:, i + j*num_bits:i + (j+1)*num_bits].unsqueeze(2).clone()
                    enc_chunks.append(chunk)
                if info_bits_present:
                    concatenated_chunks = torch.cat(enc_chunks, 2)
                    if self.shared:
                        output = torch.cat([self.gnet_dict[d](concatenated_chunks), u[:, i + (ell-1)*num_bits:i + (ell)*num_bits].unsqueeze(2)], dim=2)
                    else:
                        output = torch.cat([self.gnet_dict[d][bit_position](concatenated_chunks), u[:, i + (ell-1)*num_bits:i + (ell)*num_bits].unsqueeze(2)], dim=2)
                    output = output.permute(0,2,1).reshape(msg_bits.shape[0], -1, 1).squeeze(2)

                else:
                    output = self.encode_chunks_plotkin(enc_chunks, ell)
                u = torch.cat((u[:, :i], output, u[:, i + ell*num_bits:]), dim=1)

        power_constrained_u = self.power_constraint(u)
        if binary:
            stequantize = STEQuantize.apply
            power_constrained_u = stequantize(power_constrained_u)
        return power_constrained_u

    def power_constraint(self, codewords):
        return F.normalize(codewords, p=2, dim=1)*np.sqrt(self.N)

    def encode_chunks_plotkin(self, enc_chunks, ell = None):

        # message shape is (batch, k)
        # BPSK convention : 0 -> +1, 1 -> -1
        # Therefore, xor(a, b) = a*b

        # to change for other kernels

        if ell is None:
            ell = self.ell
        assert len(enc_chunks) == ell
        chunk_size = enc_chunks[0].shape[1]
        batch_size = enc_chunks[0].shape[0]

        u = torch.cat(enc_chunks, 1).squeeze(2)
        n = int(np.log2(ell))

        for d in range(0, n):
            num_bits = 2**d * chunk_size
            for i in np.arange(0, chunk_size*ell, 2*num_bits):
                # [u v] encoded to [(u,v) xor v]
                u = torch.cat((u[:, :i], u[:, i:i+num_bits].clone() * u[:, i+num_bits: i+2*num_bits], u[:, i+num_bits:]), dim=1)
        return u
            
    def deeppolar_parallel_decode(self, noisy_code):
        # Successive cancellation decoder for polar codes
        assert noisy_code.shape[1] == self.N

        depth = self.n_ell

        decoded_llrs = self.infty*torch.ones(noisy_code.shape[0], self.N, device = noisy_code.device)
        # function is recursively called (DFS)
        # arguments: Beliefs at the input of node (LLRs at top node), depth of children, bit_position (zero at top node)
        decoded_llrs  = self.KO_parallel_decode_depth(noisy_code.unsqueeze(2), depth, 0, decoded_llrs)
        decoded_llrs = decoded_llrs[:, self.info_positions]
        return decoded_llrs, torch.sign(decoded_llrs)

    def deeppolar_parallel_decode_depth(self, llrs, depth, bit_position, decoded_llrs):
        # Function to call recursively, for SC decoder

        # half_index = self.ell ** (depth - 1)
        half_index = np.prod([self.depth_map[d] for d in range(1, depth)]) if depth > 1 else 1
        ell = self.depth_map[depth]
        left_bit_position = self.depth_map[depth] *  bit_position 

        # Check if >1 information bits are present in the current projection. If not, don't use NNs - use polar encoding and minsum SC decoding.
        # proj_size = self.ell**(depth)
        proj_size = np.prod([self.depth_map[d] for d in range(1, depth+1)])

        proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)
        get_num_info_proj = lambda proj : sum([int(x in self.info_positions) for x in proj])
        get_info_proj = lambda proj : [x for x in proj if x in self.info_positions]

        num_info_in_proj = get_num_info_proj(proj)
        info_in_proj = get_info_proj(proj)

        subproj_len = len(proj) // ell
        subproj = [proj[i:i+subproj_len] for i in range(0, len(proj), subproj_len)]
        num_info_in_subproj = [get_num_info_proj(x) for x in subproj]
        num_nonzero_subproj = sum([int(x != 0) for x in num_info_in_subproj])
        unfrozen = np.array([i for i, x in enumerate(num_info_in_subproj) if x >= 1])

        dec_chunks = torch.cat([llrs[:, (j)*half_index:(j+1)*half_index].clone() for j in range(ell)], 2)
        Lu = self.fnet_dict[depth][bit_position](dec_chunks)

        if depth == 1:
            u = torch.tanh(Lu/2)
            decoded_llrs[:, left_bit_position + unfrozen] = Lu.squeeze(1)
        else:
            for index, current_position in enumerate(unfrozen):
                bit_position_offset = left_bit_position + current_position                
                decoded_llrs = self.deeppolar_parallel_decode_depth(Lu[:, :, index:index+1], depth-1, bit_position_offset, decoded_llrs)

        return decoded_llrs
            
    def deeppolar_decode(self, noisy_code):
        assert noisy_code.shape[1] == self.N

        depth = self.n_ell

        decoded_llrs = self.infty*torch.ones(noisy_code.shape[0], self.N, device = noisy_code.device)
        
        # don't want to go into useless frozen subtrees.
        partial_sums = torch.ones(noisy_code.shape[0], self.n_ell+1, self.N, device=noisy_code.device)

        # function is recursively called (DFS)
        # arguments: Beliefs at the input of node (LLRs at top node), depth of children, bit_position (zero at top node)

        decoded_llrs, partial_sums = self.deeppolar_decode_depth(noisy_code.unsqueeze(2), depth, 0, decoded_llrs, partial_sums)
        decoded_llrs = decoded_llrs[:, self.info_positions]

        return decoded_llrs, torch.sign(decoded_llrs)
    
    def deeppolar_decode_depth(self, llrs, depth, bit_position, decoded_llrs, partial_sums):
        # Function to call recursively, for SC decoder

        # half_index = self.ell ** (depth - 1)
        half_index = np.prod([self.depth_map[d] for d in range(1, depth)]) if depth > 1 else 1
        ell = self.depth_map[depth]
        left_bit_position = self.depth_map[depth] *  bit_position 

        # Check if >1 information bits are present in the current projection. If not, don't use NNs - use polar encoding and minsum SC decoding.
        # proj_size = self.ell**(depth)
        # size of the projection of tht subtree
        proj_size = np.prod([self.depth_map[d] for d in range(1, depth+1)])

        # This chunk - finds infrozen positions in this kernel.
        proj = np.arange(bit_position*proj_size, (bit_position+1)*proj_size)
        get_num_info_proj = lambda proj : sum([int(x in self.info_positions) for x in proj])
        get_info_proj = lambda proj : [x for x in proj if x in self.info_positions]

        num_info_in_proj = get_num_info_proj(proj)
        info_in_proj = get_info_proj(proj)

        subproj_len = len(proj) // ell
        subproj = [proj[i:i+subproj_len] for i in range(0, len(proj), subproj_len)]
        num_info_in_subproj = [get_num_info_proj(x) for x in subproj]
        num_nonzero_subproj = sum([int(x != 0) for x in num_info_in_subproj])
        unfrozen = np.array([i for i, x in enumerate(num_info_in_subproj) if x >= 1])

        if num_nonzero_subproj > 0:
            info_bits_present = True      
        else:
            info_bits_present = False 

        if depth in self.args.polar_depths:
            info_bits_present = False
                
        # This will be input to decoder
        dec_chunks = [llrs[:, (j)*half_index:(j+1)*half_index].clone() for j in range(ell)]
        # n = 2 tree case
        if depth == 1:
            if self.args.decoder_type == 'KO_last_parallel':
                concatenated_chunks = torch.cat(dec_chunks, 2)
                Lu = self.fnet_dict[depth][bit_position](concatenated_chunks)[:, 0, unfrozen]
                u_hat = torch.tanh(Lu/2)
                decoded_llrs[:, left_bit_position + unfrozen] = Lu
                partial_sums[:, depth-1, left_bit_position + unfrozen] = u_hat

            else:
                for current_position in range(ell):
                    bit_position_offset = left_bit_position + current_position
                    if current_position > 0:
                        # I am adding previously decoded bits . (either onehot or normal)
                        if self.args.onehot:
                            prev_decoded = get_onehot(partial_sums[:, depth-1, (current_position -1)*half_index:(current_position)*half_index].unsqueeze(2).sign()).detach().clone()
                        else:
                            prev_decoded = partial_sums[:, depth-1, (current_position -1)*half_index:(current_position)*half_index].unsqueeze(2).clone()
                        dec_chunks.append(prev_decoded)

                    if bit_position_offset in self.frozen_positions: # frozen 
                        # don't update decoded llrs. It already has ones*prior.
                        # actually don't need this. can skip.
                        partial_sums[:, depth-1, bit_position_offset] = torch.ones_like(partial_sums[:, depth-1, bit_position_offset])
                    else: # information bit
                        # This is the decoding.
                        concatenated_chunks = torch.cat(dec_chunks, 2)
                        if self.shared:
                            Lu = self.fnet_dict[depth][current_position](concatenated_chunks)
                        else:
                            Lu = self.fnet_dict[depth][bit_position][current_position](concatenated_chunks)

                        u_hat = torch.tanh(Lu/2).squeeze(2)
                        decoded_llrs[:, bit_position_offset] = Lu.squeeze(2).squeeze(1)
                        partial_sums[:, depth-1, bit_position_offset] = u_hat.squeeze(1)

            # Encoding back the decoded bits - for higher layers.
            # # Compute decoded codeword
            i = left_bit_position * half_index
            # num_bits = self.ell**(depth-1)
            num_bits = 1

            enc_chunks = []
            for j in range(ell):
                chunk = torch.sign(partial_sums[:, depth-1, i + j*num_bits:i + (j+1)*num_bits]).unsqueeze(2).detach().clone()
                enc_chunks.append(chunk)
            if info_bits_present:
                concatenated_chunks = torch.cat(enc_chunks, 2)
                if 'KO' in self.args.encoder_type:
                    if self.shared:
                        output = torch.cat([self.gnet_dict[depth](concatenated_chunks), partial_sums[:, depth-1, i + (ell-1)*num_bits:i + (ell)*num_bits].unsqueeze(2)], dim=2)
                    else:
                        # bit position of the previous depth.
                        output = torch.cat([self.gnet_dict[depth][bit_position](concatenated_chunks), partial_sums[:, depth-1, i + (ell-1)*num_bits:i + (ell)*num_bits].unsqueeze(2)], dim=2)
                    output = output.permute(0,2,1).reshape(llrs.shape[0], -1, 1).squeeze(2)
                else:
                    output = self.encode_chunks_plotkin(enc_chunks, ell)
            else:
                output = self.encode_chunks_plotkin(enc_chunks, ell)
            partial_sums[:, depth, i : i + num_bits*ell] = output.clone()
            
            return decoded_llrs, partial_sums

        # General case
        else:
            for current_position in range(ell):
                bit_position_offset = left_bit_position + current_position

                if current_position > 0:
                    if self.args.onehot:
                        prev_decoded = get_onehot(partial_sums[:, depth-1, (current_position -1)*half_index:(current_position)*half_index].unsqueeze(2).sign()).detach().clone()
                    else:
                        prev_decoded = partial_sums[:, depth-1, (current_position -1)*half_index:(current_position)*half_index].unsqueeze(2).clone()
                    dec_chunks.append(prev_decoded)
                concatenated_chunks = torch.cat(dec_chunks, 2)

                if current_position in unfrozen:
                    # General decoding ....
                    # add the decoded bit here
                    if self.shared:
                        Lu = self.fnet_dict[depth][current_position](concatenated_chunks).squeeze(2)
                    else:
                        # if current_position == 0:
                        #     Lu = self.fnet_dict[depth][bit_position][current_position](llrs)
                        # else:
                        Lu = self.fnet_dict[depth][bit_position][current_position](concatenated_chunks)
                    decoded_llrs, partial_sums = self.deeppolar_decode_depth(Lu, depth-1, bit_position_offset, decoded_llrs, partial_sums)
                else:
                    Lu = self.infty*torch.ones_like(llrs)


            # Compute decoded codeword
            if depth < self.n_ell :
                i = left_bit_position * half_index
                # num_bits = self.ell**(depth-1)
                num_bits = np.prod([self.depth_map[d] for d in range(1, depth)])
                enc_chunks = []
                for j in range(ell):
                    chunk = torch.sign(partial_sums[:, depth-1, i + j*num_bits:i + (j+1)*num_bits]).unsqueeze(2).detach().clone()
                    enc_chunks.append(chunk)
                if info_bits_present:
                    concatenated_chunks = torch.cat(enc_chunks, 2)
                    if 'KO' in self.args.encoder_type:
                        if self.shared:
                            output = torch.cat([self.gnet_dict[depth](concatenated_chunks), partial_sums[:, depth-1, i + (ell-1)*num_bits:i + (ell)*num_bits].unsqueeze(2)], dim=2)
                        else:
                            # bit position of the previous depth.
                            output = torch.cat([self.gnet_dict[depth][bit_position](concatenated_chunks), partial_sums[:, depth-1, i + (ell-1)*num_bits:i + (ell)*num_bits].unsqueeze(2)], dim=2)
                        output = output.permute(0,2,1).reshape(llrs.shape[0], -1, 1).squeeze(2)
                    else:
                        output = self.encode_chunks_plotkin(enc_chunks, ell)
                else:
                    output = self.encode_chunks_plotkin(enc_chunks, ell)
                partial_sums[:, depth, i : i + num_bits*ell] = output.clone()

                return decoded_llrs, partial_sums
            else: # encoding not required for last level - we have already decoded all bits.
                return decoded_llrs, partial_sums


    def kernel_decode(self, ell, fnet_dict, noisy_code, info_positions = None):
        input_shape = noisy_code.shape[-1]
        noisy_code = noisy_code.unsqueeze(2)
        assert input_shape == ell
        u = torch.ones(noisy_code.shape[0], self.N, dtype=torch.float).to(self.device)
        decoded_llrs = self.infty*torch.ones(noisy_code.shape[0], self.N, device = noisy_code.device)
        half_index = 1
        dec_chunks = [noisy_code[:, (j)*half_index:(j+1)*half_index].clone() for j in range(ell)]

        for current_position in range(ell):
            if current_position > 0:
                if self.args.onehot:
                    prev_decoded = get_onehot(u[:, (current_position -1)*half_index:(current_position)*half_index].unsqueeze(2).clone().sign()).detach().clone()
                else:
                    prev_decoded = u[:, (current_position -1)*half_index:(current_position)*half_index].unsqueeze(2).clone()
                dec_chunks.append(prev_decoded)
            if current_position in info_positions:
                if current_position in info_positions:
                    concatenated_chunks = torch.cat(dec_chunks, 2)
                    Lu = fnet_dict[current_position](concatenated_chunks)
                    decoded_llrs[:, current_position] = Lu.squeeze(2).squeeze(1)
                    u_hat = torch.tanh(Lu/2).squeeze(2)
                    u[:, current_position] = u_hat.squeeze(1)
        return decoded_llrs[:, info_positions], u[:, info_positions]

    def kernel_parallel_decode(self, ell, fnet_dict, noisy_code, info_positions = None):
        input_shape = noisy_code.shape[-1]
        noisy_code = noisy_code.unsqueeze(2)
        assert input_shape == ell
        u = torch.ones(noisy_code.shape[0], self.N, dtype=torch.float).to(self.device)
        decoded_llrs = self.infty*torch.ones(noisy_code.shape[0], self.N, device = noisy_code.device)
        half_index = 1
        dec_chunks = torch.cat([noisy_code[:, (j)*half_index:(j+1)*half_index].clone() for j in range(ell)], 2)

        decoded_llrs = fnet_dict(dec_chunks).squeeze(1)
        u = torch.tanh(decoded_llrs/2).squeeze(1)
        return decoded_llrs[:, info_positions], u[:, info_positions]
    

#     ѵ�������Ż�����
# �γ�ѧϰ�������ӱ��븴�ӶȺ�����ˮƽ
# ��Ͼ���ѵ����ʹ��torch.cuda.amp����ѵ��
# �Ľ�����ʧ��������Ͻ����غ��Զ�����ʧ
    def enhanced_loss_function(pred_llrs, true_bits, snr):
        # ������������ʧ
        ce_loss = F.binary_cross_entropy_with_logits(pred_llrs, (true_bits+1)/2)
        
        # �Զ���ɿ�����ʧ
        reliability = torch.sigmoid(torch.abs(pred_llrs))
        rel_loss = -torch.mean(reliability)
        
        # SNR����Ӧ��Ȩ
        snr_weight = 1.0 / (1.0 + torch.exp(-0.1*(snr-10)))
        total_loss = ce_loss + snr_weight * rel_loss
        
        return total_loss
# ѵ�������Ż���
# ʵ�ֿγ�ѧϰ����
# ���Ӹ����ӵ����򻯼���
# �Ľ���ʧ����
##ADD down