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': 20, 'U_init': None, 'init_mom': False}, 5, 'AltMin', '-ob'],

    'altmingd': [meta.apply_alt_min_gd, 
      {'N_step': 1000, 'U_init': None, 'stepsize': 0.2}, 5, 'AltMinGD', '--sm'],

    'mom': [meta.apply_method_of_moments, {}, 5, 'MoM', ':vr'],

    'mom2': [meta.apply_2nd_ord_method_of_moments, {}, 5, 'MoM2', ':^y'],

    'gd': [meta.apply_grad_descent, 
      {'N_step': 500, 'U_init': None, 'stepsize': 0.1}, 5, 'BM-GD', '-.Dg'],
  }

  d = 100
  r = 5
  t = 20 
  m_list = 5**np.arange(1, 8)
  sigma = 1.0
  dist_x = 'normal'
  dist_y = 'normal'

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

  metrics_list = expt_utils.algos_vs_var_metrics(
      d, r, t, m_list, sigma, algos_dict,
      dist_x=dist_x, dist_y=dist_y,
      loglog=True, metrics=metrics, plot_metrics=None,
      results_dir='results/neurips_submit/vs_m', script_file=script_file,
      )