"""
Experiment on face data.
"""
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

from future import standard_library
standard_library.install_aliases()
from builtins import str
from builtins import range
from past.utils import old_div
import freqopttest.data as data
import freqopttest.tst as tst
import freqopttest.glo as glo
import freqopttest.util as util 
import freqopttest.kernel as kernel 
from . import exglobal
try:
    import pickle as pickle 
except:
    import pickle

# need independent_jobs package 
# https://github.com/karlnapf/independent-jobs
# The independent_jobs and freqopttest have to be in the globl search path (.bashrc)
import independent_jobs as inj
from independent_jobs.jobs.IndependentJob import IndependentJob
from independent_jobs.results.SingleResult import SingleResult
from independent_jobs.aggregators.SingleResultAggregator import SingleResultAggregator
from independent_jobs.engines.BatchClusterParameters import BatchClusterParameters
from independent_jobs.engines.SerialComputationEngine import SerialComputationEngine
from independent_jobs.engines.SlurmComputationEngine import SlurmComputationEngine
from independent_jobs.tools.Log import logger
import math
import autograd.numpy as np
import os
import sys 

def job_met_opt(sample_source, tr, te, r):
    """MeanEmbeddingTest with test locations optimzied."""
    # MeanEmbeddingTest. optimize the test locations
    with util.ContextTimer() as t:
        met_opt_options = {'n_test_locs': J, 'max_iter': 200, 
                'locs_step_size': 5.0, 'gwidth_step_size': 0.1, 'seed': r+92856,
                'tol_fun': 1e-4}
        test_locs, gwidth, info = tst.MeanEmbeddingTest.optimize_locs_width(tr, alpha, **met_opt_options)
        met_opt = tst.MeanEmbeddingTest(test_locs, gwidth, alpha)
        met_opt_test  = met_opt.perform_test(te)

    result = {'test_method': met_opt, 'test_result': met_opt_test, 'time_secs': t.secs}
    return result

def job_met_gwgrid(sample_source, tr, te, r):
    """MeanEmbeddingTest. Optimize only the Gaussian width with grid search
    Fix the test locations."""

    # optimize on the training set
    with util.ContextTimer() as t:
        T_randn = tst.MeanEmbeddingTest.init_locs_2randn(tr, J, seed=r+92856)
        med = util.meddistance(tr.stack_xy(), 1000)
        list_gwidth = np.hstack( ( (med**2) *(2.0**np.linspace(-5, 5, 40) ) ) )
        list_gwidth.sort()
        besti, powers = tst.MeanEmbeddingTest.grid_search_gwidth(tr, T_randn,
                list_gwidth, alpha)

        best_width2 = list_gwidth[besti]
        met_grid = tst.MeanEmbeddingTest(T_randn, best_width2, alpha)
        test_result = met_grid.perform_test(te)
    result = {'test_method': met_grid, 'test_result': test_result, 'time_secs': t.secs}
    return result

def job_scf_opt(sample_source, tr, te, r):
    """SmoothCFTest with frequencies optimized."""
    with util.ContextTimer() as t:
        op = {'n_test_freqs': J, 'max_iter': 200, 'freqs_step_size': 1.0, 
                'gwidth_step_size': 0.1, 'seed': r+92856, 'tol_fun': 1e-4}
        test_freqs, gwidth, info = tst.SmoothCFTest.optimize_freqs_width(tr, alpha, **op)
        scf_opt = tst.SmoothCFTest(test_freqs, gwidth, alpha)
        scf_opt_test = scf_opt.perform_test(te)
    
    result = {'test_method': scf_opt, 'test_result': scf_opt_test, 'time_secs': t.secs}
    return result

