import os
import pickle
import time
import timeit

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np
import torch
import tempfile
import horovod.torch as hvd
from horovod.ray import RayExecutor

from ray_shuffling_data_loader.torch_dataset import (TorchShufflingDataset)
from ray_shuffling_data_loader.data_generation import (generate_data,
                                                       DATA_SPEC)

import argparse

DEFAULT_DATA_DIR = "s3://shuffling-data-loader-benchmarks/data/"

numpy_to_torch_dtype = {
    np.bool: torch.bool,
    np.uint8: torch.uint8,
    np.int8: torch.int8,
    np.int16: torch.int16,
    np.int32: torch.int32,
    np.int64: torch.int64,
    np.float16: torch.float16,
    np.float32: torch.float32,
    np.float64: torch.float64,
    np.complex64: torch.complex64,
    np.complex128: torch.complex128
}

# Training settings
parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
parser.add_argument(
    "--batch-size",
    type=int,
    default=250000,
    metavar="N",
    help="input batch size for training (default: 64)")
parser.add_argument(
    "--test-batch-size",
    type=int,
    default=250000,
    metavar="N",
    help="input batch size for testing (default: 1000)")
parser.add_argument(
    "--epochs",
    type=int,
    default=10,
    metavar="N",
    help="number of epochs to train (default: 10)")
parser.add_argument(
    "--lr",
    type=float,
    default=0.01,
    metavar="LR",
    help="learning rate (default: 0.01)")
parser.add_argument(
    "--momentum",
    type=float,
    default=0.5,
    metavar="M",
    help="SGD momentum (default: 0.5)")
parser.add_argument(
    "--no-cuda",
    action="store_true",
    default=False,
    help="disables CUDA training")
parser.add_argument(
    "--seed",
    type=int,
    default=42,
    metavar="S",
    help="random seed (default: 42)")
parser.add_argument(
    "--log-interval",
    type=int,
    default=10,
    metavar="N",
    help=("how many batches to wait before logging training "
          "status"))
parser.add_argument(
    "--fp16-allreduce",
    action="store_true",
    default=False,
    help="use fp16 compression during allreduce")
parser.add_argument(
    "--use-adasum",
    action="store_true",
    default=False,
    help="use adasum algorithm to do reduction")
parser.add_argument(
    "--gradient-predivide-factor",
    type=float,
    default=1.0,
    help=("apply gradient predivide factor in optimizer "
          "(default: 1.0)"))
parser.add_argument("--num-workers", type=int, default=None)
parser.add_argument("--num-hosts", type=int, default=None)
parser.add_argument("--num-workers-per-host", type=int, default=None)
parser.add_argument("--cpus-per-worker", type=int, default=1)
parser.add_argument("--mock-train-step-time", type=float, default=1.0)

# Synthetic training data generation settings.
parser.add_argument("--cache-files", action="store_true", default=False)
parser.add_argument("--num-rows", type=int, default=2 * (10**7))
parser.add_argument("--num-files", type=int, default=25)
parser.add_argument("--max-row-group-skew", type=float, default=0.0)
parser.add_argument("--num-row-groups-per-file", type=int, default=5)
parser.add_argument("--data-dir", type=str, default=DEFAULT_DATA_DIR)

# Shuffling data loader settings.
parser.add_argument("--num-reducers", type=int, default=32)
parser.add_argument("--max-concurrent-epochs", type=int, default=2)
parser.add_argument("--address", default="auto")


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)


