main.py

# main.py
# This script serves as the entry point for training and testing the TransCoder model.
# It includes functions for setting up configurations, training the model, testing the model,
# and managing checkpoints. The script also handles noise generation, loss calculation,
# and performance evaluation for communication systems.
import os
import sys

# sys.path.append('./ldpc_nn')
import torch
import torch.nn as nn
import numpy as np
from torch import optim
import pandas as pd

from utils import *
from layers import *
from model import TransCoder
from parameters import *
from outer_code import *
from copy import deepcopy
import logging, time, random

from scipy.stats import gaussian_kde

def set_seed(seed=42):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

def ModelAvg(w):
    """
    Perform model averaging by averaging the weights of multiple models.

    Args:
        w (list): List of state dictionaries of models.

    Returns:
        dict: Averaged state dictionary.
    """
    w_avg = copy.deepcopy(w[0])
    for k in w_avg.keys():
        # Check if a weight requires gradient. If not, skip the averaging for that weight.
        if not w_avg[k].requires_grad:
            continue
        for i in range(1, len(w)):
            w_avg[k] += w[i][k]
        w_avg[k] = torch.div(w_avg[k], len(w))
    return w_avg

def power_norm(inputs):  # Normalize all through block
    """
    Normalize the power of the input tensor across the batch for every codeword to be 1.

    Args:
        inputs (torch.Tensor): Input tensor to normalize.

    Returns:
        torch.Tensor: Power-normalized tensor.
    """
    bs = inputs.size()[0]
    inputs_abs = inputs.clone().detach().abs().pow(2).reshape(bs, -1).mean(-1).sqrt().reshape(bs, 1)
    outputs = torch.div(inputs.reshape(bs, -1), inputs_abs)
    return outputs

def get_printing_metrics(ys0, belief_preds, priors, outer_code_type, model, iVec, batch_size):
    """
    Calculate and return performance metrics such as success rate, bit errors, and block errors.

    Args:
        ys0 (torch.Tensor): Ground truth labels.
        belief_preds (torch.Tensor): Predicted probabilities.
        priors (torch.Tensor): Priors from the outer code.
        outer_code_type (str): Type of outer code used.
        model (nn.Module): The TransCoder model.
        iVec (torch.Tensor): Original bit sequence.
        batch_size (int): Batch size.

    Returns:
        tuple: Success rate, bit errors, and block errors.
    """
    if outer_code_type != 'Turbo' or i == 0:
        preds = belief_preds
        logits, decodeds = preds.max(dim=-1)
        decodeds = decodeds.view(-1)
        succes_rate_iter_i = (sum(decodeds == ys0.to(model.device)) / len(ys0)).item()
    else:
        succes_rate_iter_i = 0
    message_llr = torch.split(priors, model.outer_code.k, dim=1)[0].contiguous()
    decoded_llr = message_llr.view(batch_size, -1)
    decoded = torch.signbit(decoded_llr).to(torch.float)
    if outer_code_type == 'Polar_usc':
        error_1 = (model.outer_code.X_train[0].view(batch_size, -1).to(model.device) - priors).abs().view(
            batch_size, -1)[:, model.outer_code.info_indices].sum(-1)
    else:
        error_1 = (iVec.to(model.device) - decoded).abs().view(batch_size, -1).sum(-1)
    bitErrors_i = error_1
    blockErrors_i = (error_1 != 0).to(float)
    return succes_rate_iter_i, bitErrors_i, blockErrors_i


def get_loss_values(bVec, padd_symbol_len, ones_symbols, num_blocks, M, map_vec):
    """
    Prepare loss-related values such as reshaped bit vectors and mapped labels.

    Args:
        bVec (torch.Tensor): Input bit vector.
        padd_symbol_len (int): Length of padding symbols.
        ones_symbols (torch.Tensor): Tensor of ones for padding.
        num_blocks (int): Number of blocks.
        M (int): Block length.
        map_vec (torch.Tensor): Mapping vector for look-up.

    Returns:
        tuple: Ground truth labels for loss calculation.
    """
    if padd_symbol_len > 0:
        bVec_new = bVec.clone()
        bVec_new = torch.cat((bVec_new, ones_symbols), dim=1)
        bVec_new = bVec_new.reshape(-1, num_blocks, M)
        bVec_mc = torch.matmul(bVec_new, map_vec)
        ys0 = bVec_mc.long().contiguous().view(-1)
        ys1 = ys0.clone()

    else:
        bVec_new = bVec.clone()
        bVec_new = bVec_new.reshape(-1, num_blocks, M)
        bVec_mc = torch.matmul(bVec_new, map_vec)
        ys0 = bVec_mc.long().contiguous().view(-1)
        ys1 = ys0.clone()
    return ys0, ys1

