# Author:  Lisandro Dalcin
# Contact: dalcinl@gmail.com
"""Convert NumPy and MPI datatypes."""
# pylint: disable=too-many-locals
# pylint: disable=too-many-branches
# pylint: disable=too-many-statements
# pylint: disable=too-many-return-statements

from .. import MPI

try:
    from numpy import dtype as _np_dtype
except ImportError:  # pragma: no cover
    _np_dtype = None


def _get_datatype(dtype):
    return MPI.Datatype.fromcode(dtype.char)


def _get_typecode(datatype):
    # pylint: disable=protected-access
    return MPI._typecode(datatype)


def _get_alignment(datatype):
    # pylint: disable=protected-access
    return MPI._typealign(datatype)


def _is_aligned(datatype, offset=0):
    """Determine whether an MPI datatype is aligned."""
    if datatype.is_predefined:
        if offset == 0:
            return True
        alignment = _get_alignment(datatype) or 0
        return offset % alignment == 0

    combiner = datatype.combiner
    basetype, _, info = datatype.decode()
    types, disps = [basetype], [0]
    try:
        if combiner == MPI.COMBINER_RESIZED:
            disps = [info["extent"]]
        if combiner == MPI.COMBINER_STRUCT:
            types = info["datatypes"]
            disps = info["displacements"]
        if combiner == MPI.COMBINER_HVECTOR:
            disps = [info["stride"] if info["count"] > 1 else 0]
        if combiner == MPI.COMBINER_HINDEXED:
            disps = info["displacements"]
        if combiner == MPI.COMBINER_HINDEXED_BLOCK:
            disps = info["displacements"]
        return all(_is_aligned(t, offset + d) for (t, d) in zip(types, disps))
    finally:
        for _tp in types:
            if not _tp.is_predefined:
                _tp.Free()


def from_numpy_dtype(dtype):
    """Convert NumPy datatype to MPI datatype."""
    if _np_dtype is None:
        raise RuntimeError("NumPy is not available") from None
    dtype = _np_dtype(dtype)

    if dtype.hasobject:
        raise ValueError("NumPy datatype with object entries")
    if not dtype.isnative:
        raise ValueError("NumPy datatype with non-native byteorder")

    # struct data type
    fields = dtype.fields
    if fields:
        blocklengths = []
        displacements = []
        datatypes = []
        try:
            for name in dtype.names or ():
                ftype, fdisp, *_ = fields[name]
                blocklengths.append(1)
                displacements.append(fdisp)
                datatypes.append(from_numpy_dtype(ftype))
            datatype = MPI.Datatype.Create_struct(
                blocklengths, displacements, datatypes
            )
        finally:
            for mtp in datatypes:
                mtp.Free()
        try:
            return datatype.Create_resized(0, dtype.itemsize)
        finally:
            datatype.Free()

    # subarray data type
    subdtype = dtype.subdtype
    if subdtype:
        base, shape = subdtype
        datatype = from_numpy_dtype(base)
        try:
            if len(shape) == 1:
                return datatype.Create_contiguous(shape[0])
            starts = (0,) * len(shape)
            return datatype.Create_subarray(shape, shape, starts)
        finally:
            datatype.Free()

    # elementary data type
    datatype = _get_datatype(dtype)
    return datatype.Dup()


