import os
import subprocess
import json
import argparse
import math
import numpy as np

parser = argparse.ArgumentParser(description='manual to this script')
parser.add_argument('-dp', type=str, default="gaussian", help="DP mechanism")  # laplace or gaussian
parser.add_argument('-dataset', type=str, default="mnist", help="dataset")  # mnist or femnist
parser.add_argument('-q', type=float, default=0.01, help="local dataset sampling rate")  # 0~1

# the following parameters are used to find related optimal files generated by cal_optimal.py
parser.add_argument('-N', type=int, default=10, help="client number")  # 10, 20, 30, 40, 50
parser.add_argument('-b', type=int, default=10, help="ECP size, b randomly selected clients")  # at most N =50
parser.add_argument('-C', type=float, default=10000, help="gradient clip")
parser.add_argument('-lr', type=float, default=0.05, help="learning rate")
parser.add_argument('-T', type=int, default=100, help="total updates")  # global iterations
parser.add_argument('-E', type=int, default=1, help="local updates")  # local iterations
parser.add_argument('-mu', type=float, default=0.0001, help="weight decay")
parser.add_argument('-device_train', type=str, default='cuda:2')
parser.add_argument('-device_test', type=str, default='cuda:3')
parser.add_argument('-S', type=int, default=10, help="smooth time")
parser.add_argument('-percent', type=float, default=1.0, help="data size")

args = parser.parse_args()
lr = args.lr
S = args.S
T = args.T
E = args.E
b = args.b
mu = args.mu
C = args.C
N = args.N
dp = args.dp
q = args.q
S = args.S
device_tain = args.device_train
device_test = args.device_test
dataset = args.dataset
percent = args.percent

optimal_result_dir = os.path.join(os.getcwd(), 'estimate', dataset, 'clients_%d' % N, 'result')
log_path = os.path.join(os.getcwd(), 'log', '{}_{}_{}'.format(dataset, dp, N))
os.makedirs(log_path, exist_ok=True)

fileName = "{}_dp_train_lr={},T=*,E={},N={},b=*,C={},mu={},epsilon=*,q={}".format(dp, lr, 1, N, C, mu, q, percent)
mark = "{}_dp_train_lr={},T=*,E={},N={},b=*,C={},mu={},epsilon=*,q={},percent={}".format(dp, lr, 1, N, C, mu, q, percent)  # log mark


def logger(mark, process):
    with open(os.path.join(log_path, '{}.log'.format(mark)), 'wb') as f:
        for ln in process.stdout:
            f.write(ln)
            f.flush()
    for ln in process.stderr.readlines():
        print(ln)


def laplace_run(SCRIPT, T, b, epsilon, enalbe_print=True):
    print(SCRIPT + "-T {} -b {} -epsilon {} ".format(T, b, epsilon))
    process = subprocess.Popen(SCRIPT + "-T {} -b {} -epsilon {} ".format(T, b, epsilon),
                               shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1)
    logger(mark.replace("*", "{}").format(T, b, epsilon), process)
    print(mark.replace("*", "{}").format(T, b, epsilon))
    if enalbe_print:
        result_dir = os.path.join(os.getcwd(), 'result', dataset, 'clients_%d' % N)
        with open(os.path.join(result_dir, fileName.replace("*", "{}").format(T, b, epsilon) + '.loss.json')) as f:
            rs = np.array(json.load(f)).mean(axis=0)
            if len(rs) > 10:
                print(rs[:5], '...', rs[-5:])
            else:
                print(rs)


def gaussian_run(SCRIPT, T, b, epsilon, sigma, enalbe_print=True):
    print(SCRIPT + "-T {} -b {} -epsilon {} -sigma {} ".format(T, b, epsilon, sigma))
    process = subprocess.Popen(SCRIPT + "-T {} -b {} -epsilon {} -sigma {} ".format(T, b, epsilon, sigma),
                               shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1)
    logger(mark.replace("*", "{}").format(T, b, epsilon), process)
    print(mark.replace("*", "{}").format(T, b, epsilon))
    result_dir = os.path.join(os.getcwd(), 'result', dataset, 'clients_%d' % N)
    if enalbe_print:
        with open(os.path.join(result_dir, fileName.replace("*", "{}").format(T, b, epsilon) + '.loss.json')) as f:
            rs = np.array(json.load(f)).mean(axis=0)
            if len(rs) > 10:
                print(rs[:5], '...', rs[-5:])
            else:
                print(rs)


