import numpy as np
from optimizer import inexact_newton, localSGD, fedAC, newton
from function import logistic
from data import loada9a
from datetime import datetime
import os

# this defines some functions for tuning and re-running the optimization algorithms


def find_gap(R):
    if R == 1:
        gap = 1
    elif R == 5:
        gap = 1
    elif R == 10:
        gap = 2
    elif R == 25:
        gap = 5
    elif R == 50:
        gap = 5
    elif R == 75:
        gap = 5
    elif R == 100:
        gap = 10
    elif R == 500:
        gap = 25
    elif R == 1000:
        gap = 50
    elif R == 10000:
        gap = 500

    return gap


def tune_lr(alg, alg_name, lrs, **algargs):
    R = algargs["R"]
    print(f"[*] Running {alg_name} for R = {R}")

    # Identify a reasonable gap
    gap = find_gap(R)

    # Obtaining the best learning rate
    min_loss = np.Infinity
    min_lr = np.Infinity
    min_W = None
    min_losses = None

    for lr in lrs:
        print(f"[*] Using learning rate {lr}")
        algargs['lr'] = lr
        algargs['gap'] = gap
        W, losses = alg(**algargs)
        best_loss = np.min(losses)
        print(f"[+] Obtained best loss {best_loss}")

        if best_loss < min_loss:
            min_lr = lr
            min_loss = best_loss
            min_W = W
            min_losses = losses

    print(f"[+] Best lr={min_lr} obtaining loss={min_loss}")

    return min_lr, min_loss, min_W, min_losses


def repeat_run(alg, alg_name, rep, **algargs):
    R = algargs["R"]
    lr = algargs["lr"]
    print(f"[*] Running {alg_name} for R = {R} with lr = {lr}")
    Losses = []
    for i in range(rep):
        print(f"[*] In repetetion number {i+1}")
        # Identify a reasonable gap
        gap = find_gap(R)
        algargs['gap'] = gap
        W, losses = alg(**algargs)
        best_loss = np.min(losses)
        print(f"[+] Obtained best loss {best_loss}")
        Losses.append(best_loss)
    return Losses
