# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import subprocess
import os
import errno

def download_file(url, local_fname=None, force_write=False):
    # requests is not default installed
    import requests
    if local_fname is None:
        local_fname = url.split('/')[-1]
    if not force_write and os.path.exists(local_fname):
        return local_fname

    dir_name = os.path.dirname(local_fname)

    if dir_name != "":
        if not os.path.exists(dir_name):
            try: # try to create the directory if it doesn't exists
                os.makedirs(dir_name)
            except OSError as exc:
                if exc.errno != errno.EEXIST:
                    raise

    r = requests.get(url, stream=True)
    assert r.status_code == 200, "failed to open %s" % url
    with open(local_fname, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk: # filter out keep-alive new chunks
                f.write(chunk)
    return local_fname

def get_gpus():
    """
    return a list of GPUs
    """
    try:
        re = subprocess.check_output(["nvidia-smi", "-L"], universal_newlines=True)
    except OSError:
        return []
    return range(len([i for i in re.split('\n') if 'GPU' in i]))

# <EcoSys> callback functions for nvprof start/stop
class cudaProfilerStart(object):
    """Start nvprof at certain iteration

    Parameters
    ----------
    nbatch: int
        The batch number where nvprof begins.
    nepoch: int
        The epoch number where nvprof begins.
    ----------

    """

    def __init__(self, nbatch, nepoch = 0):
        self.nbatch = nbatch
        self.nepoch = nepoch

    def __call__(self, param):
        import numba.cuda as cuda
        if self.nbatch == param.nbatch and self.nepoch == param.epoch:
            cuda.profile_start()

class cudaProfilerStop(object):
    """Start nvprof at certain iteration

    Parameters
    ----------
    nbatch: int
        The batch number where nvprof ends.
    nepoch: int
        The epoch number where nvprof ends.
    ----------

    """

    def __init__(self, nbatch, nepoch = 0):
        self.nbatch = nbatch
        self.nepoch = nepoch

    def __call__(self, param):
        import numba.cuda as cuda
        if self.nbatch == param.nbatch and self.nepoch == param.epoch:
            cuda.profile_stop()

# </EcoSys>