if __name__ == "__main__":
    file_name_opt = "{}_optimal.json".format(dp)
    file_name_params = "{}_estimate.hyper_params.json".format(dp)
    if not os.path.exists(os.path.join(optimal_result_dir, file_name_opt)):
        print(os.path.join(optimal_result_dir, file_name_opt))
        print("WARNING: optimal results are not generated, please run cal_optimal.py first!")
        exit(1)
    with open(os.path.join(optimal_result_dir, file_name_params)) as f:
        hyper_params = json.load(f)
    xi = hyper_params['xi_1'] if dp == 'laplace' else hyper_params['xi_2']
    print("Load hyper parameters: ", file_name_params)

    with open(os.path.join(optimal_result_dir, file_name_opt)) as f:
        params = json.load(f)
    print("Load optimal file: ", file_name_opt)
    # print(params)

    SCRIPT = "/home/magnolia/FedSGD-DPC/venv/bin/python dp_train.py " \
             "-dp {} -dataset {} -q {} -mu {} -C {} -xi {} -N {} -E {} -lr {} -S {} -device_train {} -device_test {} -percent {} ".\
        format(dp, dataset, q, mu, C, xi, N, 1, lr, S, device_tain, device_test, percent)

    SCRIPT2 = "/home/magnolia/FedSGD-DPC/venv/bin/python dp_train.py " \
              "-dp {} -dataset {} -q {} -mu {} -C {} -xi {} -N {} -E {} -lr {} -S {} -device_train {} -device_test {} -percent {} ". \
        format(dp, dataset, q, mu, 10000, xi, N, 1, lr, S, device_tain, device_test, percent)

    # without DP
    if dp == 'laplace':
        laplace_run(SCRIPT2, 100, N, -1, False)
    else:
        gaussian_run(SCRIPT2, 100, N, -1, 0, False)

    # with DP
    for epsilon, rs in params.items():
        # optimal
        epsilon = float(epsilon.split("=")[1])
        print("\n" + "x="*100)
        print("-set epsilon={}".format(epsilon))
        joint = rs['joint']
        if dp == 'laplace':
            T, b, U, Tl = math.floor(joint['opt_T']), math.floor(joint['opt_b']), joint['opt_U'], joint['opt_Tl']
            print("-"*50)
            print("optimal: b={}, T={}, U={:.5f}".format(b, T, U))
            laplace_run(SCRIPT, T, b, epsilon)
        else:
            T, b, U, Tl, sigma = math.floor(joint['opt_T']), math.floor(joint['opt_b']), joint['opt_U'], joint['opt_Tl'], joint['opt_sigma']
            print("-" * 50)
            print("optimal: b={}, T={}, U={:.5f}".format(b, T, U))
            gaussian_run(SCRIPT, T, b, epsilon, sigma)

        # baseline
        if dp == 'laplace':
            # # run fix b
            fix = rs['fix_b']
            for b, T, U, Tl in zip(fix['b_list'], fix['T_list'], fix['U_list'], fix['Tl_list']):
                b, T = math.ceil(b), math.ceil(T)
                print("-" * 50)
                print("fix b: b={}, T={}, U={:.5f}".format(b, T, U))
                laplace_run(SCRIPT, T, b, epsilon)

            # run fix T
            fix = rs['fix_T']
            for b, T, U, Tl in zip(fix['b_list'], fix['T_list'], fix['U_list'], fix['Tl_list']):
                b, T = math.ceil(b), math.ceil(T)
                print("-" * 50)
                print("fix T: T={}, b={}, U={:.5f}".format(T, b, U))
                if epsilon <= 1 and T <= 100:
                    laplace_run(SCRIPT, T, b, epsilon)
                else:
                    laplace_run(SCRIPT, T, b, epsilon)
        else:
            rr = rs['fix_b']
            for b, T, U, Tl, sigma in zip(rr['b_list'], rr['T_list'], rr['U_list'], rr['Tl_list'], rr['sigma_list']):
                b, T = math.ceil(b), math.ceil(T)
                print("-" * 50)
                print("b={}, T={}, U={:.5f}".format(b, T, U))
                gaussian_run(SCRIPT, T, b, epsilon, sigma)