def train_main(args, filenames):
    # Horovod: initialize library.
    hvd.init()
    torch.manual_seed(args.seed)

    if torch.cuda.is_available() and not args.no_cuda:
        # Horovod: pin GPU to local rank.
        torch.cuda.set_device(hvd.local_rank())
        torch.cuda.manual_seed(args.seed)

    # Horovod: limit # of CPU threads to be used per worker.
    torch.set_num_threads(1)
    rank = hvd.rank()
    train_dataset = create_dataset(
        filenames,
        batch_size=args.batch_size,
        rank=rank,
        num_epochs=args.epochs,
        world_size=hvd.size(),
        num_reducers=args.num_reducers,
        max_concurrent_epochs=args.max_concurrent_epochs)
    model = Net()
    # By default, Adasum doesn"t need scaling up learning rate.
    lr_scaler = hvd.size() if not args.use_adasum else 1

    if torch.cuda.is_available() and not args.no_cuda:
        # Move model to GPU.
        model.cuda()
        # If using GPU Adasum allreduce, scale learning rate by local_size.
        if args.use_adasum and hvd.nccl_built():
            lr_scaler = hvd.local_size()

    # Horovod: scale learning rate by lr_scaler.
    optimizer = optim.SGD(
        model.parameters(), lr=args.lr * lr_scaler, momentum=args.momentum)

    # Horovod: broadcast parameters & optimizer state.
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    hvd.broadcast_optimizer_state(optimizer, root_rank=0)

    # Horovod: (optional) compression algorithm.
    compression = (hvd.Compression.fp16
                   if args.fp16_allreduce else hvd.Compression.none)

    # Horovod: wrap optimizer with DistributedOptimizer.
    optimizer = hvd.DistributedOptimizer(
        optimizer,
        named_parameters=model.named_parameters(),
        compression=compression,
        op=hvd.Adasum if args.use_adasum else hvd.Average,
        gradient_predivide_factor=args.gradient_predivide_factor)

    def _train(epoch):
        model.train()
        # Horovod: set epoch to sampler for shuffling.
        train_dataset.set_epoch(epoch)
        start_epoch = timeit.default_timer()
        last_batch_time = start_epoch
        batch_wait_times = []
        for batch_idx, (data, target) in enumerate(train_dataset):
            batch_wait_times.append(timeit.default_timer() - last_batch_time)
            if torch.cuda.is_available() and not args.no_cuda:
                if isinstance(data, list):
                    data = [t.cuda() for t in data]
                target = target.cuda()
            optimizer.zero_grad()
            # output = model(data)
            if batch_idx % args.log_interval == 0:
                print(
                    f"Processing batch {batch_idx} in epoch {epoch} on worker "
                    f"{rank}.")
            time.sleep(args.mock_train_step_time)
            # TODO(Clark): Add worker synchronization barrier here.
            # loss = F.nll_loss(output, target)
            # loss.backward()
            # optimizer.step()
            last_batch_time = timeit.default_timer()
        epoch_duration = timeit.default_timer() - start_epoch
        avg_batch_wait_time = np.mean(batch_wait_times)
        std_batch_wait_time = np.std(batch_wait_times)
        max_batch_wait_time = np.max(batch_wait_times)
        min_batch_wait_time = np.min(batch_wait_times)
        print(f"\nEpoch {epoch}, worker {rank} stats over "
              f"{len(batch_wait_times)} steps: {epoch_duration:.3f}")
        print(f"Mean batch wait time: {avg_batch_wait_time:.3f}s +- "
              f"{std_batch_wait_time}")
        print(f"Max batch wait time: {max_batch_wait_time:.3f}s")
        print(f"Min batch wait time: {min_batch_wait_time:.3f}s")
        return batch_wait_times

    print(f"Starting training on worker {rank}.")
    batch_wait_times = []
    for epoch in range(args.epochs):
        batch_wait_times.extend(_train(epoch))
    batch_wait_times.pop(0)
    print(f"Done training on worker {rank}.")
    avg_batch_wait_time = np.mean(batch_wait_times)
    std_batch_wait_time = np.std(batch_wait_times)
    max_batch_wait_time = np.max(batch_wait_times)
    min_batch_wait_time = np.min(batch_wait_times)
    print(f"\nWorker {rank} training stats over {args.epochs} epochs:")
    print(f"Mean batch wait time: {avg_batch_wait_time:.3f}s +- "
          f"{std_batch_wait_time}")
    print(f"Max batch wait time: {max_batch_wait_time:.3f}s")
    print(f"Min batch wait time: {min_batch_wait_time:.3f}s")
    # TODO(Clark): Add logic to the dataset abstraction so we don't have to do
    # this.
    if rank == 0:
        print("Waiting in rank 0 worker to let other workers consume queue...")
        time.sleep(10)
        print("Done waiting in rank 0 worker.")


