"""
To avoid any issues of memory hanging around between inputs,
we run each input as a separate process.

A little ugly but effective
"""
import gc
import glob
import json
import os
import random
import time

import numpy as np
import torch

from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists

from validate_config import validate_trials_config
from pt_trial_util import create_csv_writer
from tqdm import tqdm
import model_util


def extend_simrd_config(dest_dir, sim_conf_filename, model_name, specific_params, log_name):
    if not check_file_exists(dest_dir, sim_conf_filename):
        prepare_out_file(dest_dir, sim_conf_filename)
        write_json(dest_dir, sim_conf_filename, dict())

    conf = read_json(dest_dir, sim_conf_filename)
    if model_name not in conf:
        conf[model_name] = []
    conf[model_name].append({
            'name': model_util.get_model_family(model_name),
            'batch_size': str(specific_params['batch_size']),
            'layers': specific_params.get('layers', model_util.get_model_layers(model_name)),
            'type': model_util.get_model_type(model_name),
            'log': log_name,
            'has_start': True
    })
    write_json(dest_dir, sim_conf_filename, conf)


def save_trial_log(dest_dir, sim_conf_filename, model_name, specific_params, is_baseline=False):
    """
    Find the last DTR log produced in the trial (if any exist)
    and move it to the directory
    """
    all_logs = glob.glob(os.path.join(os.getcwd(), '*.log'))
    if not all_logs:
        return

    # if we delete all logs in advance, there should be at most one log
    assert len(all_logs) == 1
    most_recent = all_logs[0]

    # rename and move
    # (new name just appends info to the old one)
    batch_size = specific_params['batch_size']
    budget = specific_params['memory_budget']
    if budget < 0:
        budget = 'inf'
    new_name = '{}-{}-{}-{}'.format(model_name, batch_size, budget,
                                    os.path.basename(most_recent))
    filename = prepare_out_file(dest_dir, new_name)
    os.rename(most_recent, filename)
    if is_baseline and sim_conf_filename is not None:
        extend_simrd_config(dest_dir, sim_conf_filename, model_name, specific_params, filename)


def delete_logs():
    for log in glob.glob(os.path.join(os.getcwd(), '*.log')):
        os.remove(log)


def run_single_measurement(model_name, produce_model, run_model, teardown, inp, criterion, extra_params, use_dtr, use_profiling):
    """
    This function initializes a model and performs
    a single measurement of the model on the given input.

    While it might seem most reasonable to initialize
    the model outside of the loop, DTR's logs have shown
    that certain constants in the model persist between loop iterations;
    performing these actions in a separate *function scope* turned out to be the only
    way to prevent having those constants hang around.

    Returns a dict of measurements
    """
    torch.cuda.reset_max_memory_allocated()
    # resetting means the count should be reset to
    # only what's in scope, meaning only the input
    input_mem = torch.cuda.max_memory_allocated()
    model = produce_model(extra_params=extra_params)
    params = []
    for m in model:
        if hasattr(m, 'parameters'):
            params.extend(m.parameters())

    model_mem = torch.cuda.max_memory_allocated()

    optimizer = torch.optim.SGD(model[0].parameters(), 1e-3, momentum=0.9, weight_decay=1e-4)

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    # start timing
    torch.cuda.synchronize()
    start_time = time.time()
    if use_dtr:
        torch.reset_profile()
    start.record()
    # with torch.autograd.profiler.profile(use_cuda=True) as prof:
    run_model(criterion, *model, *inp, optimizer=optimizer)
    end.record()
    start_sync = time.time()
    torch.cuda.synchronize()
    end_sync = time.time()
    end_time = time.time()
    # end timing

    if use_dtr:
        # operators-only time, tracked by DTR
        cuda_time = torch.compute_time()

    base_compute_time = -1
    remat_compute_time = -1
    search_time = -1
    cost_time = -1
    if use_profiling:
        base_compute_time = torch.base_compute_time()
        remat_compute_time = torch.remat_compute_time()
        search_time = torch.search_time()
        cost_time = torch.cost_time()
        torch.reset_profile()

    total_mem = torch.cuda.max_memory_allocated()
    teardown(*model)
    torch.cuda.reset_max_memory_allocated()

    del model

    if use_dtr:
        torch.toggle_log(False)

    del params

    batch_size = len(inp[0])
    ips = batch_size / (end_time - start_time)

    result = {
        'time': end_time - start_time,
        'sync_time': end_sync - start_sync,
        'gpu_time': start.elapsed_time(end),
        'input_mem': input_mem,
        'model_mem': model_mem,
        'total_mem': total_mem,
        'base_compute_time': base_compute_time,
        'remat_compute_time': remat_compute_time,
        'search_time': search_time,
        'cost_time': cost_time,
        'batch_size': batch_size,
        'ips': ips
    }
    if use_dtr:
        result['cuda_time'] = cuda_time
    else:
        result['cuda_time'] = -1.0


    return result


