# run.py

import json
import torch
import argparse
import logging
import os
# import wandb # Delete or comment out this line
import datetime
import numpy as np
import random  # Added import of random for set_seed

from pathlib import Path
import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from src.modules import *
from src.other_utils.utils import *

###########################################################################
# This logging configuration is independent; if you do not need to record model_accuracy.log, you can remove it.
# It is separate from the training curve data saved by the main program.
# logging.basicConfig(filename='model_accuracy.log',
#                     filemode='a',
#                     format='%(asctime)s - %(levelname)s - %(message)s',
#                     level=logging.INFO)

# def log_accuracy(accuracy):
#     logging.info(f'Model Accuracy: {accuracy:.4f}')

# model_accuracy = 0.85
# log_accuracy(model_accuracy)
# print("Accuracy has been recorded in model_accuracy.log")
##############################################################################################

ServerClass = {"async": AsyncServer,
               "buffer": BufferedServer,
               "concurrency": ConServer,
               "shuffle": ShuffleServer,
               "mifa": MIFAServer,
               "proposed": ProposedServer,
               "central": CentralServer}


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='cnn', help='neural network used in training')
    parser.add_argument('--batch_size', type=int, default=64, help='input batch size for training (default: 64)')
    parser.add_argument('--lr', type=float, default=0.05, help='learning rate (default: 0.01)')
    parser.add_argument('--lr_decay', type=float, default=1.0, help='learning rate decay factor (default: 1.0)')
    parser.add_argument('--lr_decay_interval', type=int, default=1, help='learning rate decay interval (default: 1)')
    parser.add_argument('--n_workers', type=int, default=10, help='number of workers')
    parser.add_argument('--alg', type=str, default='async',
                        choices=['async', 'buffer', 'concurrency', 'shuffle', 'mifa', 'proposed', 'central'])
    parser.add_argument('--tot_step', type=int, default=1000, help='number of total computation steps')
    parser.add_argument('--test_freq', type=int, default=1, help='frequency of testing the global model')
    parser.add_argument('--init_seed', type=int, default=42, help="Random seed")  #
    parser.add_argument('--logdir', type=str, required=False, default="./results/raw/", help='Directory for raw CSV log files')
    parser.add_argument('--modeldir', type=str, required=False, default="./results/checkpoints/", help='Directory for saved model checkpoints')
    parser.add_argument('--datadir', type=str, required=False, default="~/data/", help="Data directory")
    parser.add_argument('--log_file_name', type=str, default=None, help='The log file name')  #

    parser.add_argument('--dataset', type=str, default='fmnist', help='dataset used for training')
    parser.add_argument('--partition', type=str, default='homo', help='the data partitioning strategy')
    parser.add_argument('--n_sample', type=int, default=2000,
                        help='number of data samples on each device')  # TODO: delete
    parser.add_argument('--alpha', type=float, default=0.5,
                        help='The parameter for the dirichlet distribution for data partitioning')
    parser.add_argument('--in_dim', type=int, default=50, help='for synthetic dataset')
    parser.add_argument('--device', type=str, default='cuda:0', help='The device to run the program')
    parser.add_argument('--lamb', type=float, default=0.0, help='lambda for l2 regularization')
    parser.add_argument('--resume', default=False, action='store_true')
    parser.add_argument('--save_net', default=False, action='store_true')

    # delay setup
    parser.add_argument('--comput_dist', type=str, default="uniform", help='distribution of computation speeds')
    parser.add_argument('--tau_min', type=int, default=1, help='Maximal speed')
    parser.add_argument('--tau_max', type=int, default=5, help='Minimal speed')
    parser.add_argument('--tau_mean', type=int, default=1, help='Mean speed')
    parser.add_argument('--tau_std', type=int, default=1, help='Std of speed')
    parser.add_argument('--buffer_size', type=int, default=5,
                        help='Buffer size for buffered server; concurrency for concurrency server')
    parser.add_argument('--init_gradient', type=lambda x: str(x).lower() in ('1','true','yes','y'), default=True, help='Whether to initialize the gradient table before training (True/False).')

    parser.add_argument('--s', type=float, default=10)

    args = parser.parse_args()

    args.datadir = os.path.expanduser(args.datadir)

    if args.dataset == "synthetic":
        args.model = "linear"
    return args


