import numpy as np
import pandas as pd
import tensorflow as tf
import time
import argparse
from matplotlib import pyplot as plt

from random_effect_logistic_regression_utils import generate_data, timestamp
import sys
sys.path.append('../models')
from random_effect_logistic_regression import random_effect_logistic_regression as RELR
from random_effect_logistic_regression import bayesian_random_effect_logistic_regression as BRELR

# Turn GPUs off
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

EPS = 1e-6


def d(f, params):
    # Take the derivative of f
    # returned value is a function df/d[beta0, beta, alpha]
    def df(x,y,level):
        with tf.GradientTape(persistent=True) as g:
            g.watch(params)
            target = f(x,y,level)
        est0 = g.gradient(target, params)
        est = np.concatenate([e.numpy().reshape([-1]) for e in est0], axis=0)
        return est
    return df

def get_mlmc_cost(N, max_level, b, w0):
    # compute the cost of MLMC estimation 
    # when the size of x (and that of y) is N
    if max_level==0:
        levels = np.array([0])
        weights = np.array([1.])
    else:
        weights = 2.**(-(b+1)/2*np.arange(max_level))
        weights /= sum(weights)
        weights = np.concatenate([[w0], (1-w0)*weights])
        levels = np.arange(max_level+1)
    cost = N * weights[0] + N * sum( weights[1:] * (2**levels[1:] + 2**(levels[1:]-1)) )
    return cost


# Argparse
parser = argparse.ArgumentParser()
parser.add_argument("--output_file", default="../../out/random_effect_logistic_regression/learning_curve_{}.csv".format(timestamp()), type=str,
    help="Output file name. \nAn example usage: `python random_effect_logistic_regression_learning_log.py --output_file example.csv`")
args = parser.parse_args()

input_files = args.input_files.split()
output_file = args.output_file
x_axis = args.x_axis
time_discretization = args.time_discretization



### Initializations
N_total = 100000
B,T,D = (1000, 2, 3) if tf.test.is_gpu_available() else (100, 2, 3)#(100, 2, 3)

cost_nmc  = B * 2**9
cost_mlmc = get_mlmc_cost(B, max_level=9, b=1.8, w0=0.9)
cost_sumo = B * 9
B_mlmc = np.math.ceil(B * (cost_nmc / cost_mlmc))
B_sumo = np.math.ceil(B * (cost_nmc / cost_sumo))


alpha = np.float64(1.)
beta0 = np.float64(0.)
beta  = np.array([0.25, 0.50, 0.75]) #np.random.randn(D) / np.sqrt(D)
model = RELR(D=D)
# True model parameters

X,Y,_ = generate_data(N_total, D, T, beta0, beta, alpha)

objectives = {
    "iwelbo1": lambda x, y: model.IWELBO(x, y, n_MC=1),
    "iwelbo8": lambda x, y: model.IWELBO(x, y, n_MC=8),
    "iwelbo64": lambda x, y: model.IWELBO(x, y, n_MC=64),
    "iwelbo512": lambda x, y: model.IWELBO(x, y, n_MC=512),
    "iwelbo512_mlmc": lambda x, y: model.IWELBO_MLMC(x, y, max_level=9, b=1.8, w0=0.9, randomize=False),
    "iwelbo512_randmlmc": lambda x, y: model.IWELBO_MLMC(x, y, max_level=9, b=1.8, w0=0.9, randomize=True),
    "iwelbo512_sumo": lambda x, y: model.IWELBO_SUMO(x, y, K_max=512),
    "jvi8": lambda x, y: model.JVI_IWELBO(x, y, n_MC=8),
    "jvi64": lambda x, y: model.JVI_IWELBO(x, y, n_MC=64),
    "jvi512": lambda x, y: model.JVI_IWELBO(x, y, n_MC=512),
}

# for parallelization
#obj_id = int(input())
#objectives = {k:objectives[k] for i, k in enumerate(objectives.keys()) if i == obj_id}
#print(objectives)

n_train_steps = {
    "iwelbo1": 2000,
    "iwelbo8": 2000,
    "iwelbo64": 2000,
    "iwelbo512": 17000,
    "iwelbo512_mlmc": 3000,
    "iwelbo512_randmlmc": 3000,
    "iwelbo512_sumo": 2000,
    "jvi8": 2000,
    "jvi64": 2000,
    "jvi512": 17000,
}

data = [] 

n_repeat = 100 # TODO: change to 20
params_repeated = {name:[] for name in objectives.keys()}

