main.py

from methods.ecmp import RUN_ECMP
from methods.lp import RUN_LP
from methods.fc import RUN_FC
from methods.lstm import RUN_LSTM
from methods.cnn import RUN_CNN
from methods.drl import RUN_DRL
import networkx as nx
import argparse


def main(method, train_flag, metric_flag, topo, scale=1150383, k=4, obj='MLU', train_size=4500, batch_size=64, alpha=0, learning_rate=0.01, seed=76, epoch=100, weight_decay = 0.001, buf_capacity=10000, learning_rate_actor=0.0001, learning_rate_critic=0.0001, gamma=0.99, tau=0.01, sigma=0.6):
    net = nx.read_graphml(f'dataset/topologies/{topo}.graphml')
    path_path = f'dataset/paths/{topo}.npy'
    traffic_path = f'dataset/traffic/{topo}.npy'
    ca_path = f'dataset/capacities/{topo}.npy'
    save_dir = f'results/{method}/{topo}/'

    # algorithm model
    if train_flag == True:
        if method == 'ECMP':
            running = RUN_ECMP(net=net, path_path=path_path, traffic_path=traffic_path, ca_path=ca_path, scale=scale, k=k, save_dir=save_dir, obj=obj, train_size=train_size, batch_size=batch_size, alpha=alpha, learning_rate=learning_rate, seed=seed, epoch=epoch, weight_decay=weight_decay, buf_capacity=buf_capacity, learning_rate_actor=learning_rate_actor, learning_rate_critic=learning_rate_critic, gamma=gamma, tau=tau, sigma=sigma)
        elif method == 'LP':
            running = RUN_LP(net=net, path_path=path_path, traffic_path=traffic_path, ca_path=ca_path, scale=scale, k=k, save_dir=save_dir, obj=obj, train_size=train_size, batch_size=batch_size, alpha=alpha, learning_rate=learning_rate, seed=seed, epoch=epoch, weight_decay=weight_decay, buf_capacity=buf_capacity, learning_rate_actor=learning_rate_actor, learning_rate_critic=learning_rate_critic, gamma=gamma, tau=tau, sigma=sigma)
        elif method == 'FC':
            running = RUN_FC(net=net, path_path=path_path, traffic_path=traffic_path, ca_path=ca_path, scale=scale, k=k, save_dir=save_dir, obj=obj, train_size=train_size, batch_size=batch_size, alpha=alpha, learning_rate=learning_rate, seed=seed, epoch=epoch, weight_decay=weight_decay, buf_capacity=buf_capacity, learning_rate_actor=learning_rate_actor, learning_rate_critic=learning_rate_critic, gamma=gamma, tau=tau, sigma=sigma)
        elif method == 'LSTM':
            running = RUN_LSTM(net=net, path_path=path_path, traffic_path=traffic_path, ca_path=ca_path, scale=scale, k=k, save_dir=save_dir, obj=obj, train_size=train_size, batch_size=batch_size, alpha=alpha, learning_rate=learning_rate, seed=seed, epoch=epoch, weight_decay=weight_decay, buf_capacity=buf_capacity, learning_rate_actor=learning_rate_actor, learning_rate_critic=learning_rate_critic, gamma=gamma, tau=tau, sigma=sigma)
        elif method == 'CNN':
            running = RUN_CNN(net=net, path_path=path_path, traffic_path=traffic_path, ca_path=ca_path, scale=scale, k=k, save_dir=save_dir, obj=obj, train_size=train_size, batch_size=batch_size, alpha=alpha, learning_rate=learning_rate, seed=seed, epoch=epoch, weight_decay=weight_decay, buf_capacity=buf_capacity, learning_rate_actor=learning_rate_actor, learning_rate_critic=learning_rate_critic, gamma=gamma, tau=tau, sigma=sigma)
        elif method == 'DRL':
            running = RUN_DRL(net=net, path_path=path_path, traffic_path=traffic_path, ca_path=ca_path, scale=scale, k=k, save_dir=save_dir, obj=obj, train_size=train_size, batch_size=batch_size, alpha=alpha, learning_rate=learning_rate, seed=seed, epoch=epoch, weight_decay=weight_decay, buf_capacity=buf_capacity, learning_rate_actor=learning_rate_actor, learning_rate_critic=learning_rate_critic, gamma=gamma, tau=tau, sigma=sigma)
        running.run()

    if metric_flag == True:
        mlu_path = f'results/{method}/{topo}/MLU.npy'
        tp_path = f'results/{method}/{topo}/TP.npy'
        closs_path = f'results/{method}/{topo}/CLoss.npy'
        mlu, tp, closs = running.check_metric(mlu_path, tp_path, closs_path)
        print(f'Testing Metrics MLU/TP/CongLoss: {mlu}/{tp}/{closs}')