def getting_noise_samples(snr_type, snr_noisy, batch_size, K, device, snr_high=None, snr_low=None):
    """
    Generate noise samples based on the specified SNR type and range.

    Args:
        snr_type (str): Type of SNR ('all' or specific value).
        snr_noisy (bool): Whether to use noisy SNR.
        batch_size (int): Batch size.
        K (int): Length of the codeword.
        device (torch.device): Device to generate noise on.
        snr_high (float, optional): Upper bound of SNR range.
        snr_low (float, optional): Lower bound of SNR range.

    Returns:
        tuple: Noise samples, average SNR, and noise standard deviation.
    """
    noise_shape = (batch_size, K)
    if snr_type != 'all':
        snr = snr_type
        noise_shape = (batch_size, K)
        this_sigma = 10 ** (-snr * 1.0 / 10 / 2)
        this_sigma = torch.tensor(this_sigma).to(device).repeat(batch_size).view(-1, 1)
    elif snr_noisy:
        std_high = 10 ** (-snr_high * 1.0 / 10 / 2)
        std_low = 10 ** (-snr_low * 1.0 / 10 / 2)
        this_sigma = (std_low - std_high) * torch.rand(noise_shape).to(device) + std_high
        snr = torch.mean(this_sigma).item()
    else:
        std_high = 10 ** (-snr_high * 1.0 / 10 / 2)
        std_low = 10 ** (-snr_low * 1.0 / 10 / 2)
        this_sigma = (std_low - std_high) * torch.rand((batch_size, 1)).to(device) + std_high
        snr = torch.mean(this_sigma).item()

    noise_samples = this_sigma * torch.randn(noise_shape,
                                             dtype=torch.float, requires_grad=False).to(device)
    return noise_samples, snr, this_sigma

def directory_name(configs, num_blocks):
    """
    Generate a directory name based on the configuration parameters.

    Args:
        configs (dict): Configuration dictionary.
        num_blocks (int): Number of blocks.

    Returns:
        dict: Updated configuration dictionary with the directory name.
    """
    configs["saveDir"] = 'Results/CodeType_' + str(configs["outer_code_type"]) + \
                         "_K_in_" + str(configs["K_in"]) + \
                         "_K_" + str(configs["K"]) + \
                         "_K_out_" + str(num_blocks * configs["block_rate"]) + \
                         '_snr_' + str(configs["snr"]) + '_n_iter_' + str(configs["num_iters"]) + \
                         '_BlockRate_' + str(configs["block_rate"]) + '_bp_' + str(configs["bp_iters"]) + configs[
                             "loss_type"] + \
                         str(configs["usual_koef"]) + '_' + str(configs["first_koef"]) + '_lr'
    configs["saveDir"] += str(configs["lr"])
    if configs["opt_method"] != 'adamW':
        configs["saveDir"] += '_' + configs["opt_method"]
    if configs["use_lr_schedule"] != 'linear':
        configs["saveDir"] += '_' + configs["use_lr_schedule"]
    if configs["ldpc_type"] != 'first':
        configs["saveDir"] += configs["ldpc_type"]
    if configs["multilabel"] != 'all' and configs["multilabel"] is not None:
        configs["saveDir"] += '_' + configs["multilabel"]
    if not configs["backprop_outer"]:
        configs["saveDir"] += '_outf'
    if configs["no_inner"]:
        configs["saveDir"] += '_noin'
    if configs["af_module"]:
        configs["saveDir"] += '_af'
    if configs["weights"]:
        configs["saveDir"] += f'_w{configs["w_n"]}'
    if configs["loss_koef"] != 1:
        configs["saveDir"] += f'_lk{configs["loss_koef"]}'
    if configs["no_Tmodel"]:
        configs["saveDir"] += f'_noT'
    if configs["no_Rmodel"]:
        configs["saveDir"] += f'_noR'
    if configs["linear_layers"]:
        configs["saveDir"] += f'_ll'
    if configs["label_smoothing"]:
        configs["saveDir"] += f'_ls'
    # if configs["method_bp"] != 'MSA':
    #     configs["saveDir"] += f'_{configs["method_bp"]}'
    configs["saveDir"] += '/'
    return configs

def model_path_name(configs):
    """
    Generate a model path name based on the configuration parameters.

    Args:
        configs (dict): Configuration dictionary.

    Returns:
        str: Model path name.
    """
    model_path_add = ''
    if configs["reloc1"]:
        model_path_add += '_reloc1'
    if not configs["model_average"]:
        model_path_add += '_notaver'
    if configs["d_k_trx"] != 16:
        model_path_add += '_' + str(configs["d_k_trx"])
    if configs["N_rec"] != 3:
        model_path_add += '_nr' + str(configs["N_rec"])
    if configs["ldpc_type"] != 'first':
        model_path_add += configs["ldpc_type"]
    if configs["snr_noisy"]:
        model_path_add += '_snrall'
    if configs["ebno_train"]:
        model_path_add += '_ebn'
    if configs["high_bp"] != 6:
        model_path_add += f'_{configs["high_bp"]}'
    if configs["low_bp"] != 2:
        model_path_add += f'_{configs["low_bp"]}'
    if configs["rayleigh"]:
        model_path_add += '_rayleigh'
    if configs["snr_low"] == configs["snr_high"]:
        model_path_add += f'_s{configs["snr_low"]}'
    return model_path_add

