import argparse
import logging
import numpy as np
import os
import pandas as pd
import random
import subprocess
from pathlib import Path

from hyperopt import hp
from hyperopt.pyll.stochastic import sample

from hfta.hfht import (tune_hyperparameters, attach_common_args,
                       rearrange_algorithm_kwargs, handle_integers,
                       generate_fusible_param_flags, generate_nonfusible_param)
from hfta.workflow import extract_logging_level
from hfta.hfht.utils import fuse_dicts


def main(args):
  random.seed(args.seed)
  np.random.seed(args.seed)
  rng_state = np.random.RandomState(seed=args.seed)

  fusibles = {
      'lr': hp.uniform('lr', 0.0001, 0.01),
      'beta1': hp.uniform('beta1', 0.001, 0.999),
      'beta2': hp.uniform('beta2', 0.001, 0.999),
      'weight_decay': hp.uniform('weight_decay', 0.0, 0.5),
      'gamma': hp.uniform('gamma', 0.1, 0.9),
      'step_size': hp.choice('step_size', (5, 10, 20, 40)),
  }
  nonfusibles = {
      'batch_size': hp.choice('batch_size', (8, 16, 32)),
      'feature_transform': hp.choice('feature_transform', (True, False)),
  }

  def _run(results_dir, epochs, iters_per_epoch, params, env_vars=None):
    # Build the cmd.
    cmd = [
        'python',
        'train_classification.py',
        '--epochs',
        str(epochs),
        '--iters-per-epoch',
        str(iters_per_epoch),
        '--dataset',
        args.dataset,
        '--dataset_type',
        args.dataset_type,
        '--num_points',
        str(args.num_points),
        '--device',
        args.device,
        '--eval',
        '--seed',
        str(args.seed),
        '--batch_size',
        str(generate_nonfusible_param(params, 'batch_size')),
    ]
    if results_dir is not None:
      cmd.extend(['--outf', results_dir])
    if generate_nonfusible_param(params, 'feature_transform'):
      cmd.append('--feature_transform')
    cmd.extend(
        generate_fusible_param_flags(
            params,
            ['lr', 'beta1', 'beta2', 'weight_decay', 'gamma', 'step_size'],
        ))
    if args.mode == 'hfta':
      cmd.append('--hfta')
    if args.amp:
      cmd.append('--amp')

    # Launch the training process.
    succeeded = True
    try:
      logging.info('--> Running cmd = {}'.format(cmd))
      subprocess.run(
          cmd,
          stdout=subprocess.DEVNULL if results_dir is None else open(
              os.path.join(results_dir, 'stdout.txt'),
              'w',
          ),
          stderr=subprocess.DEVNULL if results_dir is None else open(
              os.path.join(results_dir, 'stderr.txt'),
              'w',
          ),
          check=True,
          cwd=os.path.join(
              os.path.abspath(os.path.expanduser(os.path.dirname(__file__))),
              '../pointnet/'),
          env=env_vars,
      )
    except subprocess.CalledProcessError as e:
      logging.error(e)
      succeeded = False
    return succeeded

  def try_params(ids, epochs, params, env_vars=None):
    """ Running the training process for pointnet classification task.

    Args:
      ids: Either a single int ID (for serial), or a list of IDs (for HFTA).
      epochs: number of epochs to run.
      params: maps hyperparameter name to its value(s). For HFTA, the values are
        provided as a list.
      env_vars: optional, dict(str, str) that includes extra environment that
        needs to be forwarded to the subprocess call

    Returns:
      result(s): A single result dict for serial or a list of result dicts for
        HFTA in the same order as ids.
      early_stop(s): Whether the training process early stopped. A single bool
        for serial or a list of bools for HFTA in the same order as ids.
    """
    epochs = int(round(epochs))
    ids_str = (','.join([str(i) for i in ids]) if isinstance(
        ids,
        (list, tuple),
    ) else str(ids))
    # Allocate result dir.
    results_dir = os.path.join(args.outdir, ids_str)
    Path(results_dir).mkdir(parents=True, exist_ok=True)
    # Run training.
    succeeded = _run(
        results_dir,
        epochs,
        args.iters_per_epoch,
        params,
        env_vars=env_vars,
    )
    if not succeeded:
      raise RuntimeError('_run failed!')
    # Gather the results.
    results_frame = pd.read_csv(os.path.join(results_dir, 'eval.csv'))
    if isinstance(ids, (list, tuple)):
      results = [{'acc': acc} for acc in results_frame['acc'].tolist()]
      assert len(results) == len(ids)
      return results, [False] * len(ids)
    else:
      return {'acc': results_frame['acc'][0]}, False

  def dry_run(
      B=None,
      nonfusibles_kvs=None,
      epochs=None,
      iters_per_epoch=None,
      env_vars=None,
  ):
    params = [{
        **handle_integers(sample(fusibles, rng=rng_state)),
        **nonfusibles_kvs
    } for _ in range(max(B, 1))]
    if B > 0:
      params = fuse_dicts(params)
    else:
      params = params[0]
    return _run(None, epochs, iters_per_epoch, params, env_vars=env_vars)

  tune_hyperparameters(
      space={
          **fusibles,
          **nonfusibles
      },
      try_params_callback=try_params,
      dry_run_callback=dry_run,
      mode=args.mode,
      algorithm=args.algorithm,
      nonfusibles=nonfusibles.keys(),
      dry_run_repeats=args.dry_run_repeats,
      dry_run_epochs=args.dry_run_epochs,
      dry_run_iters_per_epoch=args.dry_run_iters_per_epoch,
      metric='acc',
      goal='max',
      algorithm_configs={
          'hyperband': args.hyperband_kwargs,
          'random': args.random_kwargs,
      },
      seed=args.seed,
      outdir=args.outdir,
  )