def str2bool(value):
    if isinstance(value, bool):
        return value
    if value.lower() in ('true', 'yes', '1'):
        return True
    elif value.lower() in ('false', 'no', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")

# Command-line argument parsing
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run specific TE algorithm with parameters. Use '-h' to see specifics of each parameter.")
    
    # Mandatory
    parser.add_argument('--method', type=str, required=True, help="Mandatory: The method to use (ECMP, LP, FC-TE, LSTM-TE, CNN-TE, DRL-TE).")
    parser.add_argument('--train_flag', type=str2bool, required=True, help="Mandatory: Flag to indicate whether training specified DNNs or runining LP or ECMP algorithms.")
    parser.add_argument('--metric_flag', type=str2bool, required=True, help="Mandatory: Flag to derermine whether showing TE metrics (MLU, Throughput and Congestion Loss).")
    parser.add_argument('--topo', type=str, required=True, help="Mandatory: Topology name (e.g. G_test in our Repo).")

    # Optional
    parser.add_argument('--scale', type=float, default=1150383, help="Optional (Network): Scaling factor of the original traffic. (default: 1150383)")
    parser.add_argument('--k', type=int, default=4, help="Optional (Network): number of candidate paths per flow. (default: 4)")
    parser.add_argument('--obj', type=str, default='MLU', help="Optinal (Network, LP onlys): Objective function name (MLU, TP and CLoss). (default: MLU)")
    parser.add_argument('--train_size', type=int, default=4500, help="Optinal (Training): Size of traning set. (default: 4500)")
    parser.add_argument('--batch_size', type=int, default=64, help="Optinal (Training): Batch size. (default: 64)")
    parser.add_argument('--alpha', type=float, default=0, help="Optinal (Training): Alpha value to balance MLU and CLoss in loss function (MLU + alpha * CLoss). (default: 0)")
    parser.add_argument('--learning_rate', type=float, default=0.01, help="Optinal (Training): Learning rate. (default: 0.01)")
    parser.add_argument('--seed', type=int, default=76, help="Optinal (Training): Random seed. (default: 76)")
    parser.add_argument('--epoch', type=int, default=100, help="Optinal (Training): Number of training epochs. (default: 100)")
    parser.add_argument('--weight_decay', type=float, default=0.001, help="Optinal (Training): Weight decay for DNN training optimizer. (default: 0.001)")
    parser.add_argument('--buf_capacity', type=int, default=10000, help="Optinal (RL only): Replay buffer capacity. (default: 10000)")
    parser.add_argument('--learning_rate_actor', type=float, default=0.0001, help="Optinal (RL only): Learning rate for actor. (default: 0.0001)")
    parser.add_argument('--learning_rate_critic', type=float, default=0.0001, help="Optinal (RL only): Learning rate for critic. (default: 0.0001)")
    parser.add_argument('--gamma', type=float, default=0.99, help="Optinal (RL only): Gamma value -- Discount factor in DRL. (default: 0.99)")
    parser.add_argument('--tau', type=float, default=0.01, help="Optinal (RL only): Tau value -- Soft update rate for target network in DRL. (default: 0.01)")
    parser.add_argument('--sigma', type=float, default=0.6, help="Optinal (RL only): Sigma value -- Exploration noise in DRL. (default: 0.6)")
    
    args = parser.parse_args()

    main(
        method=args.method,
        train_flag=args.train_flag,
        metric_flag=args.metric_flag,
        topo=args.topo,
        scale=args.scale,
        k=args.k,
        obj=args.obj,
        train_size=args.train_size,
        batch_size=args.batch_size,
        alpha=args.alpha,
        learning_rate=args.learning_rate,
        seed=args.seed,
        epoch=args.epoch,
        buf_capacity=args.buf_capacity,
        learning_rate_actor=args.learning_rate_actor,
        learning_rate_critic=args.learning_rate_critic,
        gamma=args.gamma,
        tau=args.tau,
        sigma=args.sigma
    )