import os
import random as pyrandom

import pymoo
from torch import nn
from collections.abc import Callable
from typing import Union
from loguru import logger
import torch
import numpy as np

import main_config
from projectcode.server.database import get_task_by_network_id, pickle_load_from_str
from projectcode.weightmanagement.gcs import download_blob
from search import micro_encoding

layer_types_not_initialized = set()


def get_init_weights(func: Callable, bias: Union[bool, float] = 0.01) -> Callable:
    """
    Constructs a weight initialization function to be applied to a pytorch module object

    Args:
        func: callable
        bias: if True, apply the func to the bias. Otherwise, set the bias to the float value. Do not pass false.

    Returns:
        callable

    """

    def init_weights(m):
        if not type(m) == nn.BatchNorm2d:
            try:
                try:
                    func(m.weight)
                    if type(bias) == bool and bias and m.bias is not None:
                        func(m.bias)
                    else:
                        if type(bias) == bool and not bias:
                            logger.warning(
                                "Received false as apply to bias, assuming bias to be 0.00"
                            )
                            apply_to_bias = 0.0
                        else:
                            apply_to_bias = bias

                        m.bias.data.fill_(apply_to_bias)

                except AttributeError as e:
                    layer_types_not_initialized.add(str(e))

            except:
                print(m)
                raise

    return init_weights


def initialize_null(model: torch.nn.Module) -> torch.nn.Module:
    """
    Sets all weights to negative 100, so we can check if weights were set correctly via a future
    initialization method.

    Args:
        model:

    Returns:
        model
    """

    def zero_out(m):
        return torch.nn.init.constant_(m, -101.0)

    init_weights = get_init_weights(zero_out, bias=True)
    model.apply(init_weights)
    return model


def assert_non_null_weights(statedict: dict):
    """
    Checks if a statedict has weights defined as null (-100)
    Args:
        statedict:

    Returns:

    """
    for w in statedict.items():
        assert not np.any(w[1].cpu().numpy() == -101.0), (
            w[0] + "contains parameters that were likely not initialized"
        )


def determine_more_fit_parent(parent1: dict, parent2: dict) -> list:
    """
    TODO: refactor this to the weight manager class to allow for better metric selection

    Given two parents, determine the more fit parent.

    Args:
        parent1:
        parent2:

    Returns:
        list([more fit parent, less fit parent])

    """
    metric = pyrandom.choice(["flops", "error"])

    if parent1[metric] < parent2[metric]:
        parent_better = parent1
        parent_worse = parent2
    else:
        parent_better = parent2
        parent_worse = parent1

    return [parent_better, parent_worse]


def decode_individual(individual) -> tuple:
    """
    Takes an individual, with a genome as a numpy array, and decodes it into a
    human friendly genome.

    Args:
        individual:

    Returns:

    """

    return micro_encoding.decode(micro_encoding.convert(individual.X))


def convert_mutation_map(individual) -> tuple:
    """
    Converts a mutation map that describes what genes mutated into a format that matches
    the genome length.

    Args:
        individual:

    Returns:
        normal_cell_map, reduce_cell_map
    """
    genome_length = len(decode_individual(individual).normal)

    normal_muta_map = np.all(
        individual.mutation_map[: genome_length * 2].reshape(genome_length, 2), axis=1
    )
    reduce_muta_map = np.all(
        individual.mutation_map[genome_length * 2 :].reshape(genome_length, 2), axis=1
    )

    return normal_muta_map, reduce_muta_map


def read_parent_by_id(id: int, expr_root: str = None, args = None) -> dict:
    """
    based on the id, find the directory of the parent and parse its relevant data for wt inheritance.

    Args:
        id: one of the parent ids saved on the individual after traced mating
        expr_root: this should match args.save_path

    Returns:
        A dictionary describing the parent
    """

    if main_config.distributed_cloud:
        task = get_task_by_network_id(id, run_code=args.code)
        if task is None:
            logger.warning("Could not find expected parent id")
        spec = pickle_load_from_str(task.pkl_data)

        results = pickle_load_from_str(task.result_pkl_data)

        genome = micro_encoding.decode(spec["genome"])
        flops = results["flops"]
        acc = results["valid_acc"]

        wt_path = download_blob(results["wt_blob_name"])

    else:
        # parent is a dict
        arch_path = os.path.join(expr_root, "arch_{}".format(id))

        path = os.path.join(arch_path, "log.txt")
        wt_path = os.path.join(arch_path, "weights.pt")

        with open(path, "r") as f:
            s = f.read()

        genome = eval(s.split("\n")[1][15:])
        flops = float(s.split("\n")[3][8:-2])
        acc = float(s.split("\n")[4][11:])

    error = 100 - acc
    weights = torch.load(wt_path)

    has_nan = False
    for k, v in weights.items():
        if torch.any(torch.isnan(v.cpu())):
            has_nan = True

        weights[k] = torch.nan_to_num(v.cpu())
    if has_nan:
        logger.warning("Nan values detected when trying to load weights for parent with id %i" % id)

    parent = {"genome": genome, "flops": flops, "error": error, "weights": weights}
    return parent
