# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os
import json
import numpy as np
import math
from packaging import version as pkg_version

from .config import ElasticityConfig, ElasticityConfigError, ElasticityError, \
    ElasticityIncompatibleWorldSize
from .constants import ELASTICITY, ENABLED, ENABLED_DEFAULT, LATEST_ELASTICITY_VERSION, \
    MINIMUM_DEEPSPEED_VERSION, DEEPSPEED_ELASTICITY_CONFIG
from ..git_version_info import version as __version__
from ..utils import logger

# Thirty eight smallest highly composite numbers. The list should
# be enough to support up to 720K batch size.
HCN_LIST = [
    1, 2, 4, 6, 12, 24, 36, 48, 60, 120, 180, 240, 360, 720, 840, 1260, 1680, 2520, 5040, 7560, 10080, 15120, 20160,
    25200, 27720, 45360, 50400, 55440, 83160, 110880, 166320, 221760, 277200, 332640, 498960, 554400, 665280, 720720
]


def get_candidate_batch_sizes(base_list, max_acceptable_batch_size):
    candidate_batch_size = []
    for base in base_list:
        if base >= max_acceptable_batch_size:
            candidate_batch_size.append(base)
        else:
            value = max_acceptable_batch_size // base
            index = np.argmax(np.asarray(HCN_LIST) > value)
            candidate_batch_size.append(HCN_LIST[index - 1] * base)
    candidate_batch_size = list(set(candidate_batch_size))
    logger.info(f"Candidate batch size: {candidate_batch_size}")
    return candidate_batch_size