def job_scf_gwgrid(sample_source, tr, te, r):

    with util.ContextTimer() as t:
        rand_state = np.random.get_state()
        np.random.seed(r+92856)

        d = tr.dim()
        T_randn = np.random.randn(J, d)
        np.random.set_state(rand_state)

        # grid search to determine the initial gwidth
        mean_sd = tr.mean_std()
        scales = 2.0**np.linspace(-4, 4, 20)
        list_gwidth = np.hstack( (mean_sd*scales*(d**0.5), 2**np.linspace(-8, 8, 20) ))
        list_gwidth.sort()
        besti, powers = tst.SmoothCFTest.grid_search_gwidth(tr, T_randn,
                list_gwidth, alpha)
        # initialize with the best width from the grid search
        best_width = list_gwidth[besti]
        scf_gwgrid = tst.SmoothCFTest(T_randn, best_width, alpha)
        test_result = scf_gwgrid.perform_test(te)

    result = {'test_method': scf_gwgrid, 'test_result': test_result, 'time_secs': t.secs}
    return result

def job_quad_mmd(sample_source, tr, te, r):
    """Quadratic mmd with grid search to choose the best Gaussian width.
    One-sample U-statistic. This should NOT be used anymore."""
    # If n is too large, pairwise meddian computation can cause a memory error. 
            
    with util.ContextTimer() as t:
        med = util.meddistance(tr.stack_xy(), 1000)
        list_gwidth = np.hstack( ( (med**2) *(2.0**np.linspace(-4, 4, 40) ) ) )
        list_gwidth.sort()
        list_kernels = [kernel.KGauss(gw2) for gw2 in list_gwidth]

        # grid search to choose the best Gaussian width
        besti, powers = tst.QuadMMDTest.grid_search_kernel(tr, list_kernels, alpha)
        # perform test 
        best_ker = list_kernels[besti]
        mmd_test = tst.QuadMMDTest(best_ker, n_permute=1000, alpha=alpha, 
                use_1sample_U=True)
        test_result = mmd_test.perform_test(te)

    result = {'test_method': mmd_test, 'test_result': test_result, 'time_secs': t.secs}
    return result


def job_quad_mmd_2U(sample_source, tr, te, r):
    """Quadratic mmd with grid search to choose the best Gaussian width.
    Use two-sample U statistics to compute k(X,Y).
    """
    # If n is too large, pairwise meddian computation can cause a memory error. 
            
    with util.ContextTimer() as t:
        med = util.meddistance(tr.stack_xy(), 1000)
        list_gwidth = np.hstack( ( (med**2) *(2.0**np.linspace(-4, 4, 40) ) ) )
        list_gwidth.sort()
        list_kernels = [kernel.KGauss(gw2) for gw2 in list_gwidth]

        # grid search to choose the best Gaussian width
        besti, powers = tst.QuadMMDTest.grid_search_kernel(tr, list_kernels, alpha)
        # perform test 
        best_ker = list_kernels[besti]
        mmd_test = tst.QuadMMDTest(best_ker, n_permute=1000, alpha=alpha,
                use_1sample_U=False)
        test_result = mmd_test.perform_test(te)
    result = {'test_method': mmd_test, 'test_result': test_result, 'time_secs': t.secs}
    return result

def job_lin_mmd(sample_source, tr, te, r):
    """Linear mmd with grid search to choose the best Gaussian width."""
    # should be completely deterministic

    # If n is too large, pairwise meddian computation can cause a memory error. 
    with util.ContextTimer() as t:
        X, Y = tr.xy()
        Xr = X[:min(X.shape[0], 1000), :]
        Yr = Y[:min(Y.shape[0], 1000), :]
        
        med = util.meddistance(np.vstack((Xr, Yr)) )
        widths = [ (med*f) for f in 2.0**np.linspace(-1, 4, 40)]
        list_kernels = [kernel.KGauss( w**2 ) for w in widths]
        # grid search to choose the best Gaussian width
        besti, powers = tst.LinearMMDTest.grid_search_kernel(tr, list_kernels, alpha)
        # perform test 
        best_ker = list_kernels[besti]
        lin_mmd_test = tst.LinearMMDTest(best_ker, alpha)
        test_result = lin_mmd_test.perform_test(te)

    result = {'test_method': lin_mmd_test, 'test_result': test_result, 'time_secs': t.secs}
    return result