def timing_loop(model_name, i, config, use_dtr,
                specific_params, writer, trial_run=False, trial_run_outfile=None, memory_budget=-1.0):
    dry_run = config['dry_run']
    measurements = []
    print(f'Running {model_name} : {specific_params}')

    # remove any logs hanging around (so we only have to look for one)
    delete_logs()

    # we only save logs for the final input on DTR
    save_log = use_dtr and specific_params.get('save_logs', config['save_logs']) and i == config['n_inputs'] - 1
    if use_dtr:
        torch.toggle_log(False)

    # whether to report profiling info
    use_profiling = use_dtr and specific_params.get('use_profiling', False)

    use_cudnn = model_util.use_cudnn(model_name)
    with torch.backends.cudnn.flags(enabled=use_cudnn, benchmark=use_cudnn):
        criterion = model_util.get_criterion(model_name)
        produce_model, gen_input, run_model, teardown = model_util.prepare_model(model_name,
                                                                                 specific_params['batch_size'],
                                                                                 use_dtr=use_dtr)
        inp = gen_input(i, specific_params.get('extra_params', dict()))

        n_reps = specific_params.get('n_reps', config['n_reps'])

        if use_profiling:
            torch.toggle_profile(use_profiling)

        progress = tqdm(range(dry_run + n_reps))
        for j in progress:
            progress.set_description(f'Rep [{j}]' + '' if j > dry_run else f'Dry run [{j}]')
            gc.collect()
            # Annotate where the final run starts in the log
            if save_log and j == dry_run + n_reps - 1:
                torch.toggle_log(True)
                torch.annotate_log('START')

            res = run_single_measurement(model_name, produce_model, run_model,
                                         teardown, inp, criterion, extra_params=specific_params.get('extra_params', dict()), use_dtr=use_dtr, use_profiling=use_profiling)
            if j >= dry_run:
                measurements.append(res)

        # Dump results
        model_name_replace_dict = {
            'tv_resnet152': 'resnet152',
            'tv_resnet50': 'resnet50',
        }

        train_ips_list = []
        batch_size = None
        for res in measurements:
            batch_size = res['batch_size']
            train_ips_list.append(res['ips'])

        out_file = "speed_results.tsv"
        with open(out_file, "a") as fout:
            val_dict = {
                'network': model_name_replace_dict.get(model_name, model_name),
                'algorithm': 'dtr',
                'budget': specific_params['memory_budget'],
                'batch_size': batch_size,
                'ips': np.median(train_ips_list) if train_ips_list else -1,
            }
            print(val_dict)
            fout.write(json.dumps(val_dict) + "\n")
        print(f"save results to {out_file}")



    # write to csv file only when this trial is not
    # for getting a baseline memory usage
    if trial_run:
        write_json(os.getcwd(), trial_run_outfile, {
            'mem' : max(map(lambda data: data['total_mem'], measurements))
        })
        return

    if save_log:
        save_trial_log(config['log_dest'], config.get('simrd_config', None),
                       model_name,
                       specific_params,
                       is_baseline=specific_params['memory_budget'] == -1)

    # clean up after ourselves
    delete_logs()

    # do all the writing after the trial is over
    for j in range(len(measurements)):
        data = measurements[j]
        # do unit conversions now: times in ms,
        # memory in MB
        writer.writerow({
            'time': data['time']*1e3,
            'sync_time': data['sync_time']*1e3,
            # pytorch's cuda elapsed time is already in ms
            'gpu_time': float(data['gpu_time']),
            # 'cuda_time' : float(data['cuda_time']) * 1e-6,
            'input_mem': data['input_mem']*1e-6,
            'model_mem': data['model_mem']*1e-6,
            'total_mem': data['total_mem']*1e-6,
            'memory_budget': memory_budget,
            # profiling (reported in nanoseconds)
            'base_compute_time': data['base_compute_time']*1e-6,
            'remat_compute_time': data['remat_compute_time']*1e-6,
            'search_time': data['search_time']*1e-6,
            'cost_time': data['cost_time']*1e-6,
            'rep': j - dry_run,
            'input': i,
            **specific_params
         })