def get_valid_gpus(batch_size, micro_batches, min_valid_gpus, max_valid_gpus):
    valid_gpus = []
    for micro_batch in micro_batches:
        if batch_size % micro_batch == 0:

            max_gpus = batch_size // micro_batch
            if min_valid_gpus <= max_gpus <= max_valid_gpus:
                valid_gpus.append(max_gpus)

            # find all factors less than max_gpus / 2
            for i in range(1, max_gpus // 2 + 1):
                if i > max_valid_gpus:
                    break
                if i < min_valid_gpus:
                    continue
                if max_gpus % i == 0:
                    valid_gpus.append(i)
    valid_gpus = set(valid_gpus)
    valid_gpus = sorted(list(valid_gpus))
    return valid_gpus


def get_best_candidates(candidate_batch_sizes, micro_batches, min_gpus, max_gpus, prefer_larger):

    max_valid_gpus = 0
    valid_gpus = None
    final_batch_size = int(min(micro_batches))

    for batch_size in candidate_batch_sizes:

        current_valid_gpus = get_valid_gpus(batch_size, micro_batches, min_gpus, max_gpus)

        if (len(current_valid_gpus) > max_valid_gpus or (len(current_valid_gpus) == max_valid_gpus and
                                                         ((prefer_larger and batch_size > final_batch_size) or
                                                          (not prefer_larger and batch_size < final_batch_size)))):
            max_valid_gpus = len(current_valid_gpus)
            valid_gpus = current_valid_gpus
            final_batch_size = batch_size

    return final_batch_size, valid_gpus


def _get_compatible_gpus_v01(micro_batches,
                             max_acceptable_batch_size,
                             min_gpus=None,
                             max_gpus=None,
                             prefer_larger=True):
    '''We use two heuristics to compute the batch size
        1. We use the Lowest Common Multiple of the micro-batches
    as the base batch size and scale it by a HCN such that the result is
    the largest batch size less than the max_acceptable batch size
        2. We use each of the micro batches as a base and scale it
    by a HCN such that the result is the largest batch size less than the
    max_acceptable batch size.

    We then use brute force to count the number of compatible GPU count for
    each of the aforementioned cases, and return the batch size with the most number of
    compatible GPU counts in the min-max GPU range if provided, other wise
    we return the batch size with the most number of total compatible GPU counts.

    Returns:
        final_batch_size
        valid_gpus
    '''
    min_gpus = min_gpus or 1
    max_gpus = max_gpus or max_acceptable_batch_size // min(micro_batches)

    if not all(mb <= max_acceptable_batch_size for mb in micro_batches):
        raise ValueError(f"All micro batches must be less than \
            or equal to max_acceptable_batch_size: {max_acceptable_batch_size}")

    lcm = np.lcm.reduce(micro_batches)

    base_list = []
    base_list.extend(micro_batches)
    base_list.append(lcm)

    candidate_batch_sizes = get_candidate_batch_sizes(base_list, max_acceptable_batch_size)

    final_batch_size, valid_gpus = get_best_candidates(candidate_batch_sizes, micro_batches, min_gpus, max_gpus,
                                                       prefer_larger)

    return final_batch_size, valid_gpus


def _get_compatible_gpus_v02(micro_batches,
                             max_acceptable_batch_size,
                             current_num_gpus,
                             min_gpus=None,
                             max_gpus=None,
                             prefer_larger=True,
                             num_gpus_per_node=1,
                             model_parallel_size=1):
    '''
    Returns:
        final_batch_size
        valid_gpus
        micro-batch size
    '''
    if num_gpus_per_node % model_parallel_size != 0:
        raise ElasticityError(
            f"In Elasticity v0.2, number of GPUs per node:" \
            f"{num_gpus_per_node} should be divisible by " \
            f"model parallel size {model_parallel_size}")

    def get_microbatch(final_batch_size):
        candidate_microbatch = None

        for micro_batch in micro_batches:
            if final_batch_size // current_num_gpus % micro_batch == 0:
                if candidate_microbatch is None:
                    candidate_microbatch = micro_batch
                if prefer_larger and candidate_microbatch < micro_batch:
                    candidate_microbatch = micro_batch
        return candidate_microbatch

    dp_size_per_node = num_gpus_per_node // model_parallel_size

    final_batch_size, valid_world_size = _get_compatible_gpus_v01(
        micro_batches,
        int(max_acceptable_batch_size / dp_size_per_node),
        int(min_gpus / num_gpus_per_node),
        int(max_gpus / num_gpus_per_node),  # Passing number of max nodes as Elasticity v2 works at node level
        prefer_larger=prefer_larger)

    final_batch_size = int(final_batch_size) * dp_size_per_node
    valid_dp_world_size = [i * dp_size_per_node for i in valid_world_size]
    if current_num_gpus // model_parallel_size in valid_dp_world_size:
        candidate_microbatch = get_microbatch(final_batch_size)
        return final_batch_size, valid_dp_world_size, candidate_microbatch

    current_dp_size = (current_num_gpus / num_gpus_per_node) * dp_size_per_node
    candidate_batch_sizes = []
    for micro_batch in micro_batches:
        min_batch_size = micro_batch * current_dp_size

        factor = math.floor(max_acceptable_batch_size / float(min_batch_size))
        candidate_batch_sizes.append(factor * min_batch_size)

    used_microbatch = None
    if prefer_larger:
        candidate_batch_size = max(candidate_batch_sizes)
    else:
        candidate_batch_size = min(candidate_batch_sizes)

    candidate_microbatch = get_microbatch(candidate_batch_size)

    return candidate_batch_size, [int(current_dp_size)], candidate_microbatch


def _compatible_ds_version_check(target_deepspeed_version: str):
    min_version = pkg_version.parse(MINIMUM_DEEPSPEED_VERSION)
    target_version = pkg_version.parse(target_deepspeed_version)

    err_str = f"Target deepspeed version of {target_deepspeed_version} is not compatible " \
        f"with minimum version {MINIMUM_DEEPSPEED_VERSION} supporting elasticity."
    if target_version < min_version:
        raise ElasticityError(err_str)
    return True


def elasticity_enabled(ds_config: dict):
    if ELASTICITY not in ds_config:
        return False
    return ds_config[ELASTICITY].get(ENABLED, ENABLED_DEFAULT)


def ensure_immutable_elastic_config(runtime_elastic_config_dict: dict):
    """
    Ensure the resource scheduler saw the same elastic config we are using at runtime
    """
    if DEEPSPEED_ELASTICITY_CONFIG in os.environ:
        scheduler_elastic_config_dict = json.loads(os.environ[DEEPSPEED_ELASTICITY_CONFIG])
        scheduler_elastic_config = ElasticityConfig(scheduler_elastic_config_dict)
        runtime_elastic_config = ElasticityConfig(runtime_elastic_config_dict)
        err_str = "Elastic config '{}={}' seen by resource scheduler does not match config passed to runtime {}={}"
        if runtime_elastic_config.max_acceptable_batch_size != scheduler_elastic_config.max_acceptable_batch_size:
            raise ElasticityConfigError(
                err_str.format('max_acceptable_batch_size', scheduler_elastic_config.max_acceptable_batch_size,
                               'max_acceptable_batch_size', runtime_elastic_config.max_acceptable_batch_size))
        if runtime_elastic_config.micro_batches != scheduler_elastic_config.micro_batches:
            raise ElasticityConfigError(
                err_str.format('micro_batches', scheduler_elastic_config.micro_batches, 'micro_batches',
                               runtime_elastic_config.micro_batches))
        if runtime_elastic_config.version != scheduler_elastic_config.version:
            raise ElasticityConfigError(
                err_str.format('version', scheduler_elastic_config.version, 'version', runtime_elastic_config.version))
    else:
        logger.warning("Unable to find DEEPSPEED_ELASTICITY_CONFIG environment variable, cannot " \
            "guarantee resource scheduler will scale this job using compatible GPU counts.")


def compute_elastic_config(ds_config: dict, target_deepspeed_version: str, world_size=0, return_microbatch=False):
    """Core deepspeed elasticity API. Given an elastic config (similar to the example below)
    DeepSpeed will compute a total train batch size corresponding valid GPU count list that
    provides a high level of elasticity. Elasticity in this case means we are safe to scale
    the training job up/down across the GPU count list *without* any negative impacts on
    training convergence. This is achievable primarily due to DeepSpeed's gradient accumulation
    feature which allows us to decompose a global training batch size into:
    micro-batch-size * gradient-accumulation-steps * world-size.

    "elasticity": {
        "enabled": true,
        "max_train_batch_size": 2000,
        "micro_batch_sizes": [2,4,6],
        "min_gpus": 1,
        "max_gpus" : 10000
        "min_time": 20
        "version": 0.1
    }

    Intended to be called both by scheduling infrastructure and deepspeed runtime.
    For the same `ds_config` we should return deterministic results.

    Args:
        ds_config (dict): DeepSpeed config dictionary/json
        target_deepspeed_version (str): When called from scheduling
            infrastructure we want to ensure that the target deepspeed version is
            compatible with the elasticity version used in the backend.
        world_size (int, optional): Intended/current DP world size, will do some sanity
            checks to ensure world size is actually valid with the config.
        return_microbatch (bool, optional): whether to return micro batch size or not.

    Raises:
        ElasticityConfigError: Missing required elasticity config or elasticity disabled
        ElasticityError: If target deepspeed version is not compatible with current version

    Returns:
        final_batch_size (int): total batch size used for training
        valid_gpus (list(int)): list of valid GPU counts with this config
        micro_batch_size (int, optional): if world_size is provided will return
            specific micro batch size
    """
    if not isinstance(ds_config, dict):
        raise ValueError("Expected ds_config to be a dictionary but received " \
            f"a {type(ds_config)}, containing: {ds_config}")

    if ELASTICITY not in ds_config:
        raise ElasticityConfigError(f"'{ELASTICITY}' is missing from config json," \
            " please add it if running an elastic training job.")

    elastic_config_dict = ds_config[ELASTICITY]
    if not elastic_config_dict.get(ENABLED, ENABLED_DEFAULT):
        raise ElasticityConfigError("Elasticity is disabled, please enable it " \
            "('enabled':true) if running an elastic training job.")

    elastic_config = ElasticityConfig(elastic_config_dict)
    model_parallel_size = elastic_config.model_parallel_size
    num_gpus_per_node = elastic_config.num_gpus_per_node

    if model_parallel_size > 1 and float(elastic_config.version) != 0.2:
        raise ElasticityConfigError(f"Elasticity V{elastic_config.version} " \
            f"does not support model-parallel training. Given model-parallel size: " \
            f"{model_parallel_size}")

    if float(elastic_config.version) > LATEST_ELASTICITY_VERSION:
        raise ElasticityConfigError("Attempting to run elasticity version " \
            f"{elastic_config.version} but runtime only supports up " \
            f"to {LATEST_ELASTICITY_VERSION}")

    # Ensure target deepspeed version works with intended elasticity version
    if not _compatible_ds_version_check(target_deepspeed_version):
        raise ElasticityError("Unable to run elasticity on target deepspeed version of" \
            f" {target_deepspeed_version}, currently {__version__}")

    if float(elastic_config.version) == 0.1:
        final_batch_size, valid_gpus = _get_compatible_gpus_v01(
            micro_batches=elastic_config.micro_batches,
            max_acceptable_batch_size=elastic_config.max_acceptable_batch_size,
            min_gpus=elastic_config.min_gpus,
            max_gpus=elastic_config.max_gpus,
            prefer_larger=elastic_config.prefer_larger_batch_size)
        # ensure batch size is int dtype
        final_batch_size = int(final_batch_size)
    elif float(elastic_config.version) == 0.2:
        if world_size != 0:
            current_num_gpus = world_size
        else:
            if "WORLD_SIZE" in os.environ and \
                os.getenv('WORLD_SIZE').isnumeric():
                current_num_gpus = int(os.getenv('WORLD_SIZE'))
            else:
                WORLD_SIZE = os.getenv('WORLD_SIZE')
                raise ElasticityConfigError(
                    'Elasticity V 0.2 needs WORLD_SIZE '\
                    'to compute valid batch size. '\
                    'Either give it as argument to function compute_elastic_config '\
                    'or set it as an environment variable. '\
                    f'Value of WORLD_SIZE as environment variable is {WORLD_SIZE}')

        final_batch_size, valid_gpus, candidate_microbatch_size = _get_compatible_gpus_v02(
            micro_batches=elastic_config.micro_batches,
            max_acceptable_batch_size=elastic_config.max_acceptable_batch_size,
            current_num_gpus=current_num_gpus,
            min_gpus=elastic_config.min_gpus,
            max_gpus=elastic_config.max_gpus,
            prefer_larger=elastic_config.prefer_larger_batch_size,
            num_gpus_per_node=num_gpus_per_node,
            model_parallel_size=model_parallel_size)
        # ensure batch size is int dtype
        final_batch_size = int(final_batch_size)
    else:
        raise NotImplementedError(f"Unable to find elastic logic for version: {elastic_config.version}")

    logger.info(f"Valid World Size (GPUs / Model Parallel Size): {valid_gpus}")

    if world_size > 0:
        if world_size not in valid_gpus:
            raise ElasticityIncompatibleWorldSize(f"World size ({world_size}) is not valid " \
        f"with the current list of valid GPU counts: {valid_gpus}")

        # Pick largest valid micro batch size
        micro_batch_size = None
        for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True):
            if final_batch_size // world_size % mbsz == 0:
                micro_batch_size = mbsz
                break
        assert micro_batch_size is not None, "Unable to find divisible micro batch size" \
            f" world_size={world_size}, final_batch_size={final_batch_size}, and " \
            f" micro_batches={elastic_config.micro_batches}."
        return final_batch_size, valid_gpus, micro_batch_size

    if return_microbatch:
        # Pick a valid micro batch size
        if float(elastic_config.version) == 0.2:
            return final_batch_size, valid_gpus, candidate_microbatch_size
        else:
            micro_batch_size = None
            for mbsz in sorted(list(set(elastic_config.micro_batches)), reverse=True):
                if final_batch_size // world_size % mbsz == 0:
                    micro_batch_size = mbsz
                    break
            assert micro_batch_size is not None, "Unable to find divisible micro batch size" \
                    f" world_size={world_size}, final_batch_size={final_batch_size}, and " \
                    f" micro_batches={elastic_config.micro_batches}."
            return final_batch_size, valid_gpus, micro_batch_size

    return final_batch_size, valid_gpus