def create_dataset(filenames, *, batch_size, rank, num_epochs, world_size,
                   num_reducers, max_concurrent_epochs):
    print(f"Creating Torch shuffling dataset for worker {rank} with "
          f"{batch_size} batch size, {num_epochs} epochs, {num_reducers} "
          f"reducers, and {world_size} trainers.")
    feature_columns = list(DATA_SPEC.keys())
    feature_types = [
        numpy_to_torch_dtype[dtype] for _, _, dtype in DATA_SPEC.values()
    ]
    label_column = feature_columns.pop()
    label_type = feature_types.pop()
    return TorchShufflingDataset(
        filenames,
        num_epochs,
        world_size,
        batch_size,
        rank,
        num_reducers=num_reducers,
        max_concurrent_epochs=max_concurrent_epochs,
        feature_columns=feature_columns,
        feature_types=feature_types,
        label_column=label_column,
        label_type=label_type)


if __name__ == "__main__":
    args = parser.parse_args()
    from ray_shuffling_data_loader.stats import human_readable_size
    import ray
    print("Connecting to Ray cluster...")
    ray.init(address=args.address)

    num_rows = args.num_rows
    num_files = args.num_files
    num_row_groups_per_file = args.num_row_groups_per_file
    max_row_group_skew = args.max_row_group_skew
    data_dir = args.data_dir

    cache_path = os.path.join(tempfile.gettempdir(), "data_cache")
    filenames = None
    if args.cache_files and os.path.exists(cache_path):
        try:
            with open(cache_path, "rb") as f:
                filenames, num_bytes = pickle.load(f)
        except Exception as exc:
            print(f"Cache load failed - {exc}")

    if not filenames:

        print(f"Generating {num_rows} rows over {num_files} files, with "
              f"{num_row_groups_per_file} row groups per file and at most "
              f"{100 * max_row_group_skew:.1f}% row group skew.")
        filenames, num_bytes = generate_data(num_rows, num_files,
                                             num_row_groups_per_file,
                                             max_row_group_skew, data_dir)
        if args.cache_files:
            with open(os.path.join(tempfile.gettempdir(), "data_cache"),
                      "wb") as f:
                pickle.dump((filenames, num_bytes), f)
    print(f"Generated {len(filenames)} files containing {num_rows} rows "
          f"with {num_row_groups_per_file} row groups per file, totalling "
          f"{human_readable_size(num_bytes)}.")

    print("Create Ray executor")
    worker_kwargs = {}
    num_workers = args.num_workers
    num_hosts = args.num_hosts
    num_workers_per_host = args.num_workers_per_host
    if num_workers is not None:
        if num_hosts is not None:
            raise ValueError(
                "Only one of --num-workers and --num-hosts should be used.")
        worker_kwargs["num_workers"] = num_workers
    elif num_hosts is not None:
        worker_kwargs["num_hosts"] = num_hosts
        if num_workers_per_host is None:
            raise ValueError("When giving --num-hosts, --num-workers-per-host "
                             "must also be given.")
        worker_kwargs["num_workers_per_host"] = num_workers_per_host
    cpus_per_worker = args.cpus_per_worker
    settings = RayExecutor.create_settings(timeout_s=30)
    executor = RayExecutor(
        settings,
        use_gpu=True,
        gpus_per_worker=1,
        cpus_per_worker=cpus_per_worker,
        **worker_kwargs)
    executor.start()
    executor.run(train_main, args=[args, filenames])
    executor.shutdown()

    print("Done consuming batches.")