def main(config_dir, experiment_mode, model_name, input_idx, params_file, out_file,
         trial_run=False, trial_run_outfile=None):
    if 'DTR_MODEL_NAME' in os.environ:
        model_name = os.environ['DTR_MODEL_NAME']
    config, msg = validate_trials_config(config_dir)
    if config is None:
        print(msg)
        return 1

    use_dtr = (experiment_mode == 'dtr')

    i = int(input_idx)
    is_trial = trial_run == 'True'

    if config['set_seed']:
        torch.manual_seed(config['seed'] + i)
        random.seed(config['seed'] + i)

    cwd = os.getcwd()

    # handle specific params, esp. for DTR
    specific_params = read_json(cwd, params_file)
    if 'DTR_MEMORY_BUDGET' in os.environ:
        specific_params['memory_budget'] = float(os.environ['DTR_MEMORY_BUDGET'])

    assert 'batch_size' in specific_params
    if use_dtr:
        assert 'memory_budget' in specific_params
        if specific_params['memory_budget'] > 0:
            print(f'Setting budget to {int(specific_params["memory_budget"])}')
            torch.set_memory_budget(int(specific_params['memory_budget']))
    if is_trial:
        timing_loop(model_name, i, config, use_dtr, specific_params, None, True, trial_run_outfile)
        return

    with open(out_file, 'a', newline='') as csvfile:
        writer = create_csv_writer(csvfile, specific_params)
        timing_loop(model_name, i, config, use_dtr, specific_params, writer, memory_budget=specific_params.get('memory_budget', -1))


if __name__ == '__main__':
    invoke_main(main, 'config_dir', 'experiment_mode',
                'model_name', 'input_idx', 'params_file',
                'out_file', 'trial_run', 'trial_run_outfile')