# Define our custom Job, which inherits from base class IndependentJob
class Ex5Job(IndependentJob):
   
    def __init__(self, aggregator, prob_label, rep, n, job_func):
        walltime = 60*59*24 
        memory = int(tr_proportion*n*1e-1) + 50

        IndependentJob.__init__(self, aggregator, walltime=walltime,
                               memory=memory)
        self.prob_label = prob_label
        self.rep = rep
        self.n = n
        self.job_func = job_func

    # we need to define the abstract compute method. It has to return an instance
    # of JobResult base class
    def compute(self):
        
        r = self.rep
        sample_source, nmax = get_sample_source(self.prob_label)
        d = sample_source.dim()
        job_func = self.job_func
        logger.info("computing. %s. r=%d "%(job_func.__name__, r ))

        tst_data = sample_source.sample(self.n, seed=r)
        tr, te = tst_data.split_tr_te(tr_proportion=tr_proportion, seed=r+20 )
        prob_label = self.prob_label
        job_result = job_func(sample_source, tr, te, r)

        # create ScalarResult instance
        result = SingleResult(job_result)
        # submit the result to my own aggregator
        self.aggregator.submit_result(result)
        logger.info("done. ex2: %s, r=%d "%(job_func.__name__, r))

        # save result
        func_name = job_func.__name__
        fname = '%s-%s-J%d_r%d_d%d_a%.3f_trp%.2f.p' \
                %(prob_label, func_name, J, r, d, alpha, tr_proportion)
        glo.ex_save_result(ex, job_result, prob_label, fname)


# This import is needed so that pickle knows about the classes
# pickle is used when collecting the results from the submitted jobs.
from freqopttest.ex.ex5_face import job_met_opt
from freqopttest.ex.ex5_face import job_met_gwgrid
from freqopttest.ex.ex5_face import job_scf_opt
from freqopttest.ex.ex5_face import job_scf_gwgrid
from freqopttest.ex.ex5_face import job_quad_mmd
from freqopttest.ex.ex5_face import job_quad_mmd_2U
from freqopttest.ex.ex5_face import job_lin_mmd
from freqopttest.ex.ex5_face import Ex5Job

#--- experimental setting -----
ex = 5

# number of test locations / test frequencies J
J = 1
alpha = 0.01
tr_proportion = 0.5
# repetitions 
reps = 500
#method_job_funcs = [ job_met_opt, job_scf_opt, job_lin_mmd, job_hotelling]
#method_job_funcs = [ job_met_opt, job_scf_opt, job_quad_mmd_2U,
#        job_lin_mmd]
method_job_funcs = [ job_met_opt, job_met_gwgrid, job_scf_opt, job_scf_gwgrid,
       job_quad_mmd_2U, job_lin_mmd]
#method_job_funcs = [ job_met_opt]
#method_job_funcs = [ job_lin_mmd, job_hotelling]

# If is_rerun==False, do not rerun the experiment if a result file for the current
# setting of (ni, r) already exists.
is_rerun = False
#---------------------------

label2fname = {
    'S64_HANESU_AFANDI':'S64_HANESU_AFANDI.p',
    'S48_HANESU_AFANDI':'S48_HANESU_AFANDI.p',
    'crop48_HANESU_AFANDI':'crop48_HANESU_AFANDI.p',
    'crop48_h0':'crop48_h0.p',
        }

def get_sample_source(prob_label):
    """Return a (SampleSource, n) representing the problem"""

    if prob_label not in label2fname:
        raise ValueError('Unknown problem label. Need to be one of %s'%str(list(label2fname.keys())) )

    if prob_label == 'crop48_h0':
        one_sample = glo.load_data_file(label2fname[prob_label])
        n_max = old_div(one_sample.shape[0],2)
        return data.SSNullResample(one_sample), n_max

    fname = label2fname[prob_label]
    tst_data = glo.load_data_file(fname)
    n = tst_data.X.shape[0]
    ss = data.SSResample(tst_data)
    return ss, n

