import numpy as _np
from mpi4py import MPI as _MPI

from jax import core, typeof
from jax.ffi import ffi_lowering
from jax.interpreters import ad, batching
from jax.core import ShapedArray

from mpi4jax._src.utils import (
    HashableMPIType,
    default_primitive_impl,
    to_dtype_handle,
    to_mpi_handle,
    to_mpi_ptr,
    unpack_hashable,
    wrap_as_hashable,
    ordered_effect,
    NOTSET,
    raise_if_token_is_set,
)
from mpi4jax._src.jax_compat import (
    register_lowering,
    get_token_effect,
    set_token_effect,
    Primitive,
)
from mpi4jax._src.decorators import (
    translation_rule_cpu,
    translation_rule_cuda,
    translation_rule_xpu,
)
from mpi4jax._src.validation import enforce_types
from mpi4jax._src.utils import get_default_comm

# The Jax primitive
mpi_sendrecv_p = Primitive("sendrecv_mpi")  # Create the primitive
mpi_sendrecv_impl = default_primitive_impl(mpi_sendrecv_p)


# This function applies the primitive to an AST
@enforce_types(
    source=_np.integer,
    dest=_np.integer,
    sendtag=_np.integer,
    recvtag=_np.integer,
    comm=(type(None), _MPI.Intracomm, HashableMPIType),
    status=(type(None), _MPI.Status, HashableMPIType),
)
def sendrecv(
    sendbuf,
    recvbuf,
    source,
    dest,
    *,
    sendtag=0,
    recvtag=_MPI.ANY_TAG,
    comm=None,
    status=None,
    token=NOTSET,
):
    """Perform a sendrecv operation.

    .. warning::

        Unlike mpi4py's sendrecv, this returns a *new* array with the received data.

    Arguments:
        sendbuf: Array or scalar input to send.
        recvbuf: Array or scalar input with the correct shape and dtype. This can
           contain arbitrary data and will not be overwritten.
        source (int): Rank of the source MPI process.
        dest (int): Rank of the destination MPI process.
        sendtag (int): Tag of this message for sending.
        recvtag (int): Tag of this message for receiving.
        comm (mpi4py.MPI.Comm): The MPI communicator to use (defaults to
            a clone of :obj:`COMM_WORLD`).
        status (mpi4py.MPI.Status): Status object, can be used for introspection.

    Returns:
        DeviceArray: Received data.

    """
    raise_if_token_is_set(token)

    if comm is None:
        comm = get_default_comm()

    comm = wrap_as_hashable(comm)

    if status is not None:
        status = wrap_as_hashable(status)

    return mpi_sendrecv_p.bind(
        sendbuf,
        recvbuf,
        source=source,
        dest=dest,
        sendtag=sendtag,
        recvtag=recvtag,
        comm=comm,
        status=status,
        _must_transpose=False,
    )


def _mpi_sendrecv_xla_encode(
    ctx,
    sendbuf,
    recvbuf,
    source,
    dest,
    sendtag,
    recvtag,
    comm,
    status,
    _must_transpose=False,
):
    """Common lowering for all platforms using jax.ffi.ffi_lowering."""
    from mpi4jax._src.xla_bridge import MPI_STATUS_IGNORE_ADDR

    if _must_transpose:
        raise RuntimeError(
            "sendrecv cannot be used with forward-mode (vjp) autodiff, because "
            "the gradient might be located on a different mpi rank than the "
            "desired one. Use reverse-mode (jvp) differentiation instead."
        )

    comm = unpack_hashable(comm)
    status = unpack_hashable(status)

    send_aval, recv_aval, *_ = ctx.avals_in
    send_nptype = send_aval.dtype
    recv_nptype = recv_aval.dtype

    send_nitems = _np.prod(send_aval.shape, dtype=_np.int64)
    send_dtype_handle = _np.int64(to_dtype_handle(send_nptype))

    recv_nitems = _np.prod(recv_aval.shape, dtype=_np.int64)
    recv_dtype_handle = _np.int64(to_dtype_handle(recv_nptype))

    if status is None:
        status_ptr = _np.int64(MPI_STATUS_IGNORE_ADDR)
    else:
        status_ptr = _np.int64(to_mpi_ptr(status))

    comm_handle = _np.int64(to_mpi_handle(comm))

    token = get_token_effect(ctx, ordered_effect)
    operands = (sendbuf, token)

    # sendrecv takes (sendbuf, token) but ctx.avals_in has (sendbuf, recvbuf)
    # Output is (result, token) and ctx.avals_out has (result,)
    ctx_with_token = ctx.replace(
        avals_in=(send_aval, core.abstract_token),
        avals_out=(*ctx.avals_out, core.abstract_token),
    )

    lowering_rule = ffi_lowering(
        "mpi_sendrecv_ffi",
        has_side_effect=True,
    )

    results = lowering_rule(
        ctx_with_token,
        *operands,
        sendcount=send_nitems,
        dest=_np.int64(dest),
        sendtag=_np.int64(sendtag),
        sendtype=send_dtype_handle,
        recvcount=recv_nitems,
        source=_np.int64(source),
        recvtag=_np.int64(recvtag),
        recvtype=recv_dtype_handle,
        comm=comm_handle,
        status=status_ptr,
    )

    results = list(results)
    token = results.pop(-1)
    set_token_effect(ctx, ordered_effect, token)

    return results