def attach_args(parser=argparse.ArgumentParser()):
  parser.add_argument(
      '--workers',
      type=int,
      help='number of data loading workers',
      default=4,
  )
  parser.add_argument(
      '--iters-per-epoch',
      type=int,
      default=int(1e9),
      help='number of epochs to train for',
  )
  parser.add_argument('--dataset', type=str, required=True, help="dataset path")
  parser.add_argument(
      '--dataset-type',
      type=str,
      default='shapenet',
      help="dataset type shapenet|modelnet40",
  )
  parser.add_argument(
      '--num-points',
      type=int,
      default=2500,
      help='num of points for dataset',
  )
  parser.add_argument(
      '--device',
      type=str,
      default='cuda',
      choices=['cpu', 'cuda', 'xla'],
      help="the device where this test is running",
  )
  parser.add_argument(
      '--amp',
      default=False,
      action='store_true',
      help='Enable AMP; only used when --device is cuda',
  )
  parser = attach_common_args(parser)
  return parser


if __name__ == '__main__':
  args = attach_args().parse_args()
  rearrange_algorithm_kwargs(args)
  logging.basicConfig(level=extract_logging_level(args))
  args.outdir = os.path.abspath(os.path.expanduser(args.outdir))
  args.dataset = os.path.abspath(os.path.expanduser(args.dataset))
  main(args)

