# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.
"""These functions are used by FullyShardedDataParallel to help consolidate and shard optimizer states."""
import copy
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Tuple, cast

import torch

from fairscale.nn.misc import FlattenParamsWrapper

if TYPE_CHECKING:
    from fairscale.nn.data_parallel import FullyShardedDataParallel

# These return keys are used by fairseq. To change, add @sshleifer as a reviewer.
UNFLAT_RETURN_KEYS = {"state", "param_groups", "uncollected_local_ids", "param_id_map"}

# This function helps shard a full optimizer state dict
def flatten_optim_state_dict(sd: Dict) -> Dict:
    """Shard a full optimizer state dict (called by FSDP.get_shard_from_optim_state_dict)"""
    param_id_map = sd["param_id_map"]
    num_local_params = len(set(param_id_map.values()))
    if sd["state"]:
        new_state: Dict = {local_id: {} for local_id in range(num_local_params)}
        singleton_state: Dict = copy.deepcopy(new_state)
    else:
        new_state = {}
    non_tensor_state = {}

    # Populate `new_state["state"]`. (Assuming sd is sorted)
    for global_id, buffers in sd["state"].items():
        local_id = param_id_map[global_id]
        for buffer_name, p in buffers.items():
            if is_singleton_tensor(p):
                singleton_state[local_id][buffer_name] = p
            elif torch.is_tensor(p):
                if buffer_name not in new_state[local_id]:
                    new_state[local_id][buffer_name] = []
                new_state[local_id][buffer_name].append(p.reshape(-1))
            elif isinstance(p, list):
                singleton_state[local_id][buffer_name] = p
            else:
                non_tensor_state[buffer_name] = p
    # Now combine all tensors in each buffer using torch.cat().
    for local_id, state in new_state.items():
        for buffer_name, tensors in state.items():
            new_state[local_id][buffer_name] = torch.cat(tensors)
        new_state[local_id].update(non_tensor_state)
        new_state[local_id].update(singleton_state[local_id])
    new_sd = {"state": new_state, "param_groups": copy.deepcopy(sd["param_groups"])}
    for k in sd.keys():  # if there are extra keys, like loss_scale, don't delete them
        if k not in UNFLAT_RETURN_KEYS:
            new_sd[k] = copy.deepcopy(sd[k])

    # add pointers from the `params` dict.
    for pg_id, _ in enumerate(sd["param_groups"]):
        # TODO: this list could be huge. Can we avoid materializing?
        new_sd["param_groups"][pg_id]["params"] = list(range(num_local_params))

    return new_sd


def check_param_counts_before_sharding(full_optim_state_dict: Dict, n_instances: int) -> None:
    n_local_params_in_opt = len(set(full_optim_state_dict["param_id_map"].values()))
    msg = (
        f"Including itself, this model has {n_instances} nested instances. When the optimizer state was saved "
        f"there were {n_local_params_in_opt}"
    )
    stateless = len(full_optim_state_dict["state"]) == 0
    assert stateless or (n_instances == n_local_params_in_opt), msg


# All functions below here help saving the list of optimizer states, one from each rank
# build_unflat_state_dict is the interface used by FSDP
def _extract_non_tensor_state(combined_state: Dict[int, Dict[str, List]], param_id: int) -> Dict:
    non_tensor_state = {}  # This state is like the `step` count in Adam, not a tensor so we dont unpad or cat it.
    for k, v in combined_state[param_id].items():
        if torch.is_tensor(v[0]):
            continue
        elif len(set(v)) == 1:
            non_tensor_state[k] = v[0]
        else:
            raise TypeError(f"Dont know how to consolidate optimizer param {k} with values {v}")
    return non_tensor_state


