# coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Launch on GCP's ML Engine."""

import datetime
import os
import shutil
import subprocess as sp
import sys
import tempfile

from googleapiclient import discovery
from oauth2client.client import GoogleCredentials

from tensor2tensor.data_generators import text_encoder
from tensor2tensor.layers import common_hparams
from tensor2tensor.utils import registry
from tensor2tensor.utils import usr_dir as usr_dir_lib
import tensorflow as tf

FLAGS = tf.flags.FLAGS

CONSOLE_URL = "https://console.cloud.google.com/mlengine/jobs/"
RUNTIME_VERSION = "1.9"


class Gcloud(object):
  """gcloud command strings."""
  # Note these can be modified by set_versions
  VM_VERSION = "tf-1-9"
  TPU_VERSION = "1.9"

  @classmethod
  def set_versions(cls, vm, tpu):
    cls.VM_VERSION = vm
    cls.TPU_VERSION = tpu

  @classmethod
  def create_vm(cls):
    create_vm_str = """
    gcloud compute instances create {name} \
      --machine-type=n1-standard-8 \
      --image-family=%s \
      --image-project=ml-images \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    """ % cls.VM_VERSION
    return create_vm_str

  DELETE_VM = "gcloud compute instances delete {name} --quiet"

  @classmethod
  def create_tpu(cls):
    create_tpu_str = """
    gcloud beta compute tpus create \
      {name} \
      --range={tpu_ip}/29 \
      --version=%s
    """ % cls.TPU_VERSION
    return create_tpu_str

  DELETE_TPU = "gcloud beta compute tpus delete {name} --quiet"

  LIST_TPU = "gcloud beta compute tpus list"
  LIST_VM = "gcloud compute instances list"

  SSH_LOCAL_PORT_FORWARD = "-L {local_port}:{host}:{remote_port}"
  SSH_TUNNEL = """
  gcloud compute ssh {name} -- -N
  """

  DEFAULT_PROJECT = "gcloud config get-value project"
  DEFAULT_REGION = "gcloud config get-value compute/region"


def shell_output(cmd_, **kwargs):
  return text_encoder.to_unicode(sp.check_output(format_cmd(cmd_, **kwargs)))


def shell_run(cmd_, **kwargs):
  return sp.check_call(format_cmd(cmd_, **kwargs))


def format_cmd(cmd_, **kwargs):
  return cmd_.format(**kwargs).strip().split()


def default_region():
  return shell_output(Gcloud.DEFAULT_REGION).strip()


def default_project():
  return shell_output(Gcloud.DEFAULT_PROJECT).strip()


def get_setup_file(name, packages=None):
  if not packages:
    packages = []
  return """
from setuptools import find_packages
from setuptools import setup
setup(
    name="{name}",
    version="0.1",
    packages=find_packages(),
    install_requires={pypi_packages}
)
""".format(name=name, pypi_packages=str(list(packages)))


def job_dir():
  # The flag --job-dir is parsed differently before and after switching to absl
  return getattr(FLAGS, "job-dir", "") or getattr(FLAGS, "job_dir", "")


def get_requirements(usr_dir):
  requirements_file = os.path.join(usr_dir, "requirements.txt")
  if not tf.gfile.Exists(requirements_file):
    return []
  with tf.gfile.Open(requirements_file) as f:
    pkg_list = f.readlines()
    return [pkg.strip() for pkg in pkg_list if "tensor2tensor" not in pkg]


def flags_as_args():
  """Convert FLAGS to list of args suitable for passing on cmd line."""
  if hasattr(FLAGS, "flag_values_dict"):
    args_dict = FLAGS.flag_values_dict()
  else:
    args_dict = dict(FLAGS.__dict__["__flags"])
  del args_dict["cloud_mlengine"]
  # Configured later
  del args_dict["t2t_usr_dir"]
  args_dict.pop("h", None)
  args_dict.pop("helpfull", None)
  args_dict.pop("helpshort", None)
  args_dict.pop("help", None)
  args = []
  for name, val in args_dict.items():
    if val is None:
      continue
    if name.startswith("autotune"):
      continue
    args.extend(["--%s=%s" % (name, str(val))])
  return args


def get_default_master_type(num_gpus=1):
  """Returns master_type for trainingInput."""
  gpus_to_master_map = {
      0: "standard",
      1: "standard_p100",
      4: "complex_model_m_p100",
      8: "complex_model_l_gpu",
  }
  if num_gpus not in gpus_to_master_map:
    raise ValueError("Num gpus must be in %s" %
                     str(sorted(list(gpus_to_master_map.keys()))))
  return gpus_to_master_map[num_gpus]


def configure_job():
  """Construct jobSpec for ML Engine job."""
  # See documentation:
  # https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#traininginput
  training_input = {
      "pythonModule": "tensor2tensor.bin.t2t_trainer",
      "args": flags_as_args(),
      "region": text_encoder.native_to_unicode(default_region()),
      "runtimeVersion": RUNTIME_VERSION,