def loss_calculation(ys0, ys1, belief_preds, model, bVec,
                     loss_type, num_iters=1, no_inner=False, table=None, loss_koef=1.0):
    """
    Calculate the loss for the model based on the specified loss type.

    Args:
        ys0 (torch.Tensor): Ground truth labels for the first iteration.
        ys1 (torch.Tensor): Ground truth labels for subsequent iterations.
        belief_preds (list): List of predicted probabilities for each iteration.
        model (nn.Module): The TransCoder model.
        bVec (torch.Tensor): Input bit vector.
        loss_type (str): Type of loss to calculate.
        num_iters (int): Number of decoding iterations. Default is 1.
        no_inner (bool): Whether to disable the TransCoder modules. Default is False.
        table (torch.Tensor, optional): Lookup table for TransCoder.
        loss_koef (float): Coefficient for scaling the loss. Default is 1.0.

    Returns:
        tuple: Total loss and list of individual losses.
    """
    device = belief_preds[0].device
    losses_arr = []

    loss_belief = 0
    loss = 0
    new_loss_type = ''
    if loss_type[:4] == 'llrs':
        for i in range(len(belief_preds)):
            bs = belief_preds[i].shape[0]
            bit_wise_one_probs, bit_wise_zero_probs = convert_class_to_bitwise_matmul(belief_preds[i], table)
            llrs = torch.log(bit_wise_one_probs) - torch.log(bit_wise_zero_probs)
            llrs = llrs.reshape(bs, model.K_extended)
            if model.padd_symbol_len > 0:
                llrs = llrs[:, :-model.padd_symbol_len]
                bit_one_probs = bit_wise_one_probs.reshape(bs,
                                                           model.K_extended)[:, :-model.padd_symbol_len]

            else:
                bit_one_probs = bit_wise_one_probs.reshape(bs,
                                                           model.K_extended)

            if i == 0 and no_inner:
                loss_belief_i = 0
            else:
                loss_belief_i = F.binary_cross_entropy(bit_one_probs, bVec.to(device))


            loss_belief += loss_belief_i
            losses_arr.append(loss_belief_i)
            new_loss_type = loss_type[4:]
    elif loss_type[:5] == 'usual':
        for i in range(len(belief_preds)):
            bs = belief_preds[i].shape[0]
            reshaped_pred = belief_preds[i].contiguous().view(-1, belief_preds[i].size(-1))
            log_pred = torch.log(reshaped_pred + 1e-10)
            if i == 0 and not no_inner:
                loss_belief_i = F.nll_loss(log_pred, ys0.to(device))#.item()
            elif i == 0 and no_inner:
                loss_belief_i = 0
            else:
                loss_belief_i = F.nll_loss(log_pred, ys1.to(device))#.item()
            bit_wise_one_probs, bit_wise_zero_probs = convert_class_to_bitwise_matmul(belief_preds[i], table)
            llrs = torch.log(bit_wise_one_probs) - torch.log(bit_wise_zero_probs)
            llrs = llrs.reshape(bs, model.K_extended)
            if model.padd_symbol_len > 0:
                llrs = llrs[:, :-model.padd_symbol_len]

            loss_belief += loss_belief_i
            losses_arr.append(loss_belief_i)
            new_loss_type = loss_type[5:]


    if loss_type == 'usual' or loss_type == 'llrs' or new_loss_type == '_outer' or new_loss_type == '_outerl':
        loss = loss_belief / len(belief_preds)

    if new_loss_type == 'f_outerl' or new_loss_type == 'f_outer':
        loss = losses_arr[0]

    if (loss_type == 'outer') or new_loss_type == 'f_outer' or new_loss_type == '_outer':
        loss += model.outer_code.loss[0] / (num_iters + 1)
        losses_arr.append(model.outer_code.loss[0].item())
        for i_beliefs in range(1, num_iters + 1):
            loss += model.outer_code.loss[i_beliefs] / (num_iters + 1)
            losses_arr.append(model.outer_code.loss[i_beliefs].item())

    if (loss_type == 'lastouter') or new_loss_type == 'f_outerl' or new_loss_type == '_outerl':
        losses_arr.append(model.outer_code.loss[0].item())
        for i_beliefs in range(1, num_iters + 1):
            losses_arr.append(model.outer_code.loss[i_beliefs].item())
        loss += model.outer_code.loss[num_iters]

    if new_loss_type != '':
        loss = loss * 0.5
    loss = loss * loss_koef

    return loss, losses_arr