def _unflatten_optim_state(
    combined_state: Dict[int, Dict],
    instance_list: List["FullyShardedDataParallel"],
    world_pad_info: List[List[List[int]]],
    singleton_state: Dict[int, Dict],
) -> Tuple[Dict[int, Dict], Dict[int, int]]:
    """Convert optimizer state for flattened parameters into original, unflatten ones."""
    # local ids are the keys in the current state (combined_state), (usually fewer)
    # global ids will be the keys in the unflattened state
    next_global_id = 0  # gets incremented
    pad_info = {id: [s[id][0] for s in world_pad_info] for id in combined_state}
    local_ids = [id for id in sorted(combined_state.keys())]

    # non_tensor_state refers to entries in sd[state][param_id] that are not tensors, like "step".
    # we check that these are identical across workers and then take the first
    non_tensor_state = [_extract_non_tensor_state(combined_state, id) for id in combined_state]

    # Local corresponds to flattened, global corresponds to unflattened.
    # Casting needed only for mypy.
    num_global_params: List[int] = []
    for m in instance_list:
        if m.flatten_parameters:
            num_flatten = cast(int, m.num_params_managed)
            num_global_params.append(num_flatten)
        else:
            num_global_params.append(len(m.non_shared_params()))
    global_to_local_id = {}
    for local_id, num_unflat in enumerate(num_global_params):
        for _ in range(num_unflat):
            global_to_local_id[next_global_id] = local_id
            next_global_id += 1
    if not combined_state:
        return {}, global_to_local_id

    # copy non tensor state (like the "step" count) to all global entries
    unflat_state = {i: copy.deepcopy(non_tensor_state[0]) for i in range(sum(num_global_params))}

    if non_tensor_state[0].keys() == combined_state[0].keys():
        return unflat_state, global_to_local_id

    local_to_global: Dict[int, List] = {i: [] for i in local_ids}
    for g, l in global_to_local_id.items():
        local_to_global[l].append(g)
    # loop over parameters in state.
    # Tensor state will be padded, concatenated, and restored to original shape with FlattenParamsWrapper.get_views
    # get_views returns multiple tensors, each of which is a new parameter with a new "global" id.
    for local_id in local_ids:
        # undo the work of shard_parameters
        for k, v in combined_state[local_id].items():
            if k in non_tensor_state[local_id]:
                continue
            assert isinstance(v, list), f"got {k}: {v} for {local_id}"
            v_unpad = [t[:-np] if np > 0 else t for t, np in zip(v, pad_info[local_id])]
            flat_buffer = torch.cat(v_unpad)
            if instance_list[local_id].flatten_parameters:
                # Unflatten. Casting needed only for mypy.
                param_views: Iterator = cast(FlattenParamsWrapper, instance_list[local_id]).get_param_views(
                    [flat_buffer]
                )
                for global_id, param_view in zip(sorted(local_to_global[local_id]), param_views):
                    assert k not in unflat_state[global_id], f"already added {k} to {global_id} {local_id}"
                    unflat_state[global_id][k] = param_view
            else:
                # Copy non-flatten state directly.
                assert len(local_to_global[local_id]) == 1, "Only support a single non-flatten parameter"
                global_id = local_to_global[local_id][0]
                unflat_state[global_id][k] = flat_buffer
            unflat_state[global_id].update(singleton_state[local_id])

    return unflat_state, global_to_local_id


def build_unflat_state_dict(
    instance_list: List["FullyShardedDataParallel"],
    world_pad_info: List[List[List[int]]],
    state: Dict[int, Dict[str, List[torch.Tensor]]],
    singleton_state: Dict[int, Dict[str, List[torch.Tensor]]],
    uncollected_opt_state: Dict[int, Dict],
    param_groups: List[Dict],
) -> Dict:
    """Build an unflattened optimizer state dict given a list of flattened optimizer state dicts from each rank."""
    assert all(len(s) == len(instance_list) for s in world_pad_info)
    assert all(len(s[0]) == 1 for s in world_pad_info)

    # Use uncollected_opt_state to update tensor_state, singleton_state
    for local_id, v in uncollected_opt_state.items():
        assert local_id not in state
        state[local_id] = {buffer_name: [x] for buffer_name, x in v.items() if not is_singleton_tensor(x)}
        singleton_state[local_id] = {buffer_name: [x] for buffer_name, x in v.items() if is_singleton_tensor(x)}
    # local ids are in the current state, global_ids will be in returned state.
    unflat_state, global_to_local_id = _unflatten_optim_state(state, instance_list, world_pad_info, singleton_state)
    # Since there are no tensors in param_groups, deepcopy is fine.
    param_groups = copy.deepcopy(param_groups)
    # Casting needed only for mypy.
    num_params = sum([cast(int, m.num_params_managed) for m in instance_list])
    param_groups[0]["params"] = list(range(num_params))
    unflat_optim_state_dict = {
        "state": dict(sorted(unflat_state.items())),  # NOTE: this is probably already sorted
        "param_id_map": global_to_local_id,
        "param_groups": param_groups,
        "uncollected_local_ids": list(uncollected_opt_state.keys()),
    }
    assert set(unflat_optim_state_dict.keys()) == UNFLAT_RETURN_KEYS
    return unflat_optim_state_dict


def is_singleton_tensor(x: Any) -> bool:
    """Is x a dimensionless tensor?"""
    return torch.is_tensor(x) and x.dim() == 0
