from misc import utils
import os
from loguru import logger
import numpy as np
from search import micro_encoding
from models.micro_models import NetworkCIFAR as Network
from projectcode.weightmanagement import manage_weights

import torch
import pymoo
import argparse
from compress_pickle import dump, load
from projectcode.weightmanagement.gcs import upload_blob
import main_config
from projectcode.training.clipper import weightClip

def pickled_train_and_evaluate(
    pkl_path: str = None, spec: dict = None, client_id: str = None
):

    if spec is None:
        logger.info("Reading pickled training spec at %s" % pkl_path)
        spec = load(pkl_path)
    else:
        logger.info("Using Spec")

    performance = train_and_evaluate(
        genome=spec["genome"],
        individual=spec["individual"],
        save=spec["save"],
        args=spec["args"],
        first_gen=spec["first_gen"],
        client_id=client_id,
    )

    return performance


def train_and_evaluate(
    genome: tuple,
    individual=None,
    args: argparse.Namespace = None,
    first_gen: bool = True,
    save: str = None,
    client_id: str = None,
):
    """
    Function to train and evaluate an individual using a TPU.

    Results are always saved in the save dir to make distributed data management easier.

    Args:
        first_gen:
        genome:
        save:
        individual:
        args:

    Returns:

    """

    if args.stream == "tpu":
        # must warp up TPU
        import torch_xla

    auxiliary = False

    assert hasattr(individual, "id")

    if not first_gen:
        # this is not the first generation, so mating should have occurred
        assert hasattr(individual, "parents")

    expr_root = ""

    save_pth = os.path.join(expr_root, "{}".format(save))
    utils.create_exp_dir(save_pth)

    CIFAR_CLASSES = 10
    learning_rate = 0.025
    momentum = 0.9
    weight_decay = 3e-4
    data_root = "../data"
    batch_size = args.batch_size
    auxiliary_weight = 0.4
    grad_clip = 5
    report_freq = 50
    train_params = {
        "auxiliary": auxiliary,
        "auxiliary_weight": auxiliary_weight,
        "grad_clip": grad_clip,
        "report_freq": report_freq,
    }

    if args.search_space == "micro":
        genotype = micro_encoding.decode(genome)
        model = Network(
            args.init_channels, CIFAR_CLASSES, args.layers, auxiliary, genotype
        )

        if not first_gen:
            # change the way the weights are set up
            model = manage_weights(model, individual, expr_root, args)

    elif args.search_space == "macro":
        raise NotImplementedError("Not supported")
    else:
        raise NameError("Unknown search space type")

    logger.info("Architecture = %s", genotype)

    try:
        max_weight = args.max_weight
    except:
        print("Could Not Determine Maximum Weight Argument")
        max_weight = 1e20

    clip = weightClip(max_weight=max_weight, min_weight= max_weight* -1)

    if args.stream == "tpu":
        from projectcode.training.tpu import get_map_fn
        import torch_xla.distributed.xla_multiprocessing as xmp

        WRAPPED_MODEL = xmp.MpModelWrapper(model)

        logger.info("Executing TPU Training")
        map_fn = get_map_fn(
            model,
            train_params,
            data_root,
            momentum,
            weight_decay,
            CIFAR_CLASSES,
            learning_rate,
            args.layers,
            batch_size,
            epochs=args.epochs,
            save_pth=save_pth,
            args=args,
            WRAPPED_MODEL = WRAPPED_MODEL,
            clip = clip
        )

        FLAGS = {}

        xmp.spawn(map_fn, args=(FLAGS,), nprocs=1, start_method="fork")

        valid_acc, n_flops = torch.load("results.pt")
    elif args.stream == "gpu":
        from projectcode.training.gpu import train_gpu
        logger.info("Executing GPU Training")
        valid_acc, n_flops = train_gpu(
            model,
            train_params,
            data_root,
            momentum,
            weight_decay,
            CIFAR_CLASSES,
            learning_rate,
            args.layers,
            batch_size,
            epochs=args.epochs,
            save_pth=save_pth,
            args=args,
            clip = clip
        )

    else:

        raise NameError("Unrecognized client stream")

    n_params = (
        np.sum(
            np.prod(v.size())
            for v in filter(lambda p: p.requires_grad, model.parameters())
        )
        / 1e6
    )

    if main_config.distributed_cloud and args.weight_init == "lammarckian":
        wt_path = f"{args.code}_{client_id}_weights_{individual.id:05d}.pt"
        torch.save(model.state_dict(), wt_path)
        blob_name = upload_blob(wt_path)
    else:
        blob_name = None
        torch.save(model.state_dict(), os.path.join(save_pth, "weights.pt"))

    result_dict = {
        "id": individual.id,
        "save_path": save_pth,
        "valid_acc": valid_acc,
        "params": n_params,
        "flops": n_flops,
        "wt_blob_name": blob_name,
    }

    dump(result_dict, os.path.join(save_pth, "result.pkl"))

    return result_dict