def get_polar_code_channel_reliability(model, configs, tol=100, mc=10**6):
    """
    Evaluate the reliability of polar code channels using Monte Carlo simulations.

    Args:
        model (nn.Module): The TransCoder model.
        configs (dict): Configuration dictionary.
        tol (int): Tolerance for errors. Default is 100.
        mc (int): Number of Monte Carlo simulations. Default is 10^6.
    """
    device = model.device
    n = int(np.log2(model.K_code))
    dtype = torch.float32
    biterrd = torch.zeros([2 ** n], dtype=dtype).to(device)
    loss = torch.zeros([2 ** n], dtype=dtype).to(device)

    count = 0
    ebno = 4
    snr = ebno + 10 * np.log10(2 * model.outer_code.k / model.outer_code.n)

    k = model.outer_code.k
    model.outer_code.k = model.K_code
    model.K_in = model.K_code
    model.outer_code.info_indices = np.arange(model.K_code)

    model.outer_code.train_regime = 1
    while count < mc:
        with torch.no_grad():
            # loss_i, errors = self.forward_design(model.bs, 2 ** n)
            bVec, iVec = model.outer_code.generate_coded_bit_sequence(model.bs)
            bVec = bVec.to(device).float()
            bVec_md = 2 * bVec - 1
            noise_samples, snr, this_sigma = getting_noise_samples(snr,
                                                                   False,
                                                                   model.bs, model.K_transmit, device)

            _ = model(0,
                        bVec_md.to(device),
                        noise_samples,
                        table=model.table,
                        isTraining=1,
                        normalization_type=configs["constraint"],
                        current_sigma=this_sigma)

        loss += torch.sum(model.outer_code.loss_x_array[:, -1, :], dim=0)
        biterrd += torch.sum(model.outer_code.errors[:, -1, :], dim=0)
        count += model.bs


    print(f'iter: {count / mc * 100}% | bits w/o {tol} errors {biterrd}')

    biterrd /= count
    loss /= count
    print(f'conditional entropies of effective bit channels:\n'
          f'1 - loss: {torch.mean(np.log(2) - loss).detach().cpu().numpy() / np.log(2) : 5.4f}')
    mi = torch.squeeze(1 - loss)
    sorted_bit_channels = torch.argsort(-mi)

    model.outer_code.k = k
    model.K_in = k

    model.outer_code.sorted_bit_channels = sorted_bit_channels
    model.outer_code.info_indices = sorted(sorted_bit_channels[:model.outer_code.k])


def test_model(configs, model, path, logging_write=True, each_batch_all=100):
    """
    Test the model's performance and log the results.

    Args:
        configs (dict): Configuration dictionary.
        model (nn.Module): The TransCoder model.
        path (str): Path to save the results.
        logging_write (bool): Whether to write logs. Default is True.
        each_batch_all (int): Number of batches to test before early stopping. Default is 100.

    Returns:
        dict: Dictionary containing performance metrics.
    """
    model.eval()
    snr = configs["snr"]
    batch_size = configs["batch_size"]
    device = configs["device"]
    M = configs["M"]
    numTestbatch = 50000

    map_vec = 2 ** (torch.arange(M))
    map_vec = torch.flip(map_vec, [0]).to(device).float()
    A_blocks = generate_LUT(M)  # Look up table for blocks the order is according to the map vector

    bitErrors = [torch.zeros(batch_size).to(device) for _ in range(model.model_num_iters+1)]
    blockErrors = [torch.zeros(batch_size).to(device) for _ in range(model.model_num_iters+1)]

    df_names = ['iters', 'Inner_BER', 'BER', 'BLER',
                'Block_errors', 'Average_power']
    df = {name: [] for name in df_names}
    df['Block_errors'] = [0] *(model.model_num_iters+1)
    df['iters'] = [0] * (model.model_num_iters + 1)
    avg_power_cum = 0
    succes_rate_iter = [0] * (model.model_num_iters+1)
    inner_ber = [0] *(model.model_num_iters+1)
    BLER = [0] * (model.model_num_iters+1)
    BER = [0] * (model.model_num_iters+1)
    for each_batch in range(numTestbatch):
        # bVec -> outer coded sequence with shape BxK, iVec -> original bit sequence with shape BxK_in
        bVec, iVec = model.outer_code.generate_coded_bit_sequence(batch_size)
        bVec_md = 2 * bVec - 1

        # Generating the noise samples
        noise_samples, _, std = getting_noise_samples(snr,
                                                       configs["snr_noisy"],
                                                       batch_size, model.K_transmit, device,
                                                       snr_high=configs["snr_high"],
                                                       snr_low=configs["snr_low"])

        with (torch.no_grad()):
            belief_preds, preds_initial, generated_blocks, priors_1, priors_history = model(each_batch,
                                                                                          bVec_md.to(device),
                                                                                          noise_samples.to(device),
                                                                                          table=A_blocks.to(device),
                                                                                          isTraining=0,
                                                                                          normalization_type=configs["constraint"],
                                                                                          current_sigma=std)

            ys0, ys1 = get_loss_values(bVec, configs["padd_symbol_len"], model.ones_symbols,
                                       model.num_blocks, M, map_vec)
            if configs["loss_type"] == 'llrs':
                loss = 0
                losses_arr = [0]
            else:
                loss, losses_arr = loss_calculation(ys0, ys1, belief_preds, model, bVec_md,
                                                    configs["loss_type"],
                                                    num_iters=configs["num_iters_test"], no_inner=configs["no_inner"],
                                                    table=A_blocks.to(device), loss_koef=configs["loss_koef"])



            # Power calculation
            avg_power_cum += avg_power_calculation(generated_blocks.view(batch_size, -1))
            std = torch.mean(std).item()
            for i in range(len(priors_history)):
                succes_rate_iter_i, bitErrors_i, blockErrors_i \
                    = get_printing_metrics(ys0, belief_preds[i], priors_history[i],
                                           configs["outer_code_type"], model, iVec, batch_size)
                succes_rate_iter[i] += succes_rate_iter_i
                inner_ber[i] = 1 - (succes_rate_iter[i] / (each_batch + 1))
                bitErrors[i] += bitErrors_i
                BER[i] = (torch.sum(bitErrors[i] / model.outer_code.k) / (each_batch + 1) / batch_size).item()
                blockErrors[i] += blockErrors_i
                BLER[i] = (torch.sum(blockErrors[i]) / (each_batch + 1) / batch_size).item()


            if each_batch % 10 == 0:
                df['Inner_BER'] = inner_ber
                df['BER'] = BER
                df['BLER'] = BLER

                for i in range(model.model_num_iters+1):
                    df['Block_errors'][i] = torch.sum(blockErrors[i]).item()
                    df['iters'][i] = sum(np.array(model.bp_iters_arr)[:i + 1])
                ans_str_add = ''
                for i in range(model.model_num_iters+1):
                    for key in df.keys():
                        if key != 'Average_power':
                            ans_str_add += key + '_' + str(df[key][i]) + ', '

                df['Average_power'] = avg_power_cum.item() / (each_batch + 1)

                ans_str = f'{configs["outer_code_type"]} {configs["K"]} {configs["K_in"]} {configs["ldpc_type"]}: Iteration = {each_batch}, ' + \
                          f'SNR = {round(std, 3)}, BS = {batch_size}, Loss: {losses_arr}, ' + ans_str_add + \
                          f'AVG Power = {"{:.5e}".format(avg_power_cum.item() / (each_batch + 1))}'

                if logging_write or each_batch == numTestbatch - 1:
                    logging.warning(ans_str)

                if torch.sum(blockErrors[-1]).item() > 300 and each_batch >= 0 and torch.sum(blockErrors[0]).item() > 300:
                    logging.warning(ans_str)
                    break

    if configs["csv_name"] is not None:
        for i in range(model.model_num_iters+1):
            csv_name = configs["csv_name"] + '_it' + str(df['iters'][i]) + '_sheet.csv'
            df_save = dict()
            df_save['model_path'] = path
            df_save['Inner_BER'] = df['Inner_BER'][i]
            df_save['BER'] = df['BER'][i]
            df_save['BLER'] = df['BLER'][i]
            df_save['Block_errors'] = df['Block_errors'][i]
            df_save['Average_power'] = df['Average_power']
            df_save['num_iters'] = i

            try:
                dfnew = pd.read_csv(csv_name)
                last_ind = list(dfnew.index)[-1] + 1
            except:
                dfnew = pd.DataFrame.from_dict({})
                last_ind = 0

            df_add = pd.DataFrame(df_save, index=[last_ind])
            dffin = pd.concat([dfnew, df_add], ignore_index=True, sort=False)#, ignore_index=True, sort=False
            dffin.to_csv(csv_name)
    return df

