import torch
import torch.distributed as dist
from typing import Iterable, Any


def average_tensor(value, device):
    if not torch.is_tensor(value):
        value = torch.tensor(value, device=device, dtype=torch.float32)
    else:
        value = value.to(device)

    if dist.is_available() and dist.is_initialized():
        dist.all_reduce(value, op=dist.ReduceOp.SUM)
    return value


def move_to_device(batch, device):
    if isinstance(batch, torch.Tensor):
        return batch.to(device, non_blocking=True)
    if isinstance(batch, (list, tuple)):
        return type(batch)(move_to_device(x, device) for x in batch)
    if isinstance(batch, dict):
        return {k: move_to_device(v, device) for k, v in batch.items()}
    return batch


def _chunk_tensor(x: torch.Tensor, chunk_size: int) -> Iterable[torch.Tensor]:
    if x.dim() == 0:
        return [x]
    return torch.split(x, chunk_size, dim=0)


def _chunk_any(obj: Any, chunk_size: int) -> Iterable[Any]:
    if torch.is_tensor(obj):
        return _chunk_tensor(obj, chunk_size)
    if isinstance(obj, (list, tuple)):
        chunks = [list(_chunk_any(o, chunk_size)) for o in obj]
        # transpose list of lists
        return type(obj)(zip(*chunks))
    if isinstance(obj, dict):
        chunked = {k: list(_chunk_any(v, chunk_size)) for k, v in obj.items()}
        return [type(obj)({k: chunked[k][i] for k in chunked}) for i in range(len(next(iter(chunked.values()))))]
    return [obj]


def iter_microbatches(batch, microbatch_size: int):
    """
    Split batch into microbatches along the first dimension when tensors/lists/tuples/dicts.
    """
    if microbatch_size is None or microbatch_size <= 0:
        yield batch
        return

    for mb in _chunk_any(batch, microbatch_size):
        yield mb


class EMA:
    """
    Exponential Moving Average for model parameters.
    """

    def __init__(self, model: torch.nn.Module, decay: float = 0.999):
        self.decay = decay
        self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()}

    def update(self, model: torch.nn.Module):
        for k, v in model.state_dict().items():
            if k not in self.shadow:
                self.shadow[k] = v.detach().clone()
                continue
            self.shadow[k].mul_(self.decay).add_(v.detach(), alpha=1 - self.decay)

    def apply_to(self, model: torch.nn.Module):
        model.load_state_dict(self.shadow, strict=False)

    def state_dict(self):
        return {"decay": self.decay, "shadow": self.shadow}

    def load_state_dict(self, state):
        self.decay = state.get("decay", self.decay)
        self.shadow = state["shadow"]
