import os
from contextlib import suppress
from dataclasses import dataclass, field, InitVar
from typing import Optional, Dict

import torch
from torch.distributed import ReduceOp

try:
    import torch_xla.core.xla_model as xm
    _HAS_XLA = True
except ImportError as e:
    xm = None
    _HAS_XLA = False

try:
    # only the very latest XLA builds have AMP
    import torch_xla.amp as xa
except ImportError as e:
    xa = None

from .device_env import DeviceEnv, DeviceEnvType, TensorList


_PT_TO_XM_OP = {
    ReduceOp.SUM: 'sum',
    ReduceOp.PRODUCT: 'mul',
    ReduceOp.MIN: 'min',
    ReduceOp.MAX: 'max',
    ReduceOp.BAND: 'and',
    ReduceOp.BOR: 'or',
}


def is_xla_available(xla_device_type=None):
    if not _HAS_XLA:
        return False
    supported_devs = xm.get_xla_supported_devices(devkind=xla_device_type)
    return len(supported_devs) >= 1


@dataclass
class DeviceEnvXla(DeviceEnv):

    def __post_init__(
            self,
            device_type: Optional[str],
            device_idx: Optional[int],
            channels_last: bool,
    ):
        if device_type is not None:
            device_type = device_type.upper()
            assert device_type in ('TPU', 'GPU', 'CPU'), "XLA device type must be one of ('TPU', 'GPU', 'CPU')"
        self.device = xm.xla_device(n=device_idx, devkind=device_type)
        self.world_size = xm.xrt_world_size()
        if self.distributed:
            assert device_idx is None, "device_index is based on local rank for distributed XLA mode"
            self.local_rank = xm.get_local_ordinal()
            self.global_rank = xm.get_ordinal()
        else:
            self.local_rank = 0
            self.global_rank = 0
        if self.amp:
            assert xa is not None, 'XLA AMP is not present on this build'
        if self.autocast is None:
            self.autocast = xa.autocast if self.amp else suppress
        if channels_last:
            self.memory_format = torch.channels_last

    @property
    def type(self) -> DeviceEnvType:
        return DeviceEnvType.XLA

    def wrap_distributed(self, *modules):
        wrapped = [m for m in modules]  # NO-OP
        return wrapped[0] if len(wrapped) == 1 else wrapped

    def wrap_parallel(self, *modules):
        assert False, "Not implemented"

    def mark_step(self):
        xm.mark_step()

    def all_reduce(self, tensor: torch.Tensor, op=ReduceOp.SUM, average=False):
        assert isinstance(tensor, torch.Tensor)  # unlike in-place variant, lists/tuples not allowed
        op = _PT_TO_XM_OP[op]
        scale = 1.0 / self.world_size if average else 1.0
        return xm.all_reduce(op, tensor, scale=scale)

    def all_reduce_(self, tensor: TensorList, op=ReduceOp.SUM, average=False):
        op = _PT_TO_XM_OP[op]
        scale = 1.0 / self.world_size if average else 1.0
        wrapped = False
        if isinstance(tensor, torch.Tensor):
            tensor = [tensor]  # bare tensors are not operated on in-place
            wrapped = True
        xm.all_reduce(op, tensor, scale=scale)
        if wrapped:
            tensor = tensor[0]
        return tensor

    def all_gather(self, tensor: torch.Tensor, cat_dim=0):
        output = xm.all_gather(tensor, cat_dim)
        return output

    def all_to_all(self, tensor, num_splits, split_dim, cat_dim=0):
        output = xm.all_to_all(tensor, split_dim, cat_dim, num_splits)
        return output

    def broadcast(self, tensor: torch.Tensor, src_rank=0):
        if self.global_rank != src_rank:
            reduce_tensor = torch.zeros_like(tensor)
            xm.all_reduce('sum', reduce_tensor)
        else:
            xm.all_reduce('sum', tensor)
        return tensor

    def broadcast_(self, tensor: torch.Tensor, src_rank=0):
        out_tensor = self.broadcast(tensor, src_rank)
        return tensor.copy_(out_tensor)

    def barrier(self):
        xm.rendezvous('timm.bits.dist_barrier')

    def state_dict_to_cpu(self, state: Dict[str, torch.Tensor]):
        cpu_state = xm._maybe_convert_to_cpu(state, convert=True)
        return cpu_state

    def state_dict_to_device(self, state: Dict[str, torch.Tensor]):
        device_state = xm.send_cpu_data_to_device(state, device=self.device)
        return device_state