def train_model(model, last_each_batch, configs):
    """
    Train the TransCoder model based on the specified configurations.

    Args:
        model (nn.Module): The TransCoder model.
        last_each_batch (int): Last completed batch index.
        configs (dict): Configuration dictionary.

    Returns:
        str: Path to the best model checkpoint.
    """
    print("-->-->-->-->-->-->-->-->-->--> start training ...")
    device = configs["device"]
    M = configs["M"]
    batch_size = configs["batch_size"]
    model.train()
    map_vec = 2 ** (torch.arange(M)).float()
    map_vec = torch.flip(map_vec, [0]).to(device)
    A_blocks = generate_LUT(M)  # Look up table for blocks the order is according to the map vector

    each_batch = last_each_batch + 1
    state_model_arr = []
    best_metric, best_loss = 1, 1
    model_path_add = model_path_name(configs)

    if configs["load_batches"] > 0:
        state_model = model.state_dict()
        state_model = copy.deepcopy(state_model)
        state_model_arr.append(state_model)
        state_model_arr.append(state_model)

    while each_batch < configs["total_batches"]:
        # Generating outer coded bit sequence for training
        bVec, iVec = model.outer_code.generate_coded_bit_sequence(batch_size)
        bVec = bVec.to(device).float()
        bs = bVec.shape[0]
        bVec_md = 2 * bVec - 1

        # Generate noise samples for training and testing
        noise_samples, snr, this_sigma = getting_noise_samples(configs["snr"],
                                                               configs["snr_noisy"],
                                                               bs, model.K_transmit, device,
                                                               snr_high=configs["snr_high"],
                                                               snr_low=configs["snr_low"])

        if np.mod(each_batch, configs["core"]) == 0 or each_batch == last_each_batch + 1:
            w_locals = []
            w0 = model.state_dict()
            w0 = copy.deepcopy(w0)

        if configs["random_bpiters"] or(configs["random_bpiters_iter"] is not None and each_batch < configs["random_bpiters_iter"]):
            model.bp_iters_arr = np.random.choice(np.arange(configs["low_bp"],configs["high_bp"]), model.model_num_iters+1)

        belief_preds_history, preds_initial, generated_blocks, priors_1, priors_history = model(each_batch, bVec_md.to(device),
                                                                                      noise_samples,
                                                                                      table=A_blocks.to(device),
                                                                                      isTraining=1,
                                                                                      normalization_type=configs["constraint"],
                                                                                      current_sigma=this_sigma)

        configs["optimizer"].zero_grad()

        ys0, ys1 = get_loss_values(bVec, model.padd_symbol_len, model.ones_symbols,
                                   model.num_blocks, M, map_vec)


        loss, losses_arr = loss_calculation(ys0, ys1, belief_preds_history, model,
                                            bVec, configs["loss_type"],
                                            no_inner=configs["no_inner"], num_iters=model.model_num_iters,
                                            table=A_blocks.to(device), loss_koef=configs["loss_koef"])
        loss.backward()

        # Catching NaN gradients
        nan_id = False
        for param in model.parameters():
            if param.grad is not None:
                if torch.any(torch.isnan(param.grad)):
                    nan_id = True
                    break
        if nan_id:
            model.load_state_dict(state_model_arr[-2])
            print('got nan at ', each_batch)
            continue
        else:
            state_model = model.state_dict()
            state_model = copy.deepcopy(state_model)
            state_model_arr.append(state_model)
            if len(state_model_arr) > 2:
                state_model_arr = state_model_arr[-2:]

        # Gradient clipping
        torch.nn.utils.clip_grad_norm(model.parameters(), configs["clip_th"])

        configs["optimizer"].step()

        # Saving the model
        w1 = model.state_dict()
        w_locals.append(copy.deepcopy(w1))
        if np.mod(each_batch, configs["core"]) != configs["core"] - 1:
            each_batch = each_batch + 1
            continue
        else:
            # Perform model averaging if enabled in the configurations
            if configs["model_average"]:
                w2 = ModelAvg(w_locals)  # Average the models
                model.load_state_dict(copy.deepcopy(w2))
            # Change the learning rate
            if configs["use_lr_schedule"]:
                configs["scheduler"].step()

        # AVG Power calculation
        avg_power = avg_power_calculation(generated_blocks.view(batch_size, -1))

        df_names = ['iters', 'Inner_BER', 'BER', 'BLER', 'Block_errors']
        df = {name: [] for name in df_names}
        with torch.no_grad():
            ans_str = ''
            for i in range(len(priors_history)):
                succes_rate, bitErrors, blockErrors \
                    = get_printing_metrics(ys0, belief_preds_history[i], priors_history[i],
                                           configs["outer_code_type"], model, iVec, batch_size)
                df['Inner_BER'].append(1 - succes_rate)
                df['iters'].append(sum(np.array(model.bp_iters_arr)[:i + 1]))
                df['BER'].append((bitErrors / model.outer_code.k).mean().item())
                block_errors = torch.sum(blockErrors)
                df['BLER'].append((block_errors / batch_size).item())
                df['Block_errors'].append(block_errors.item())

                for key in df.keys():
                    ans_str += key + ' ' + str(round(df[key][-1], 4)) + ', '


            logging.warning(f'Inner Code+{configs["outer_code_type"]} {configs["ldpc_type"]}: Iteration = {each_batch}, '
                          f'Learning Rate = {"{:.3e}".format(configs["scheduler"].get_last_lr()[0])}, '
                          f'SNR = {round(snr, 3)}, BS = {configs["batch_size"]}, Loss: {loss.item()}, '
                          f'{ans_str}' 
                          f'AVG Power: {"{:.3e}".format(avg_power.item())}')
            logging.warning(losses_arr)

            # Save tested model with best loss
            if loss.item() < best_loss:
                best_loss = loss.item()
                saveDir = configs["saveDir"] + 'best_loss' + model_path_add
                save_checkpoint(model=model, optimizer=configs["optimizer"], scheduler=configs["scheduler"],
                                eachbatch=each_batch, checkpoint_path=saveDir)
            # Save tested model with best metric (BLER)
            if configs["test_batch_num"] is not None and (each_batch + 1) % configs["test_batch_num"] == 0 and each_batch > 0.1 * configs["total_batches"]:
                configs["snr"] = (configs["snr_high"] + configs["snr_low"]) / 2 - 1.0
                configs["csv_name"] = None
                saveDir = configs["saveDir"] + 'model_' +str(each_batch + 1) + model_path_add

                configs["bp_iters_arr_test"] = [configs["bp_iters_test"]] * (configs["num_iters"]+1)
                model.bp_iters_arr = [configs["bp_iters_test"]] * (configs["num_iters"]+1)
                model.model_num_iters = configs["num_iters"]
                configs["num_iters_test"] = configs["num_iters"]
                df_test = test_model(configs, model, saveDir,
                                     logging_write=True, each_batch_all=100)
                if best_metric >= df_test['BLER'][-1]:
                    best_metric = df_test['BLER'][-1]
                    logging.warning(f'Best metric = {best_metric} iteration {each_batch + 1}')
                    save_checkpoint(model=model, optimizer=configs["optimizer"], scheduler=configs["scheduler"],
                                    eachbatch=each_batch, checkpoint_path=saveDir+'_'+str(best_metric))
                    saveDir = configs["saveDir"] + 'best_model' + model_path_add
                    save_checkpoint(model=model, optimizer=configs["optimizer"], scheduler=configs["scheduler"],
                                    eachbatch=each_batch, checkpoint_path=saveDir)
                if configs["snr_all"]:
                    configs["snr"] = 'all'
                    configs["snr_load"] = 'all'

        # Save the model checkpoint after every few iterations
        if np.mod(each_batch + 1, np.floor(configs["total_batches"] / 25)) == configs["core"] - 1:
            saveDir = configs["saveDir"] + 'model_' + str(each_batch + 1)
            saveDir += model_path_add
            save_checkpoint(model=model, optimizer=configs["optimizer"], scheduler=configs["scheduler"],
                            eachbatch=each_batch, checkpoint_path=saveDir)
        each_batch += 1
    return 'best_model' + model_path_add


