# Copyright (c) 2018-2022 NVIDIA CORPORATION. All rights reserved.
# Copyright 2019 Google LLC
#
# 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.

"""Utilities for the reinforcement trainer."""

import sys
sys.path.insert(0, '.')  # nopep8

import asyncio
import logging
import os
import tensorflow as tf

import mask_flags
from absl import flags
from utils import *

from mlperf_logging import mllog

# Flags that take multiple values.
# For these flags, expand_cmd_str appends each value in order.
# For all other flags, expand_cmd_str takes the last value.
MULTI_VALUE_FLAGS = set(['--lr_boundaries', '--lr_rates'])


flag_cache = {}
flag_cache_lock = asyncio.Lock()


def is_python_cmd(cmd):
    return cmd[0] == 'python' or cmd[0] == 'python3'


def get_cmd_name(cmd):
    path = cmd[1] if is_python_cmd(cmd) else cmd[0]
    return os.path.splitext(os.path.basename(path))[0]


async def expand_cmd_str(cmd):
    n = 2 if is_python_cmd(cmd) else 1
    cmd = list(cmd)
    args = cmd[n:]
    process = cmd[:n]
    key = ' '.join(process)

    async with flag_cache_lock:
        valid_flags = flag_cache.get(key)
        if valid_flags is None:
            valid_flags = mask_flags.extract_valid_flags(cmd)
            flag_cache[key] = valid_flags

    parsed_args = flags.FlagValues().read_flags_from_files(args)
    flag_args = {}
    position_args = []
    for arg in parsed_args:
        if arg.startswith('--'):
            if '=' not in arg:
                flag_args[arg] = None
            else:
                flag, value = arg.split('=', 1)
                if flag in MULTI_VALUE_FLAGS:
                    if flag not in flag_args:
                        flag_args[flag] = []
                    flag_args[flag].append(value)
                else:
                    flag_args[flag] = value
        else:
            position_args.append(arg)

    flag_list = []
    for flag, value in flag_args.items():
        if value is None:
            flag_list.append(flag)
        elif type(value) == list:
            for v in value:
                flag_list.append('%s=%s' % (flag, v))
        else:
            flag_list.append('%s=%s' % (flag, value))

    flag_list = sorted(mask_flags.filter_flags(flag_list, valid_flags))
    return '  '.join(process + flag_list + position_args)


def list_selfplay_dirs(base_dir):
    """Returns a sorted list of selfplay data directories.

    Training examples are written out to the following directory hierarchy:
      base_dir/model_name/device_id/timestamp/

    Args:
      base_dir: either selfplay_dir or holdout_dir.

    Returns:
      A list of model directories sorted so the most recent directory is first.
    """

    model_dirs = [os.path.join(base_dir, x)
                  for x in tf.io.gfile.listdir(base_dir)]
    return sorted(model_dirs, reverse=True)

def ensure_file_exists(src, verbose=False):
    """Touch file if it doesn't exist."""

    if not tf.io.gfile.exists(src):
       if(verbose): print('Creating file {}'.format(src))
       with tf.io.gfile.GFile(src, 'wb') as f:
           f.write('Abort\n')


def ensure_dir_exists(src, verbose=False):
    """Make directory src if it doesn't already exists."""

    if not tf.io.gfile.exists(src):
       if(verbose): print('Creating {}'.format(src))
       tf.io.gfile.makedirs(src)

def remove_tree(src, verbose=False):
    """Removes everything under src."""

    if tf.io.gfile.exists(src):
       if verbose: print('Deleting {}'.format(src))
       tf.io.gfile.rmtree(src)

def copy_file(src, dst, verbose=False):
    """Copies src file to dst."""

    if verbose: print('Copying {} to {}'.format(src, dst))
    tf.io.gfile.copy(src, dst)

def copy_dir_tree(src, dst, verbose=False):
    """Copies directory structure under src to dst."""

    if verbose:
        print('Copying dir-structure {} to {}'.format(src, dst))
    for src_dir, sub_dirs, basenames in tf.io.gfile.walk(src):
        rel_dir = os.path.relpath(src_dir, src)
        dst_dir = os.path.join(dst, rel_dir)
        for sub_dir in sorted(sub_dirs):
            path = os.path.join(dst, rel_dir, sub_dir)
            tf.io.gfile.makedirs(path)

def copy_tree(src, dst, verbose=False):
    """Copies everything under src to dst."""

    if verbose:
        print('Copying {} to {}'.format(src, dst))
    for src_dir, sub_dirs, basenames in tf.io.gfile.walk(src):
        rel_dir = os.path.relpath(src_dir, src)
        dst_dir = os.path.join(dst, rel_dir)
        for sub_dir in sorted(sub_dirs):
            path = os.path.join(dst, rel_dir, sub_dir)
            if verbose:
                print('Make dir {}'.format(path))
            tf.io.gfile.makedirs(path)
        if basenames:
            if verbose:
                print('Copying {} files from {} to {}'.format(
                    len(basenames), src_dir, dst_dir))
            for basename in basenames:
                src_path = os.path.join(src_dir, basename)
                dst_path = os.path.join(dst_dir, basename)
                if verbose:
                    print('Copying {} to {}'.format(src_path, dst_path))
                tf.io.gfile.copy(src_path, dst_path)


async def checked_run(cmd, env=None, process_will_abort=False):
    """Run the given subprocess command in a coroutine.

    Args:
        cmd: the command to run and its arguments.
        env: optional OS environment to run the process under.

    Returns:
        The output that the command wrote to stdout & stderr.

    Raises:
        RuntimeError: if the command returns a non-zero result.
    """

    # Start the subprocess.
    print('Running command {}'.format(cmd))
    logging.info('Running: %s', await expand_cmd_str(cmd))
    with logged_timer('{} finished'.format(get_cmd_name(cmd))):
        p = await asyncio.create_subprocess_exec(
            *cmd, env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT)

        # Stream output from the process stdout.
        lines = []
        while True:
            line = await p.stdout.readline()
            if not line:
                break
            line = line.decode()[:-1]
            lines.append(line)
            logging.info(line)

        # Wait for the process to finish, check it was successful & build stdout.
        await p.wait()
        output = '\n'.join(lines)[:-1]
        if p.returncode and not process_will_abort:
            raise RuntimeError('Return code {} from process: {}\n{}'.format(
                p.returncode, await expand_cmd_str(cmd), output))

        return output


def wait(aws):
    """Waits for all of the awaitable objects (e.g. coroutines) in aws to finish.

    All the awaitable objects are waited for, even if one of them raises an
    exception. When one or more awaitable raises an exception, the exception
    from the awaitable with the lowest index in the aws list will be reraised.

    Args:
        aws: a single awaitable, or list awaitables.

    Returns:
        If aws is a single awaitable, its result.
        If aws is a list of awaitables, a list containing the of each awaitable
        in the list.

    Raises:
        Exception: if any of the awaitables raises.
    """

    aws_list = aws if isinstance(aws, list) else [aws]
    results = asyncio.get_event_loop().run_until_complete(asyncio.gather(
        *aws_list, return_exceptions=True))
    # If any of the cmds failed, re-raise the error.
    for result in results:
        if isinstance(result, Exception):
            raise result
    return results if isinstance(aws, list) else results[0]