def main():
    if len(sys.argv) != 2:
        print('Usage: %s problem_label'%sys.argv[0])
        sys.exit(1)
    prob_label = sys.argv[1]
    run_dataset(prob_label)

def run_dataset(prob_label):
    """Run the experiment"""
    sample_source, n = get_sample_source(prob_label)

    # ///////  submit jobs //////////
    # create folder name string
    home = os.path.expanduser("~")
    foldername = os.path.join(home, "freqopttest_slurm", 'e%d'%ex)
    logger.info("Setting engine folder to %s" % foldername)

    # create parameter instance that is needed for any batch computation engine
    logger.info("Creating batch parameter instance")
    batch_parameters = BatchClusterParameters(
        foldername=foldername, job_name_base="e%d_"%ex, parameter_prefix="")

    # Use the following line if Slurm queue is not used.
    #engine = SerialComputationEngine()
    engine = SlurmComputationEngine(batch_parameters, do_clean_up=True)
    n_methods = len(method_job_funcs)
    # repetitions x  #methods
    aggregators = np.empty((reps, n_methods ), dtype=object)
    d = sample_source.dim()
    for r in range(reps):
        for mi, f in enumerate(method_job_funcs):
            # name used to save the result
            func_name = f.__name__
            fname = '%s-%s-J%d_r%d_d%d_a%.3f_trp%.2f.p' \
                %(prob_label, func_name, J, r, d, alpha, tr_proportion)
            if not is_rerun and glo.ex_file_exists(ex, prob_label, fname):
                logger.info('%s exists. Load and return.'%fname)
                test_result = glo.ex_load_result(ex, prob_label, fname)
                sra = SingleResultAggregator()
                sra.submit_result(SingleResult(test_result))

                aggregators[r, mi] = sra
            else:
                # result not exists or rerun
                job = Ex5Job(SingleResultAggregator(), prob_label, r, n, f)
                agg = engine.submit_job(job)
                aggregators[r, mi] = agg

    # let the engine finish its business
    logger.info("Wait for all call in engine")
    engine.wait_for_all()

    # ////// collect the results ///////////
    logger.info("Collecting results")
    test_results = np.empty((reps, n_methods), dtype=object)
    for r in range(reps):
        for mi, f in enumerate(method_job_funcs):
            logger.info("Collecting result (%s, r=%d)" % (f.__name__, r ))
            # let the aggregator finalize things
            aggregators[r, mi].finalize()

            # aggregators[i].get_final_result() returns a SingleResult instance,
            # which we need to extract the actual result
            test_result = aggregators[r, mi].get_final_result().result
            test_results[r, mi] = test_result

            func_name = f.__name__
            fname = '%s-%s-J%d_r%d_d%d_a%.3f_trp%.2f.p' \
                %(prob_label, func_name, J, r, d, alpha, tr_proportion)
            glo.ex_save_result(ex, test_result, prob_label, fname)

    func_names = [f.__name__ for f in method_job_funcs]
    func2labels = exglobal.get_func2label_map()
    method_labels = [func2labels[f] for f in func_names if f in func2labels]
    # save results 
    results = {'results': test_results, 'n': n, 'data_fname':label2fname[prob_label],
            'alpha': alpha, 'J': J, 'sample_source': sample_source, 
            'tr_proportion': tr_proportion, 'method_job_funcs': method_job_funcs, 
            'prob_label': prob_label, 'method_labels': method_labels}
    
    # class name 
    fname = 'ex%d-%s-me%d_J%d_rs%d_nma%d_d%d_a%.3f_trp%.2f.p' \
        %(ex, prob_label, n_methods, J, reps, n, d, alpha, tr_proportion)
    glo.ex_save_result(ex, results, fname)
    logger.info('Saved aggregated results to %s'%fname)


if __name__ == '__main__':
    main()



