import numpy as np

import importlib

import meta_learn as meta
importlib.reload(meta)
import expt_utils as expt_utils
importlib.reload(expt_utils)

if __name__ == '__main__':
  importlib.reload(meta)
  importlib.reload(expt_utils)

  algos_dict = {
    'altmin': [meta.apply_alt_min, 
      {'N_step': 30, 'U_init': None, 'init_mom': False}, 6, 'AltMin', '-b'],

    'altmingd': [meta.apply_alt_min_gd, 
      {'N_step': 100, 'U_init': None, 'stepsize': 0.2}, 6, 'AltMinGD', '--m'],

    'gd': [meta.apply_grad_descent, 
      {'N_step': 100, 'U_init': None, 'stepsize': 0.2}, 6, 'BM-GD', '-.g'],
  }

  d = 100
  r = 5
  t = d*4
  m = r**2
  sigma = 1.0
  dist_x = 'normal'
  dist_y = 'normal'

  metrics = None # ['dist_U']
  script_file = 'k_expts.py'

  for homogeneity in [0.0, 0.4, 0.8]:
    metrics_list = expt_utils.algos_vs_var_metrics(
        d, r, t, m, sigma, algos_dict,
        homogeneity=homogeneity,
        dist_x=dist_x, dist_y=dist_y,
        loglog=None, ylims=(0.1, 1.0), #xlims=(0.0, 1.0),
        metrics=metrics, plot_metrics=None,
        results_dir='results/neurips_submit/vs_k', script_file=script_file,
        )