for name, obj in objectives.items():
for i in range(n_repeat):
    print("training {}.... #iter:{} ".format(name,i))
    # initialize parameters
    model.beta0 = tf.Variable(0.0,  dtype=tf.float64)
    model.beta  = tf.Variable(np.zeros([model.D]),  dtype=tf.float64)
    model.alpha = tf.Variable(0.0,  dtype=tf.float64)
    # pointers to the parameters of trained model
    params_list = [
        model.beta0, 
        model.beta,
        model.alpha
    ]
    optimizer = tf.keras.optimizers.Adam(0.005)

    # Training
    start = time.time()

    for t in range(n_train_steps[name] + 1):

        # Balance the cost of mlmc and nmc when level=9 (n_MC=512)
        # by changing the batch size adoptively
        if 'mlmc' in name:
            batch = np.random.choice(np.arange(N_total), B_mlmc)
        elif 'sumo' in name:
            batch = np.random.choice(np.arange(N_total), B_sumo)
        else:
            batch = np.random.choice(np.arange(N_total), B)
        x = X[batch]
        y = Y[batch]

        # Train step
        with tf.GradientTape() as g:
            g.watch(params_list)
            loss = - obj(x, y)
        dparams = g.gradient(loss, params_list)
        optimizer.apply_gradients(zip(dparams, params_list))

        # Take a log
        if t%5==0:
            data.append({
                "objective": name,
                "#iter": i,
                "step": t,
                "elapsed time": time.time() - start,
                "alpha": model.alpha.numpy(),
                "beta0": model.beta0.numpy(),
                "beta1": model.beta.numpy()[0],
                "beta2": model.beta.numpy()[1],
                "beta3": model.beta.numpy()[2],
                "squared error": sum(
                    np.concatenate([
                        [alpha - model.alpha.numpy()],
                        [beta0 - model.beta0.numpy()],
                        beta - model.beta.numpy()
                    ]) ** 2
                )
            })

        if t%200==0 and i == 0:
            print("#iter: {},\tloss: {}".format(t, loss.numpy()))

print()


print("\n======== Results ========\n")
data = pd.DataFrame(
data=data,
columns = [
    "objective", "#iter", "elapsed time", "step", 
    "alpha", "beta0", "beta1", "beta2", "beta3", "squared error"
]
)
print(data)
data.to_csv(output_file)
print("\nSaved the results to:\n{}".format(output_file))


import numpy as np
import pandas as pd
import tensorflow as tf
import time
import argparse
from matplotlib import pyplot as plt

from random_effect_logistic_regression_utils import generate_data, timestamp
import sys
sys.path.append('../models')
from random_effect_logistic_regression import random_effect_logistic_regression as RELR
from random_effect_logistic_regression import bayesian_random_effect_logistic_regression as BRELR

# Turn GPUs off
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

EPS = 1e-6


def d(f, params):
    # Take the derivative of f
    # returned value is a function df/d[beta0, beta, alpha]
    def df(x,y,level):
        with tf.GradientTape(persistent=True) as g:
            g.watch(params)
            target = f(x,y,level)
        est0 = g.gradient(target, params)
        est = np.concatenate([e.numpy().reshape([-1]) for e in est0], axis=0)
        return est
    return df

def get_mlmc_cost(N, max_level, b, w0):
    # compute the cost of MLMC estimation 
    # when the size of x (and that of y) is N
    if max_level==0:
        levels = np.array([0])
        weights = np.array([1.])
    else:
        weights = 2.**(-(b+1)/2*np.arange(max_level))
        weights /= sum(weights)
        weights = np.concatenate([[w0], (1-w0)*weights])
        levels = np.arange(max_level+1)
    cost = N * weights[0] + N * sum( weights[1:] * (2**levels[1:] + 2**(levels[1:]-1)) )
    return cost


# Argparse
parser = argparse.ArgumentParser()
parser.add_argument("--output_file", default="../../out/random_effect_logistic_regression/learning_curve_{}.csv".format(timestamp()), type=str,
    help="Output file name. \nAn example usage: `python random_effect_logistic_regression_learning_log.py --output_file example.csv`")
args = parser.parse_args()

input_files = args.input_files.split()
output_file = args.output_file
x_axis = args.x_axis
time_discretization = args.time_discretization



### Initializations
N_total = 100000
B,T,D = (1000, 2, 3) if tf.test.is_gpu_available() else (100, 2, 3)#(100, 2, 3)

cost_nmc  = B * 2**9
cost_mlmc = get_mlmc_cost(B, max_level=9, b=1.8, w0=0.9)
cost_sumo = B * 9
B_mlmc = np.math.ceil(B * (cost_nmc / cost_mlmc))
B_sumo = np.math.ceil(B * (cost_nmc / cost_sumo))


alpha = np.float64(1.)
beta0 = np.float64(0.)
beta  = np.array([0.25, 0.50, 0.75]) #np.random.randn(D) / np.sqrt(D)
model = RELR(D=D)
# True model parameters

X,Y,_ = generate_data(N_total, D, T, beta0, beta, alpha)