def to_numpy_dtype(datatype):
    """Convert MPI datatype to NumPy datatype."""

    def mpi2npy(datatype, count):
        dtype = to_numpy_dtype(datatype)
        return dtype if count == 1 else (dtype, count)

    def np_dtype(spec, **kwargs):
        if _np_dtype is None:
            return spec if not kwargs else (spec, kwargs)
        return _np_dtype(spec, **kwargs)

    if datatype == MPI.DATATYPE_NULL:
        raise ValueError("cannot convert null MPI datatype to NumPy")

    combiner = datatype.combiner

    # named elementary datatype
    if combiner == MPI.COMBINER_NAMED:
        # elementary datatype
        typecode = _get_typecode(datatype)
        if typecode is not None:
            return np_dtype(typecode)
        # pair datatype for MINLOC/MAXLOC reductions
        names = ("SHORT", "INT", "LONG", "FLOAT", "DOUBLE", "LONG_DOUBLE")
        types = [getattr(MPI, f"{name}_INT") for name in names]
        typename = names[types.index(datatype)]
        typecode = _get_typecode(getattr(MPI, typename))
        return np_dtype(f"{typecode},i", align=True)

    # user-defined datatype
    basetype, _, info = datatype.decode()
    datatypes = [basetype]
    try:
        # duplicated datatype
        if combiner == MPI.COMBINER_DUP:
            return to_numpy_dtype(basetype)

        # contiguous datatype
        if combiner == MPI.COMBINER_CONTIGUOUS:
            dtype = to_numpy_dtype(basetype)
            count = info["count"]
            return np_dtype((dtype, (count,)))

        # subarray datatype
        if combiner == MPI.COMBINER_SUBARRAY:
            dtype = to_numpy_dtype(basetype)
            sizes = info["sizes"]
            subsizes = info["subsizes"]
            starts = info["starts"]
            order = info["order"]
            if subsizes == sizes and min(starts) == max(starts) == 0:
                if order == MPI.ORDER_FORTRAN:
                    sizes = sizes[::-1]
                return np_dtype((dtype, tuple(sizes)))
            raise ValueError("cannot convert subarray MPI datatype to NumPy")

        # value-index datatype
        if combiner == MPI.COMBINER_VALUE_INDEX:  # pragma: no cover
            value = to_numpy_dtype(info["value"])
            index = to_numpy_dtype(info["index"])
            datatypes = [value, index]
            return np_dtype(
                [
                    ("f0", value),
                    ("f1", index),
                ],
                align=True,
            )

        # struct datatype
        aligned = True
        if combiner == MPI.COMBINER_RESIZED:
            if basetype.combiner == MPI.COMBINER_STRUCT:
                aligned = _is_aligned(basetype, info["extent"])
                combiner = MPI.COMBINER_STRUCT
                _, _, info = basetype.decode()
                datatypes.pop().Free()
        if combiner == MPI.COMBINER_STRUCT:
            datatypes = info["datatypes"]
            blocklengths = info["blocklengths"]
            displacements = info["displacements"]
            names = [f"f{i}" for i in range(len(datatypes))]
            formats = list(map(mpi2npy, datatypes, blocklengths))
            offsets = displacements
            itemsize = datatype.extent
            aligned &= all(map(_is_aligned, datatypes, offsets))
            return np_dtype({
                "names": names,
                "formats": formats,
                "offsets": offsets,
                "itemsize": itemsize,
                "aligned": aligned,
            })

        # vector datatype
        combiner_vector = (
            MPI.COMBINER_VECTOR,
            MPI.COMBINER_HVECTOR,
        )
        if combiner in combiner_vector:
            dtype = to_numpy_dtype(basetype)
            count = info["count"]
            blocklength = info["blocklength"]
            stride = info["stride"]
            if combiner == MPI.COMBINER_VECTOR:
                stride *= basetype.extent
                aligned = _is_aligned(basetype)
            if combiner == MPI.COMBINER_HVECTOR:
                stride = stride if count > 1 else 0
                aligned = _is_aligned(basetype, stride)
            names = [f"f{i}" for i in range(count)]
            formats = [(dtype, (blocklength,))] * count
            offsets = [stride * i for i in range(count)]
            itemsize = datatype.extent
            return np_dtype({
                "names": names,
                "formats": formats,
                "offsets": offsets,
                "itemsize": itemsize,
                "aligned": aligned,
            })

        # indexed datatype
        combiner_indexed = (
            MPI.COMBINER_INDEXED,
            MPI.COMBINER_HINDEXED,
            MPI.COMBINER_INDEXED_BLOCK,
            MPI.COMBINER_HINDEXED_BLOCK,
        )
        if combiner in combiner_indexed:
            dtype = to_numpy_dtype(basetype)
            stride = 1
            aligned = _is_aligned(basetype)
            blocklengths = []
            displacements = info["displacements"]
            if combiner in combiner_indexed[:2]:
                blocklengths = info["blocklengths"]
            if combiner in combiner_indexed[2:]:
                blocklengths = [info["blocklength"]] * len(displacements)
            if combiner in combiner_indexed[0::2]:
                stride = basetype.extent
            if combiner in combiner_indexed[1::2]:
                aligned &= all(_is_aligned(basetype, d) for d in displacements)
            names = [f"f{i}" for i in range(len(displacements))]
            formats = [(dtype, (blen,)) for blen in blocklengths]
            offsets = [disp * stride for disp in displacements]
            return np_dtype({
                "names": names,
                "formats": formats,
                "offsets": offsets,
                "aligned": aligned,
            })

        # Fortran 90 datatype
        combiner_f90 = (
            MPI.COMBINER_F90_INTEGER,
            MPI.COMBINER_F90_REAL,
            MPI.COMBINER_F90_COMPLEX,
        )
        if combiner in combiner_f90:
            datatypes.pop()
            typecode = _get_typecode(datatype)
            return np_dtype(typecode)

        raise ValueError("cannot convert MPI datatype to NumPy")
    finally:
        for _tp in datatypes:
            if not _tp.is_predefined:
                _tp.Free()
