import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader, Subset, random_split
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions.dirichlet import Dirichlet
# import torch_kfac

import sys
import os
from pathlib import Path
import time
import random
import numpy as np
from numpy.random import MT19937
from numpy.random import RandomState, SeedSequence
import argparse
import collections
import json
import pickle
import pandas as pd
import matplotlib.pyplot as plt
# import seaborn as sns
from typing import Callable, Any,Literal
from dataclasses import dataclass,asdict,field,is_dataclass
import copy
import scipy
from scipy.stats import wilcoxon
import hashlib
import pytest
from itertools import product
from joblib import Parallel, delayed
import subprocess
from dacite import from_dict
import datetime
import argparse


from src import training, task,dataset,discriminator_v1_1,save_log,util,fed_learning, robust_grad_aggregation,worker_with_byzantine

from src.training import DnnArg,VggArg
from src.discriminator_v1_1 import WeightArgs,MvArgs,KbArgs,LinearArgs,GradArgs,DiscriminatorIFArgs,WeightUpdaterArg,DnnArgDiscrim
from src.worker_with_byzantine import NBWorkersArgs,NoisyWorkersArgs,NonIidClassWorkersArgs,NonIidCorrectWorkersArgs,WorkersArgs
from src.fed_learning import FedLearnArgs,TrainFLArg
from src.robust_grad_aggregation import IFArg,GradAggArg
from src.worker_with_byzantine import ConcatArgs,NBWorkersArgsRef