[((30, 11, 30, 49), 'common.read_json', 'read_json', ({(30, 21, 30, 29): 'dest_dir', (30, 31, 30, 48): 'sim_conf_filename'}, {}), '(dest_dir, sim_conf_filename)', False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((41, 4, 41, 49), 'common.write_json', 'write_json', ({(41, 15, 41, 23): 'dest_dir', (41, 25, 41, 42): 'sim_conf_filename', (41, 44, 41, 48): 'conf'}, {}), '(dest_dir, sim_conf_filename, conf)', False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((65, 15, 65, 51), 'common.prepare_out_file', 'prepare_out_file', ({(65, 32, 65, 40): 'dest_dir', (65, 42, 65, 50): 'new_name'}, {}), '(dest_dir, new_name)', False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((66, 4, 66, 36), 'os.rename', 'os.rename', ({(66, 14, 66, 25): 'most_recent', (66, 27, 66, 35): 'filename'}, {}), '(most_recent, filename)', False, 'import os\n'), ((89, 4, 89, 43), 'torch.cuda.reset_max_memory_allocated', 'torch.cuda.reset_max_memory_allocated', ({}, {}), '()', False, 'import torch\n'), ((92, 16, 92, 49), 'torch.cuda.max_memory_allocated', 'torch.cuda.max_memory_allocated', ({}, {}), '()', False, 'import torch\n'), ((99, 16, 99, 49), 'torch.cuda.max_memory_allocated', 'torch.cuda.max_memory_allocated', ({}, {}), '()', False, 'import torch\n'), ((103, 12, 103, 48), 'torch.cuda.Event', 'torch.cuda.Event', (), '', False, 'import torch\n'), ((104, 10, 104, 46), 'torch.cuda.Event', 'torch.cuda.Event', (), '', False, 'import torch\n'), ((107, 4, 107, 28), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((108, 17, 108, 28), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((115, 17, 115, 28), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((116, 4, 116, 28), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((117, 15, 117, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((118, 15, 118, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((136, 16, 136, 49), 'torch.cuda.max_memory_allocated', 'torch.cuda.max_memory_allocated', ({}, {}), '()', False, 'import torch\n'), ((138, 4, 138, 43), 'torch.cuda.reset_max_memory_allocated', 'torch.cuda.reset_max_memory_allocated', ({}, {}), '()', False, 'import torch\n'), ((190, 16, 190, 48), 'model_util.use_cudnn', 'model_util.use_cudnn', ({(190, 37, 190, 47): 'model_name'}, {}), '(model_name)', False, 'import model_util\n'), ((291, 18, 291, 52), 'validate_config.validate_trials_config', 'validate_trials_config', ({(291, 41, 291, 51): 'config_dir'}, {}), '(config_dir)', False, 'from validate_config import validate_trials_config\n'), ((305, 10, 305, 21), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((308, 22, 308, 49), 'common.read_json', 'read_json', ({(308, 32, 308, 35): 'cwd', (308, 37, 308, 48): 'params_file'}, {}), '(cwd, params_file)', False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((328, 4, 330, 61), 'common.invoke_main', 'invoke_main', ({(328, 16, 328, 20): 'main', (328, 22, 328, 34): '"""config_dir"""', (328, 36, 328, 53): '"""experiment_mode"""', (329, 16, 329, 28): '"""model_name"""', (329, 30, 329, 41): '"""input_idx"""', (329, 43, 329, 56): '"""params_file"""', (330, 16, 330, 26): '"""out_file"""', (330, 28, 330, 39): '"""trial_run"""', (330, 41, 330, 60): '"""trial_run_outfile"""'}, {}), "(main, 'config_dir', 'experiment_mode', 'model_name',\n    'input_idx', 'params_file', 'out_file', 'trial_run', 'trial_run_outfile')", False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((26, 11, 26, 57), 'common.check_file_exists', 'check_file_exists', ({(26, 29, 26, 37): 'dest_dir', (26, 39, 26, 56): 'sim_conf_filename'}, {}), '(dest_dir, sim_conf_filename)', False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((27, 8, 27, 53), 'common.prepare_out_file', 'prepare_out_file', ({(27, 25, 27, 33): 'dest_dir', (27, 35, 27, 52): 'sim_conf_filename'}, {}), '(dest_dir, sim_conf_filename)', False, 'from common import invoke_main, read_json, write_json, prepare_out_file, check_file_exists\n'), ((64, 36, 64, 65), 'os.path.basename', 'os.path.basename', ({(64, 53, 64, 64): 'most_recent'}, {}), '(most_recent)', False, 'import os\n'), ((73, 8, 73, 22), 'os.remove', 'os.remove', ({(73, 18, 73, 21): 'log'}, {}), '(log)', False, 'import os\n'), ((110, 8, 110, 29), 'torch.reset_profile', 'torch.reset_profile', ({}, {}), '()', False, 'import torch\n'), ((123, 20, 123, 40), 'torch.compute_time', 'torch.compute_time', ({}, {}), '()', False, 'import torch\n'), ((130, 28, 130, 53), 'torch.base_compute_time', 'torch.base_compute_time', ({}, {}), '()', False, 'import torch\n'), ((131, 29, 131, 55), 'torch.remat_compute_time', 'torch.remat_compute_time', ({}, {}), '()', False, 'import torch\n'), ((132, 22, 132, 41), 'torch.search_time', 'torch.search_time', ({}, {}), '()', False, 'import torch\n'), ((133, 20, 133, 37), 'torch.cost_time', 'torch.cost_time', ({}, {}), '()', False, 'import torch\n'), ((134, 8, 134, 29), 'torch.reset_profile', 'torch.reset_profile', ({}, {}), '()', False, 'import torch\n'), ((143, 8, 143, 31), 'torch.toggle_log', 'torch.toggle_log', ({(143, 25, 143, 30): '(False)'}, {}), '(False)', False, 'import torch\n'), ((185, 8, 185, 31), 'torch.toggle_log', 'torch.toggle_log', ({(185, 25, 185, 30): '(False)'}, {}), '(False)', False, 'import torch\n'), ((191, 9, 191, 75), 'torch.backends.cudnn.flags', 'torch.backends.cudnn.flags', (), '', False, 'import torch\n'), ((192, 20, 192, 56), 'model_util.get_criterion', 'model_util.get_criterion', ({(192, 45, 192, 55): 'model_name'}, {}), '(model_name)', False, 'import model_util\n'), ((193, 56, 195, 97), 'model_util.prepare_model', 'model_util.prepare_model', (), '', False, 'import model_util\n'), ((302, 8, 302, 45), 'torch.manual_seed', 'torch.manual_seed', ({(302, 26, 302, 44): "(config['seed'] + i)"}, {}), "(config['seed'] + i)", False, 'import torch\n'), ((303, 8, 303, 39), 'random.seed', 'random.seed', ({(303, 20, 303, 38): "(config['seed'] + i)"}, {}), "(config['seed'] + i)", False, 'import random\n'), ((323, 17, 323, 60), 'pt_trial_util.create_csv_writer', 'create_csv_writer', ({(323, 35, 323, 42): 'csvfile', (323, 44, 323, 59): 'specific_params'}, {}), '(csvfile, specific_params)', False, 'from pt_trial_util import create_csv_writer\n'), ((34, 20, 34, 59), 'model_util.get_model_family', 'model_util.get_model_family', ({(34, 48, 34, 58): 'model_name'}, {}), '(model_name)', False, 'import model_util\n'), ((37, 20, 37, 57), 'model_util.get_model_type', 'model_util.get_model_type', ({(37, 46, 37, 56): 'model_name'}, {}), '(model_name)', False, 'import model_util\n'), ((49, 38, 49, 49), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((72, 38, 72, 49), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((201, 12, 201, 47), 'torch.toggle_profile', 'torch.toggle_profile', ({(201, 33, 201, 46): 'use_profiling'}, {}), '(use_profiling)', False, 'import torch\n'), ((206, 12, 206, 24), 'gc.collect', 'gc.collect', ({}, {}), '()', False, 'import gc\n'), ((247, 19, 247, 30), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((36, 52, 36, 91), 'model_util.get_model_layers', 'model_util.get_model_layers', ({(36, 80, 36, 90): 'model_name'}, {}), '(model_name)', False, 'import model_util\n'), ((209, 16, 209, 38), 'torch.toggle_log', 'torch.toggle_log', ({(209, 33, 209, 37): '(True)'}, {}), '(True)', False, 'import torch\n'), ((210, 16, 210, 43), 'torch.annotate_log', 'torch.annotate_log', ({(210, 35, 210, 42): '"""START"""'}, {}), "('START')", False, 'import torch\n'), ((236, 23, 236, 48), 'numpy.median', 'np.median', ({(236, 33, 236, 47): 'train_ips_list'}, {}), '(train_ips_list)', True, 'import numpy as np\n'), ((239, 23, 239, 43), 'json.dumps', 'json.dumps', ({(239, 34, 239, 42): 'val_dict'}, {}), '(val_dict)', False, 'import json\n')]