[(39, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (145, 'horovod.torch.init', 'hvd.init', True, 'import horovod.torch as hvd\n'), (146, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (154, 'torch.set_num_threads', 'torch.set_num_threads', False, 'import torch\n'), (155, 'horovod.torch.rank', 'hvd.rank', True, 'import horovod.torch as hvd\n'), (181, 'horovod.torch.broadcast_optimizer_state', 'hvd.broadcast_optimizer_state', True, 'import horovod.torch as hvd\n'), (239, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (240, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (241, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (242, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (267, 'ray_shuffling_data_loader.torch_dataset.TorchShufflingDataset', 'TorchShufflingDataset', False, 'from ray_shuffling_data_loader.torch_dataset import TorchShufflingDataset\n'), (286, 'ray.init', 'ray.init', False, 'import ray\n'), (336, 'horovod.ray.RayExecutor.create_settings', 'RayExecutor.create_settings', False, 'from horovod.ray import RayExecutor\n'), (337, 'horovod.ray.RayExecutor', 'RayExecutor', False, 'from horovod.ray import RayExecutor\n'), (127, 'torch.nn.Conv2d', 'nn.Conv2d', True, 'import torch.nn as nn\n'), (128, 'torch.nn.Conv2d', 'nn.Conv2d', True, 'import torch.nn as nn\n'), (129, 'torch.nn.Dropout2d', 'nn.Dropout2d', True, 'import torch.nn as nn\n'), (130, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (131, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (138, 'torch.nn.functional.dropout', 'F.dropout', True, 'import torch.nn.functional as F\n'), (140, 'torch.nn.functional.log_softmax', 'F.log_softmax', True, 'import torch.nn.functional as F\n'), (148, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (151, 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', False, 'import torch\n'), (166, 'horovod.torch.size', 'hvd.size', True, 'import horovod.torch as hvd\n'), (168, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (199, 'timeit.default_timer', 'timeit.default_timer', False, 'import timeit\n'), (221, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (222, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (223, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (224, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (252, 'time.sleep', 'time.sleep', False, 'import time\n'), (261, 'ray_shuffling_data_loader.data_generation.DATA_SPEC.keys', 'DATA_SPEC.keys', False, 'from ray_shuffling_data_loader.data_generation import generate_data, DATA_SPEC\n'), (294, 'tempfile.gettempdir', 'tempfile.gettempdir', False, 'import tempfile\n'), (296, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (308, 'ray_shuffling_data_loader.data_generation.generate_data', 'generate_data', False, 'from ray_shuffling_data_loader.data_generation import generate_data, DATA_SPEC\n'), (150, 'horovod.torch.local_rank', 'hvd.local_rank', True, 'import horovod.torch as hvd\n'), (161, 'horovod.torch.size', 'hvd.size', True, 'import horovod.torch as hvd\n'), (172, 'horovod.torch.nccl_built', 'hvd.nccl_built', True, 'import horovod.torch as hvd\n'), (173, 'horovod.torch.local_size', 'hvd.local_size', True, 'import horovod.torch as hvd\n'), (214, 'time.sleep', 'time.sleep', False, 'import time\n'), (219, 'timeit.default_timer', 'timeit.default_timer', False, 'import timeit\n'), (220, 'timeit.default_timer', 'timeit.default_timer', False, 'import timeit\n'), (263, 'ray_shuffling_data_loader.data_generation.DATA_SPEC.values', 'DATA_SPEC.values', False, 'from ray_shuffling_data_loader.data_generation import generate_data, DATA_SPEC\n'), (204, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (299, 'pickle.load', 'pickle.load', False, 'import pickle\n'), (314, 'pickle.dump', 'pickle.dump', False, 'import pickle\n'), (315, 'ray_shuffling_data_loader.stats.human_readable_size', 'human_readable_size', False, 'from ray_shuffling_data_loader.stats import human_readable_size\n'), (203, 'timeit.default_timer', 'timeit.default_timer', False, 'import timeit\n'), (312, 'tempfile.gettempdir', 'tempfile.gettempdir', False, 'import tempfile\n')]