import torch
from torch import nn
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.distributed.parallel_loader as pl
from torchvision import transforms

from projectcode.datasets.cifar10 import CifarDataset
import search.cifar10_search as my_cifar10
import datetime
from misc.flops_counter import add_flops_counting_methods
import os
import numpy as np
from collections import Callable
from loguru import logger
from misc.utils import Cutout
import argparse

# import torch_optimizer as optim

from projectcode.training.optim import Lamb


def get_map_fn(
    model: torch.nn.Module,
    train_params: dict,
    data_root: str,
    momentum: float,
    weight_decay: float,
    CIFAR_CLASSES: int,
    learning_rate: float,
    layers: int,
    batch_size: int,
    epochs: int,
    drop_path_prob: float = 0.0,
    save_pth: str = "",
    args: argparse.Namespace = None,
    WRAPPED_MODEL= None,
    clip = None,

) -> Callable:
    """
    Defines the distributed training function. Must run before training each individual to generate the
    training function for the individual. Also evaluates the individual and saves the results in a .pt file.

    TODO: Refactor as a class
    TODO: Fix the parameters by consolidating them, possibly via the args

    Args:
        model:
        train_params:
        data_root:
        momentum:
        weight_decay:
        CIFAR_CLASSES:
        learning_rate:
        layers:
        batch_size:
        epochs:
        drop_path_prob:
        save_pth: not actually used

    Returns:
        A function to be executed via
        FLAGS = {}
        xmp.spawn(map_fn, args=(FLAGS,), nprocs=8, start_method='fork')

    """
    params = train_params
    # WRAPPED_MODEL = xmp.MpModelWrapper(model)

    def train(
        train_queue: pl.ParallelLoader,
        net: torch.nn.Module,
        criterion: Callable,
        optimizer: torch.optim.Optimizer,
        params: dict,
        device: torch.device,
    ):
        """
        Training loop for a network.

        Args:
            train_queue:
            net:
            criterion:
            optimizer:
            params:
            device: pass the automatically assigned device via xm.spawn

        Returns:
            num correct, total loss, total num examples
        """

        net.train()
        train_loss = 0
        correct = 0
        total = 0

        for step, (inputs, targets) in enumerate(train_queue):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs, outputs_aux = net(inputs)
            loss = criterion(outputs, targets)

            loss.backward()
            nn.utils.clip_grad_norm_(net.parameters(), params["grad_clip"])

            xm.optimizer_step(optimizer)

            net.apply(clip)

            train_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

        return correct, train_loss, total

    def infer(
        valid_queue: pl.ParallelLoader,
        net: torch.nn.Module,
        criterion: Callable,
        device: torch.device,
    ):
        """
        Evaluation loop for a network.

        Args:
            valid_queue:
            net:
            criterion:
            device: pass the automatically assigned device via xm.spawn

        Returns:
            mean accuracy, mean loss
        """

        net.eval()
        test_loss = 0
        correct = 0
        total = 0

        with torch.no_grad():
            for step, (inputs, targets) in enumerate(valid_queue):
                inputs, targets = inputs.to(device), targets.to(device)
                outputs, _ = net(inputs)
                loss = criterion(outputs, targets)

                test_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()

        acc = 100.0 * correct / total

        return acc, test_loss / total

    def run():
        """
        Main function to setup the training loop and evaluation loop.
        See comments for detailed explanation.

        Returns:
            None, but it saves the model weights and model performance, based on the get_map_fn arguments

        """

        # xla will assign a device for each forked run of this function
        device = xm.xla_device()

        # determine if this fork is the master fork to avoid logging and print 8 times
        master = xm.is_master_ordinal()

        if master:
            logger.info("running at batch size %i" % batch_size)

        criterion = nn.CrossEntropyLoss()

        criterion.to(device)
        model = WRAPPED_MODEL.to(device)

        # standard data prep
        CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
        CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]

        train_transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.RandomCrop(32),
                transforms.RandomHorizontalFlip(),
                transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
            ]
        )

        if args.cutout > 0:
            train_transform.transforms.append(Cutout(args.cutout))

        train_data = CifarDataset(transform=train_transform)

        # distributed samples ensure data is sharded to each tpu core
        # if you do not use this, you are only using 1 of the 8 cores
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_data,
            num_replicas=xm.xrt_world_size(),
            rank=xm.get_ordinal(),
            shuffle=True,
        )

        train_queue = torch.utils.data.DataLoader(
            train_data,
            batch_size=batch_size//xm.xrt_world_size(),
            sampler=train_sampler,
            drop_last=True,
            num_workers=0,
        )

        valid_transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
            ]
        )

        valid_data = my_cifar10.CIFAR10(
            root=data_root, train=False, download=False, transform=valid_transform
        )

        valid_sampler = torch.utils.data.distributed.DistributedSampler(
            valid_data,
            num_replicas=xm.xrt_world_size(),
            rank=xm.get_ordinal(),
            shuffle=False,
        )

        valid_queue = torch.utils.data.DataLoader(
            valid_data,
            sampler=valid_sampler,
            batch_size=batch_size//xm.xrt_world_size(),
            drop_last=True,
            num_workers=0,
        )

        # standard optimizer stuff
        parameters = filter(lambda p: p.requires_grad, model.parameters())

        if args.opt == "sgd":

            optimizer = torch.optim.SGD(
                parameters,
                args.learning_rate,
                momentum=momentum,
                weight_decay=args.weight_decay,
            )
        elif args.opt == "lamb":
            optimizer = Lamb(
                parameters, lr=args.learning_rate, weight_decay=weight_decay
            )
        else:
            raise NameError("Unknown Optimizer %s" % args.opt)

        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, int(epochs))

        # training by epoch loop
        for epoch in range(epochs):

            # the model needs a droprate, so just assign it
            model.droprate = drop_path_prob * epoch / epochs

            start = datetime.datetime.now()
            st = start.strftime("%Y-%m-%d %H:%M:%S")

            if master:
                logger.info("starting epoch %i at %s" % (epoch, st))

            # parallel loader necessary to load data in parallel to each core
            para_loader = pl.ParallelLoader(train_queue, [device]).per_device_loader(
                device
            )
            correct, train_loss, total = train(
                para_loader, model, criterion, optimizer, params, device
            )

            train_acc = 100 * correct / total

            # collect the train accuracies from all cores
            train_acc = xm.mesh_reduce("avg acc", train_acc, np.mean)

            end = datetime.datetime.now()
            duration = (end - start).total_seconds()

            if master:
                logger.info("train_acc %f duration %f" % (train_acc, duration))

            scheduler.step()

        # validate using 8 cores and collect results
        valid_acc, valid_obj = infer(valid_queue, model, criterion, device)
        valid_acc = xm.mesh_reduce("val avg acc", valid_acc, np.mean)

        if master:
            logger.info("valid_acc %f" % valid_acc)

        # count flops
        _ = add_flops_counting_methods(model)
        model.eval()
        model.start_flops_count()
        random_data = torch.randn(1, 3, 32, 32)
        model(torch.autograd.Variable(random_data).to(device))
        n_flops = np.round(model.compute_average_flops_cost() / 1e6, 4)
        n_flops = xm.mesh_reduce("flops", n_flops, np.mean)

        if master:
            logger.info("flops %f" % n_flops)

        if master:
            logger.info("saving")

        # save weights and results

        xm.save([valid_acc, n_flops], "results.pt")

    def mp_fn(rank, flags):
        torch.set_default_tensor_type("torch.FloatTensor")
        run()

    return mp_fn