objectives = {
    "iwelbo1": lambda x, y: model.IWELBO(x, y, n_MC=1),
    "iwelbo8": lambda x, y: model.IWELBO(x, y, n_MC=8),
    "iwelbo64": lambda x, y: model.IWELBO(x, y, n_MC=64),
    "iwelbo512": lambda x, y: model.IWELBO(x, y, n_MC=512),
    "iwelbo512_mlmc": lambda x, y: model.IWELBO_MLMC(x, y, max_level=9, b=1.8, w0=0.9, randomize=False),
    "iwelbo512_randmlmc": lambda x, y: model.IWELBO_MLMC(x, y, max_level=9, b=1.8, w0=0.9, randomize=True),
    "iwelbo512_sumo": lambda x, y: model.IWELBO_SUMO(x, y, K_max=512),
    "jvi8": lambda x, y: model.JVI_IWELBO(x, y, n_MC=8),
    "jvi64": lambda x, y: model.JVI_IWELBO(x, y, n_MC=64),
    "jvi512": lambda x, y: model.JVI_IWELBO(x, y, n_MC=512),
}

# for parallelization
#obj_id = int(input())
#objectives = {k:objectives[k] for i, k in enumerate(objectives.keys()) if i == obj_id}
#print(objectives)

n_train_steps = {
    "iwelbo1": 2000,
    "iwelbo8": 2000,
    "iwelbo64": 2000,
    "iwelbo512": 17000,
    "iwelbo512_mlmc": 3000,
    "iwelbo512_randmlmc": 3000,
    "iwelbo512_sumo": 2000,
    "jvi8": 2000,
    "jvi64": 2000,
    "jvi512": 17000,
}

data = [] 

n_repeat = 100 # TODO: change to 20
params_repeated = {name:[] for name in objectives.keys()}

for name, obj in objectives.items():
for i in range(n_repeat):
    print("training {}.... #iter:{} ".format(name,i))
    # initialize parameters
    model.beta0 = tf.Variable(0.0,  dtype=tf.float64)
    model.beta  = tf.Variable(np.zeros([model.D]),  dtype=tf.float64)
    model.alpha = tf.Variable(0.0,  dtype=tf.float64)
    # pointers to the parameters of trained model
    params_list = [
        model.beta0, 
        model.beta,
        model.alpha
    ]
    optimizer = tf.keras.optimizers.Adam(0.005)

    # Training
    start = time.time()

    for t in range(n_train_steps[name] + 1):

        # Balance the cost of mlmc and nmc when level=9 (n_MC=512)
        # by changing the batch size adoptively
        if 'mlmc' in name:
            batch = np.random.choice(np.arange(N_total), B_mlmc)
        elif 'sumo' in name:
            batch = np.random.choice(np.arange(N_total), B_sumo)
        else:
            batch = np.random.choice(np.arange(N_total), B)
        x = X[batch]
        y = Y[batch]

        # Train step
        with tf.GradientTape() as g:
            g.watch(params_list)
            loss = - obj(x, y)
        dparams = g.gradient(loss, params_list)
        optimizer.apply_gradients(zip(dparams, params_list))

        # Take a log
        if t%5==0:
            data.append({
                "objective": name,
                "#iter": i,
                "step": t,
                "elapsed time": time.time() - start,
                "alpha": model.alpha.numpy(),
                "beta0": model.beta0.numpy(),
                "beta1": model.beta.numpy()[0],
                "beta2": model.beta.numpy()[1],
                "beta3": model.beta.numpy()[2],
                "squared error": sum(
                    np.concatenate([
                        [alpha - model.alpha.numpy()],
                        [beta0 - model.beta0.numpy()],
                        beta - model.beta.numpy()
                    ]) ** 2
                )
            })

        if t%200==0 and i == 0:
            print("#iter: {},\tloss: {}".format(t, loss.numpy()))

print()


print("\n======== Results ========\n")
data = pd.DataFrame(
data=data,
columns = [
    "objective", "#iter", "elapsed time", "step", 
    "alpha", "beta0", "beta1", "beta2", "beta3", "squared error"
]
)
print(data)
data.to_csv(output_file)
print("\nSaved the results to:\n{}".format(output_file))


# Turn GPUs off # Take the derivative of f # returned value is a function df/d[beta0, beta, alpha] # compute the cost of MLMC estimation # when the size of x (and that of y) is N # Argparse ### Initializations #(100, 2, 3) #np.random.randn(D) / np.sqrt(D) # True model parameters # for parallelization #obj_id = int(input()) #objectives = {k:objectives[k] for i, k in enumerate(objectives.keys()) if i == obj_id} #print(objectives) # TODO: change to 20 #iter:{} ".format(name,i)) # initialize parameters # pointers to the parameters of trained model # Training # Balance the cost of mlmc and nmc when level=9 (n_MC=512) # by changing the batch size adoptively # Train step # Take a log