if __name__ == '__main__':
    args = args_parser()
    configs = read_configs(args, args.outer_code_type)
    set_seed(configs['seed'])
    if configs["random_bpiters"]:
        configs["bp_iters"] = 0
    elif configs["bp_iters_arr"] is not None:
        configs["bp_iters"] = configs["bp_iters_arr"][0]

    configs["device"] = configs["device"] if torch.cuda.is_available() else 'cpu'
    print(f"The device used is :{configs['device']}")

    configs["d_model_trx"] = configs["heads_trx"] * configs["d_k_trx"]  # total number of features
    configs["model_path"] = ''

    # Initialize the TransCoder model with the given configurations
    model = TransCoder(configs).to(configs["device"])

    if configs["snr_all"]:
        snr_test = configs["snr"]
        configs["snr"] = 'all'
        configs["snr_load"] = 'all'

    # Convert SNR to Eb/N0 if needed
    if configs["ebno_train"]:
        code_rate = configs["K_in"] / configs["K"]
        if configs['protograph']:
            code_rate = 0.2
        configs["snr_low"] = configs["snr_low"] + 10 * np.log10(2 * code_rate)
        configs["snr_high"] = configs["snr_high"] + 10 * np.log10(2 * code_rate)
    print(configs["snr_low"], configs["snr_high"])
    configs = directory_name(configs, model.num_blocks)

    os.makedirs(configs["saveDir"], exist_ok=True)

    handlers = [
        logging.FileHandler(os.path.join(configs["saveDir"], 'logging.txt'))]
    handlers += [logging.StreamHandler()]
    logging.basicConfig(level=logging.NOTSET, format='%(message)s',
                        handlers=handlers)
    logging.warning(f"Path to model/logs: {configs['saveDir']}")
    logging.warning(configs)
    best_model_path=None
    model_path = configs["saveDir"] + 'model_' + str(configs["load_batches"])
    model_path += model_path_name(configs)
    if configs["train"] == 1:
        if configs["opt_method"] == 'adamW':
            configs["optimizer"] = torch.optim.AdamW(model.parameters(), lr=configs["lr"], betas=(0.9, 0.999), eps=1e-08,
                                               weight_decay=configs["wd"], amsgrad=False)
        else:
            configs["optimizer"] = torch.optim.Adam(model.parameters(), lr=configs["lr"], betas=(0.9, 0.98), eps=1e-9)
        if configs["use_lr_schedule"] == 'linear':
            lambda1 = lambda epoch: (1 - epoch / configs["total_batches"])

        configs["scheduler"] = torch.optim.lr_scheduler.LambdaLR(configs["optimizer"], lr_lambda=lambda1)
        if os.path.exists(model_path):
            model, configs["optimizer"], configs["scheduler"], each_batch = load_checkpoint(model=model,
                                                                                optimizer=configs["optimizer"],
                                                                                scheduler=configs["scheduler"],
                                                                                train=configs["train"],
                                                                                checkpoint_path=model_path,
                                                                                device=configs["device"])


        else:
            each_batch = 0
        logging.warning(model)
        logging.info(f'# of Parameters: {np.sum([np.prod(p.shape) for p in model.parameters()])}')
        model.model_path = model_path
        best_model_path = train_model(model, each_batch, configs)

    del model
    df_arr = []
    if configs["num_iters_test"] is None:
        configs["num_iters_test"] = configs["num_iters"]
    path_add = '_' + str(configs["num_iters_test"]) + '_' + str(configs["bp_iters_arr"])
    configs["batch_size"] = 10 * configs["batch_size"]

    configs["savef_arr"] = [25]
    if configs["K"] == 384 and (configs["K_in"] == 192 or configs["K_in"] == 256):
        configs["snr_arr"] = [4.5,5,5.5]#2,3,4,5,
    elif configs["K"] == 384 and configs["K_in"] == 320:
        configs["snr_arr"] = [6]#2,3,4,5,
    elif configs["K"] == 512:
        configs["snr_arr"] = [5]  # 2,3,4,5,
    elif configs["K"] == 1536:
        configs["snr_arr"] = [1,2,3]
    elif configs["K"] > 800 and configs["K_in"] < 400:
        configs["snr_arr"] = [2.5, 2.75]
    elif configs["K"] == 156:
        configs["snr_arr"] = [-1, 0, 1, 2, 3, 4, 5, 6]
    else:
        configs["snr_arr"] = [3]

    if len(configs["bp_iters_arr"]) == 0:
        configs["bp_iters_arr"] = [configs["bp_iters"]] * (configs["num_iters_test"] + 1)
    else:
        configs["bp_iters_arr"] = configs["bp_iters_arr_test"]
    model = TransCoder(configs).to(configs["device"])
    # if configs["modulation"] == 'BPSK':
    for i in range(len(configs["snr_arr"])):
        code_rate = configs["K_in"] / configs["K"]
        # if configs['protograph']:
        #     code_rate = 0.2
        if configs["K_in_polar"] is None:
            configs["snr_arr"][i] = configs["snr_arr"][i] + 10 * np.log10(2*code_rate)
        else:
            configs["snr_arr"][i] = configs["snr_arr"][i] + 10 * np.log10(2 * configs["K_in_polar"] / configs["K"])


    if configs["not_best_model"]:
        best_model_path = 'model_' + str(configs["load_batches"])
        best_model_path += model_path_name(configs)
    elif configs['test_bestloss']:
        best_model_path = 'best_loss'
        best_model_path += model_path_name(configs)
    elif best_model_path is None:
        best_model_path = 'best_model'
        best_model_path += model_path_name(configs)

    os.makedirs(configs["saveDir"] + best_model_path + path_add + '_log/', exist_ok=True)
    handlers = [
        logging.FileHandler(os.path.join(configs["saveDir"] + best_model_path + path_add + '_log/', 'logging.txt'))]
    handlers += [logging.StreamHandler()]
    logging.basicConfig(level=logging.NOTSET, format='%(message)s',
                        handlers=handlers)

    logging.warning(f"Path to model/logs: {model_path}")
    logging.warning(configs)
    if configs["outer_code_type"] == "Polar_bp"  and configs["num_iters"] != 0 and not configs["no_inner_test"]:
        test_settings = [[10,10],[4,4,4], [8], [10,10], [20], [5,5,5,5]]#[5, 5, 5,5,5], [2,2,2,2,2], [4,4,4,4,4], [3,3,3,3,3,3],] #[5, 5, 5,5],
    elif configs["outer_code_type"] == "BCH_nn"  and configs["num_iters"] != 0 and not configs["no_inner_test"]:
        test_settings = [[20]]#,[5,5,5,5,5,5,5,5],[40]]
    elif configs["K"] ==121:
        test_settings = [[20]]#[10,10,10,10],[20,20]]
    elif configs["outer_code_type"] == "Polar_usc":
        test_settings = [[1]]
    elif configs["outer_code_type"] == "Turbo" and configs["num_iters"] == 0:
        test_settings = [[1], [2], [3] ]  # [5, 5, 5,5],
    elif configs["outer_code_type"] == "Turbo" and configs["num_iters"] != 0:
        test_settings = [[1, 1, 1]]
    elif configs["num_iters"] == 0:
        test_settings = [[10],[5],  [15], [20]]
    elif configs["no_inner_test"]:
        test_settings = [[25], ]#[10], [20], [5], [15]]
    else:
        test_settings = [[10,10],[10,5,5],[15,5,5]]#[10,10,10,10], [20,20]] #[5,5,5,5],
    for item in test_settings:
        for snr in configs["snr_arr"]:
            for savef in configs["savef_arr"]:
                configs["save_first"] = savef
                configs["snr"] = snr
                if configs["outer_code_type"] == "Turbo" or configs["outer_code_type"] == "Polar_usc":
                    configs["bp_iters"] = item[0]
                    model = TransCoder(configs).to(configs["device"])

                model = load_checkpoint(model=model, optimizer=None, scheduler=None,
                                        checkpoint_path=configs["saveDir"] + best_model_path,
                                        train=0, device=configs["device"])

                configs["bp_iters_arr_test"] = item
                model.bp_iters_arr = item
                model.model_num_iters = len(item) - 1
                model.model_path = configs["saveDir"] + best_model_path + path_add
                configs["num_iters_test"] = model.model_num_iters
                name_part = 'CodeType_' + str(configs["outer_code_type"]) + \
                   "_K_in_" + str(configs["K_in"]) + "_K_" + str(configs["K"])
                os.makedirs('Results/' + name_part + '/', exist_ok=True)
                configs["csv_name"] = 'Results/' + name_part + '/' + name_part + '_snr_' + str(round(snr, 3))

                if configs["K_in_polar"] is not None:
                    model.K_in = configs["K_in_polar"]
                    model.outer_code.k = configs["K_in_polar"]


                if (configs["outer_code_type"] == "Polar_usc"): # and not configs["no_inner_test"]:

                    if configs['check_vanilla']:
                        model.outer_code.code_sc.checknode = model.outer_code.code_sc.checknode_check
                        model.outer_code.code_sc.checknode_enc = model.outer_code.code_sc.checknode_check

                    if not configs["no_inner_test"] :
                        get_polar_code_channel_reliability(model, configs, tol=100, mc=10 ** 6)


                if configs["no_inner_test"]:
                    model.no_inner = True


                df = test_model(configs, model, configs["saveDir"] + best_model_path, configs["logging_write"])
                df_arr.append(df)