import argparse
import cProfile
import pstats
import sys
import os
from typing import Dict

import torch
from torch.autograd import profiler
from torch.utils.collect_env import get_env_info


def redirect_argv(new_argv):
    sys.argv[:] = new_argv[:]


def compiled_with_cuda(sysinfo):
    if sysinfo.cuda_compiled_version:
        return 'compiled w/ CUDA {}'.format(sysinfo.cuda_compiled_version)
    return 'not compiled w/ CUDA'


env_summary = """
--------------------------------------------------------------------------------
  Environment Summary
--------------------------------------------------------------------------------
PyTorch {pytorch_version}{debug_str} {cuda_compiled}
Running with Python {py_version} and {cuda_runtime}

`{pip_version} list` truncated output:
{pip_list_output}
""".strip()


def run_env_analysis():
    print('Running environment analysis...')
    info = get_env_info()

    result: Dict[str, str] = {}

    debug_str = ''
    if info.is_debug_build:
        debug_str = ' DEBUG'

    cuda_avail = ''
    if info.is_cuda_available:
        cuda = info.cuda_runtime_version
        if cuda is not None:
            cuda_avail = 'CUDA ' + cuda
    else:
        cuda = 'CUDA unavailable'

    pip_version = info.pip_version
    pip_list_output = info.pip_packages
    if pip_list_output is None:
        pip_list_output = 'Unable to fetch'

    result = {
        'debug_str': debug_str,
        'pytorch_version': info.torch_version,
        'cuda_compiled': compiled_with_cuda(info),
        'py_version': '{}.{}'.format(sys.version_info[0], sys.version_info[1]),
        'cuda_runtime': cuda_avail,
        'pip_version': pip_version,
        'pip_list_output': pip_list_output,
    }

    return env_summary.format(**result)


def run_cprofile(code, globs, launch_blocking=False):
    print('Running your script with cProfile')
    prof = cProfile.Profile()
    prof.enable()
    exec(code, globs, None)
    prof.disable()
    return prof


cprof_summary = """
--------------------------------------------------------------------------------
  cProfile output
--------------------------------------------------------------------------------
""".strip()


def print_cprofile_summary(prof, sortby='tottime', topk=15):
    print(cprof_summary)
    cprofile_stats = pstats.Stats(prof).sort_stats(sortby)
    cprofile_stats.print_stats(topk)


def run_autograd_prof(code, globs):
    def run_prof(use_cuda=False):
        with profiler.profile(use_cuda=use_cuda) as prof:
            exec(code, globs, None)
        return prof

    print('Running your script with the autograd profiler...')
    result = [run_prof(use_cuda=False)]
    if torch.cuda.is_available():
        result.append(run_prof(use_cuda=True))
    else:
        result.append(None)

    return result


autograd_prof_summary = """
--------------------------------------------------------------------------------
  autograd profiler output ({mode} mode)
--------------------------------------------------------------------------------
        {description}
{cuda_warning}
{output}
""".strip()


def print_autograd_prof_summary(prof, mode, sortby='cpu_time', topk=15):
    valid_sortby = ['cpu_time', 'cuda_time', 'cpu_time_total', 'cuda_time_total', 'count']
    if sortby not in valid_sortby:
        warn = ('WARNING: invalid sorting option for autograd profiler results: {}\n'
                'Expected `cpu_time`, `cpu_time_total`, or `count`. '
                'Defaulting to `cpu_time`.')
        print(warn.format(sortby))
        sortby = 'cpu_time'

    if mode == 'CUDA':
        cuda_warning = ('\n\tBecause the autograd profiler uses the CUDA event API,\n'
                        '\tthe CUDA time column reports approximately max(cuda_time, cpu_time).\n'
                        '\tPlease ignore this output if your code does not use CUDA.\n')
    else:
        cuda_warning = ''

    sorted_events = sorted(prof.function_events,
                           key=lambda x: getattr(x, sortby), reverse=True)
    topk_events = sorted_events[:topk]

    result = {
        'mode': mode,
        'description': 'top {} events sorted by {}'.format(topk, sortby),
        'output': torch.autograd.profiler_util._build_table(topk_events),
        'cuda_warning': cuda_warning
    }

    print(autograd_prof_summary.format(**result))


descript = """
`bottleneck` is a tool that can be used as an initial step for debugging
bottlenecks in your program.

It summarizes runs of your script with the Python profiler and PyTorch\'s
autograd profiler. Because your script will be profiled, please ensure that it
exits in a finite amount of time.

For more complicated uses of the profilers, please see
https://docs.python.org/3/library/profile.html and
https://pytorch.org/docs/master/autograd.html#profiler for more information.
""".strip()


def parse_args():
    parser = argparse.ArgumentParser(description=descript)
    parser.add_argument('scriptfile', type=str,
                        help='Path to the script to be run. '
                        'Usually run with `python path/to/script`.')
    parser.add_argument('args', type=str, nargs=argparse.REMAINDER,
                        help='Command-line arguments to be passed to the script.')
    return parser.parse_args()


def cpu_time_total(autograd_prof):
    return sum([event.cpu_time_total for event in autograd_prof.function_events])


def main():
    args = parse_args()

    # Customizable constants.
    scriptfile = args.scriptfile
    scriptargs = [] if args.args is None else args.args
    scriptargs.insert(0, scriptfile)
    cprofile_sortby = 'tottime'
    cprofile_topk = 15
    autograd_prof_sortby = 'cpu_time_total'
    autograd_prof_topk = 15

    redirect_argv(scriptargs)

    sys.path.insert(0, os.path.dirname(scriptfile))
    with open(scriptfile, 'rb') as stream:
        code = compile(stream.read(), scriptfile, 'exec')
    globs = {
        '__file__': scriptfile,
        '__name__': '__main__',
        '__package__': None,
        '__cached__': None,
    }

    print(descript)

    env_summary = run_env_analysis()

    if torch.cuda.is_available():
        torch.cuda.init()
    cprofile_prof = run_cprofile(code, globs)
    autograd_prof_cpu, autograd_prof_cuda = run_autograd_prof(code, globs)

    print(env_summary)
    print_cprofile_summary(cprofile_prof, cprofile_sortby, cprofile_topk)

    if not torch.cuda.is_available():
        print_autograd_prof_summary(autograd_prof_cpu, 'CPU', autograd_prof_sortby, autograd_prof_topk)
        return

    # Print both the result of the CPU-mode and CUDA-mode autograd profilers
    # if their execution times are very different.
    cuda_prof_exec_time = cpu_time_total(autograd_prof_cuda)
    if len(autograd_prof_cpu.function_events) > 0:
        cpu_prof_exec_time = cpu_time_total(autograd_prof_cpu)
        pct_diff = (cuda_prof_exec_time - cpu_prof_exec_time) / cuda_prof_exec_time
        if abs(pct_diff) > 0.05:
            print_autograd_prof_summary(autograd_prof_cpu, 'CPU', autograd_prof_sortby, autograd_prof_topk)

    print_autograd_prof_summary(autograd_prof_cuda, 'CUDA', autograd_prof_sortby, autograd_prof_topk)

if __name__ == '__main__':
    main()