[((21, 2, 21, 24), 'random.seed', 'random.seed', ({(21, 14, 21, 23): 'args.seed'}, {}), '(args.seed)', False, 'import random\n'), ((22, 2, 22, 27), 'numpy.random.seed', 'np.random.seed', ({(22, 17, 22, 26): 'args.seed'}, {}), '(args.seed)', True, 'import numpy as np\n'), ((23, 14, 23, 51), 'numpy.random.RandomState', 'np.random.RandomState', (), '', True, 'import numpy as np\n'), ((185, 23, 185, 48), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((224, 11, 224, 37), 'hfta.hfht.attach_common_args', 'attach_common_args', ({(224, 30, 224, 36): 'parser'}, {}), '(parser)', False, 'from hfta.hfht import tune_hyperparameters, attach_common_args, rearrange_algorithm_kwargs, handle_integers, generate_fusible_param_flags, generate_nonfusible_param\n'), ((230, 2, 230, 34), 'hfta.hfht.rearrange_algorithm_kwargs', 'rearrange_algorithm_kwargs', ({(230, 29, 230, 33): 'args'}, {}), '(args)', False, 'from hfta.hfht import tune_hyperparameters, attach_common_args, rearrange_algorithm_kwargs, handle_integers, generate_fusible_param_flags, generate_nonfusible_param\n'), ((26, 12, 26, 42), 'hyperopt.hp.uniform', 'hp.uniform', ({(26, 23, 26, 27): '"""lr"""', (26, 29, 26, 35): '(0.0001)', (26, 37, 26, 41): '(0.01)'}, {}), "('lr', 0.0001, 0.01)", False, 'from hyperopt import hp\n'), ((27, 15, 27, 48), 'hyperopt.hp.uniform', 'hp.uniform', ({(27, 26, 27, 33): '"""beta1"""', (27, 35, 27, 40): '(0.001)', (27, 42, 27, 47): '(0.999)'}, {}), "('beta1', 0.001, 0.999)", False, 'from hyperopt import hp\n'), ((28, 15, 28, 48), 'hyperopt.hp.uniform', 'hp.uniform', ({(28, 26, 28, 33): '"""beta2"""', (28, 35, 28, 40): '(0.001)', (28, 42, 28, 47): '(0.999)'}, {}), "('beta2', 0.001, 0.999)", False, 'from hyperopt import hp\n'), ((29, 22, 29, 58), 'hyperopt.hp.uniform', 'hp.uniform', ({(29, 33, 29, 47): '"""weight_decay"""', (29, 49, 29, 52): '(0.0)', (29, 54, 29, 57): '(0.5)'}, {}), "('weight_decay', 0.0, 0.5)", False, 'from hyperopt import hp\n'), ((30, 15, 30, 44), 'hyperopt.hp.uniform', 'hp.uniform', ({(30, 26, 30, 33): '"""gamma"""', (30, 35, 30, 38): '(0.1)', (30, 40, 30, 43): '(0.9)'}, {}), "('gamma', 0.1, 0.9)", False, 'from hyperopt import hp\n'), ((31, 19, 31, 58), 'hyperopt.hp.choice', 'hp.choice', ({(31, 29, 31, 40): '"""step_size"""', (31, 42, 31, 57): '(5, 10, 20, 40)'}, {}), "('step_size', (5, 10, 20, 40))", False, 'from hyperopt import hp\n'), ((34, 20, 34, 56), 'hyperopt.hp.choice', 'hp.choice', ({(34, 30, 34, 42): '"""batch_size"""', (34, 44, 34, 55): '(8, 16, 32)'}, {}), "('batch_size', (8, 16, 32))", False, 'from hyperopt import hp\n'), ((35, 27, 35, 72), 'hyperopt.hp.choice', 'hp.choice', ({(35, 37, 35, 56): '"""feature_transform"""', (35, 58, 35, 71): '(True, False)'}, {}), "('feature_transform', (True, False))", False, 'from hyperopt import hp\n'), ((63, 7, 63, 61), 'hfta.hfht.generate_nonfusible_param', 'generate_nonfusible_param', ({(63, 33, 63, 39): 'params', (63, 41, 63, 60): '"""feature_transform"""'}, {}), "(params, 'feature_transform')", False, 'from hfta.hfht import tune_hyperparameters, attach_common_args, rearrange_algorithm_kwargs, handle_integers, generate_fusible_param_flags, generate_nonfusible_param\n'), ((123, 18, 123, 52), 'os.path.join', 'os.path.join', ({(123, 31, 123, 42): 'args.outdir', (123, 44, 123, 51): 'ids_str'}, {}), '(args.outdir, ids_str)', False, 'import os\n'), ((232, 32, 232, 63), 'os.path.expanduser', 'os.path.expanduser', ({(232, 51, 232, 62): 'args.outdir'}, {}), '(args.outdir)', False, 'import os\n'), ((233, 33, 233, 65), 'os.path.expanduser', 'os.path.expanduser', ({(233, 52, 233, 64): 'args.dataset'}, {}), '(args.dataset)', False, 'import os\n'), ((66, 8, 69, 9), 'hfta.hfht.generate_fusible_param_flags', 'generate_fusible_param_flags', ({(67, 12, 67, 18): 'params', (68, 12, 68, 74): "['lr', 'beta1', 'beta2', 'weight_decay', 'gamma', 'step_size']"}, {}), "(params, ['lr', 'beta1', 'beta2',\n    'weight_decay', 'gamma', 'step_size'])", False, 'from hfta.hfht import tune_hyperparameters, attach_common_args, rearrange_algorithm_kwargs, handle_integers, generate_fusible_param_flags, generate_nonfusible_param\n'), ((136, 32, 136, 69), 'os.path.join', 'os.path.join', ({(136, 45, 136, 56): 'results_dir', (136, 58, 136, 68): '"""eval.csv"""'}, {}), "(results_dir, 'eval.csv')", False, 'import os\n'), ((156, 15, 156, 33), 'hfta.hfht.utils.fuse_dicts', 'fuse_dicts', ({(156, 26, 156, 32): 'params'}, {}), '(params)', False, 'from hfta.hfht.utils import fuse_dicts\n'), ((231, 28, 231, 55), 'hfta.workflow.extract_logging_level', 'extract_logging_level', ({(231, 50, 231, 54): 'args'}, {}), '(args)', False, 'from hfta.workflow import extract_logging_level\n'), ((59, 12, 59, 59), 'hfta.hfht.generate_nonfusible_param', 'generate_nonfusible_param', ({(59, 38, 59, 44): 'params', (59, 46, 59, 58): '"""batch_size"""'}, {}), "(params, 'batch_size')", False, 'from hfta.hfht import tune_hyperparameters, attach_common_args, rearrange_algorithm_kwargs, handle_integers, generate_fusible_param_flags, generate_nonfusible_param\n'), ((96, 6, 96, 22), 'logging.error', 'logging.error', ({(96, 20, 96, 21): 'e'}, {}), '(e)', False, 'import logging\n'), ((124, 4, 124, 21), 'pathlib.Path', 'Path', ({(124, 9, 124, 20): 'results_dir'}, {}), '(results_dir)', False, 'from pathlib import Path\n'), ((152, 26, 152, 57), 'hyperopt.pyll.stochastic.sample', 'sample', (), '', False, 'from hyperopt.pyll.stochastic import sample\n'), ((82, 14, 82, 53), 'os.path.join', 'os.path.join', ({(82, 27, 82, 38): 'results_dir', (82, 40, 82, 52): '"""stdout.txt"""'}, {}), "(results_dir, 'stdout.txt')", False, 'import os\n'), ((86, 14, 86, 53), 'os.path.join', 'os.path.join', ({(86, 27, 86, 38): 'results_dir', (86, 40, 86, 52): '"""stderr.txt"""'}, {}), "(results_dir, 'stderr.txt')", False, 'import os\n'), ((91, 49, 91, 74), 'os.path.dirname', 'os.path.dirname', ({(91, 65, 91, 73): '__file__'}, {}), '(__file__)', False, 'import os\n')]