# Platform-specific lowering rules (all use the same FFI implementation)
mpi_sendrecv_xla_encode_cpu = translation_rule_cpu(_mpi_sendrecv_xla_encode)
mpi_sendrecv_xla_encode_cuda = translation_rule_cuda(_mpi_sendrecv_xla_encode)
mpi_sendrecv_xla_encode_xpu = translation_rule_xpu(_mpi_sendrecv_xla_encode)


# This function evaluates only the shapes during AST construction
def mpi_sendrecv_abstract_eval(
    sendbuf,
    recvbuf,
    source,
    dest,
    sendtag,
    recvtag,
    comm,
    status,
    _must_transpose=False,
):
    return ShapedArray(recvbuf.shape, recvbuf.dtype), {ordered_effect}


def mpi_sendrecv_batch_eval(
    in_args,
    batch_axes,
    source,
    dest,
    sendtag,
    recvtag,
    comm,
    status,
    _must_transpose=False,
):
    sendbuf, recvbuf = in_args

    assert batch_axes[0] == batch_axes[1]
    ax = batch_axes[0]

    res = mpi_sendrecv_p.bind(
        sendbuf,
        recvbuf,
        source=source,
        dest=dest,
        sendtag=sendtag,
        recvtag=recvtag,
        comm=comm,
        status=status,
        _must_transpose=_must_transpose,
    )
    return res, ax


def mpi_sendrecv_value_and_jvp(
    in_args,
    tan_args,
    source,
    dest,
    sendtag,
    recvtag,
    comm,
    status,
    _must_transpose=False,
):
    sendbuf, recvbuf = in_args
    send_tan, recv_tan = tan_args

    val = mpi_sendrecv_p.bind(
        sendbuf,
        recvbuf,
        source=source,
        dest=dest,
        sendtag=sendtag,
        recvtag=recvtag,
        comm=comm,
        status=status,
        _must_transpose=_must_transpose,
    )

    jvp = mpi_sendrecv_p.bind(
        send_tan,
        recv_tan,
        source=source,
        dest=dest,
        sendtag=sendtag,
        recvtag=recvtag,
        comm=comm,
        status=status,
        _must_transpose=not _must_transpose,
    )

    return val, jvp


def mpi_sendrecv_transpose_rule(
    out_tan, *x_args, source, dest, sendtag, recvtag, comm, status, _must_transpose
):
    # swap the sender and receiver
    res = mpi_sendrecv_p.bind(
        out_tan,
        out_tan,
        source=dest,
        dest=source,
        sendtag=sendtag,
        recvtag=recvtag,
        comm=comm,
        status=status,
        _must_transpose=not _must_transpose,
    )
    return (res, ad.Zero(typeof(res)))


mpi_sendrecv_p.def_impl(mpi_sendrecv_impl)
mpi_sendrecv_p.def_effectful_abstract_eval(mpi_sendrecv_abstract_eval)

batching.primitive_batchers[mpi_sendrecv_p] = mpi_sendrecv_batch_eval

ad.primitive_jvps[mpi_sendrecv_p] = mpi_sendrecv_value_and_jvp
ad.primitive_transposes[mpi_sendrecv_p] = mpi_sendrecv_transpose_rule

# assign to the primitive the correct encoder
register_lowering(mpi_sendrecv_p, mpi_sendrecv_xla_encode_cpu, platform="cpu")
register_lowering(mpi_sendrecv_p, mpi_sendrecv_xla_encode_cuda, platform="cuda")
register_lowering(mpi_sendrecv_p, mpi_sendrecv_xla_encode_xpu, platform="xpu")