def setup_logger(args):
    if args.comput_dist == "uniform":
        comput_text = f"{args.comput_dist}({args.tau_min},{args.tau_max})"
    elif args.comput_dist == "normal":
        comput_text = f"{args.comput_dist}({args.tau_mean},{args.tau_std})"

    # Dynamically generate alg_text based on args.alg and args.buffer_size
    if args.alg in ("buffer", "concurrency"):
        if args.alg == "buffer" and args.buffer_size == args.n_workers:
            alg_text = "Synchronous SGD"
        elif args.alg == "buffer" and args.buffer_size == 5:
            alg_text = "FedBuff"
        elif args.alg == "concurrency" and args.buffer_size == 5:
            alg_text = "Uniform ASGD"
        else:
            alg_text = f"{args.alg}({args.buffer_size})"
    else:
        alg_text = args.alg  # Keep unchanged, e.g., 'async', 'shuffle', 'mifa', 'proposed', 'central'

    # Convert alg_text into the format expected by draw_csv.py, e.g., async-lr0.01-normal(1,1)-alpha0.5
    file_alg_text = f"{args.alg}-lr{args.lr}"
    if args.alg in ("buffer", "concurrency"):
        file_alg_text = f"{args.alg}({args.buffer_size})-lr{args.lr}"

    time_text = datetime.datetime.now().strftime("%m%d-%H%M%S")
    mkdirs(args.logdir)

    if args.log_file_name is None:
        args.log_file_name = f'log_{file_alg_text}-n_workers{args.n_workers}-{comput_text}-dataset{args.dataset}-partition{args.partition}-model{args.model}-alpha{args.alpha}-seed{args.init_seed}'
        argument_path = args.log_file_name + '.json'
        # --- End of modification ---
    else:
        argument_path = args.log_file_name + '.json'

    with open(os.path.join(args.logdir, argument_path), 'w') as f:  # argument
        json.dump(str(args), f)
    log_path = args.log_file_name + '.log'  # log_path

    # print in the screen and save to log_file
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)
    logger = logging.getLogger(__name__)
    logging.basicConfig(level=logging.DEBUG, format='%(levelname)-8s %(message)s')  # print in the screen
    handler = logging.FileHandler(os.path.join(args.logdir, log_path))  # save to the file
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter('%(levelname)-8s %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    args.alg_text = alg_text  # This is used for the wandb name or other logging contexts and may differ from the file name
    args.comput_text = comput_text
    return logger, alg_text, comput_text


def set_seed(seed):
    torch.cuda.empty_cache()
    torch.random.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    random.seed(seed)  #
    np.random.seed(seed)  #
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    return


if __name__ == '__main__':
    args = get_args()

    # setup logger file
    # The alg_text and comput_text returned by setup_logger are used for logger information, not file names
    # log_file_name has already been set in args
    logger, alg_text_for_log, comput_text_for_log = setup_logger(args)

    # args.run_name was used by wandb; it is no longer needed now, but can be kept or removed for compatibility
    args.run_name = "{}-{}-alpha{}".format(alg_text_for_log, comput_text_for_log, args.alpha)

    # # wandb init # Delete or comment out the wandb init section
    # # wandb.init(
    # #     project="Delayed_SGD",
    # #     name=args.run_name,
    # #     config=args
    # # )
    # # logger.info("File name: %s" % args.log_file_name)

    set_seed(args.init_seed)  #

    Server = ServerClass[args.alg]
    server = Server(args, logger)

    server.run()

    # # wandb.finish() # Delete or comment out the wandb finish section