# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from contextlib import nullcontext

import jax
import jax.numpy as jnp
from jax.experimental.layout import Format
from jax.sharding import Mesh

from tpu_inference import envs

# Lazy initialized, since device might not be ready at import time.
_cpu_mesh = None


def reorder_concatenated_tensor_for_sharding(concatenated_tensor: jax.Array,
                                             split_sizes: list[int],
                                             n_shards: int, dim: int):
    """
    Reorder a replicated concatenated tensor such that when sharded on multiple chips, each shard is a concatenation of the shards of the individual tensors.
    For example, let the concatenated_tensor be:
        AAAAAAAAAAAABBBBBBBBCCCC
            12 As     8 Bs  4 Cs
    and let the split_sizes = [12, 8, 4] and n_shards = 4.
    The output is:
        AAABBCAAABBCAAABBCAAABBC
    In other words, it reorders the input tensor into 4 segements, with each segment corresponding to a shard and being AAABBC.
    Args:
        concatenated_tensor: the tensor, concatenated on the dimension specified by `dim`.
        split_sizes: each individual tensor's size on the dimension specified by `dim`.
        n_shards: num of shards.
        dim: the dimension on which the concatenated_tensor is concatenated.
    """
    # Split the concatenated tensor into individual tensors.
    if dim < 0:
        dim += concatenated_tensor.ndim
    split_tensors = []
    start_offset = 0
    old_shape = concatenated_tensor.shape
    # New shape ensures each split_tensor[i] maps to a tensor in ith shards
    new_shape = old_shape[:dim] + (n_shards, -1) + old_shape[dim + 1:]
    for split_size in split_sizes:
        split_tensor = jax.lax.slice_in_dim(concatenated_tensor,
                                            start_offset,
                                            start_offset + split_size,
                                            axis=dim)
        split_tensors.append(split_tensor.reshape(new_shape))
        start_offset += split_size
    # While maintaining 0th dim as a shard dim, we concatenate along 1th dim to
    # to create concatenated tnensor where 0th dim maps to shard dim.
    reordered_tensor = jnp.concatenate(split_tensors, axis=dim + 1)
    return reordered_tensor.reshape(old_shape)


def slice_sharded_tensor_for_concatenation(sharded_tensor: jax.Array,
                                           split_sizes: list[int],
                                           n_shards: int):
    """
    Slice the input tensor which is sharded on multiple chips (on the last dim) into individual tensors with the same sharding.
    For example, let the sharded_tensor be:
        AAABBC | AAABBC | AAABBC | AAABBC
        Shard0   Shard1   Shard2   Shard3
    and let the split_sizes = [12, 8, 4] and n_shards = 4.
    The output is a list of 3 tensors:
         AAA   |  AAA   |  AAA   |  AAA
          BB   |   BB   |   BB   |   BB
           C   |    C   |    C   |    C
        Shard0   Shard1   Shard2   Shard3
    In other words, each individual tensor is a slice of the input tensor with the same sharding.
    Args:
        sharded_tensor: the input tensor, sharded on the last dim.
        split_sizes: each individual tensor's size on the last dim.
        n_shards: num of shards.
    """
    new_shape = sharded_tensor.shape[:-1] + (n_shards, -1)
    # New shape ensures each sharded_tensor[:, i] maps to a tensor in ith shards
    sharded_tensor = sharded_tensor.reshape(new_shape)

    split_tensors = []
    start_offset = 0
    for split_size in split_sizes:
        assert split_size % n_shards == 0
        sz = split_size // n_shards  # size of this split tensor per shard
        end_offset = start_offset + sz
        # Because we are slicing over last dim, sharding dim remains intact.
        # Therefore, splitting happens locally.
        split_tensor = sharded_tensor[..., start_offset:end_offset]
        split_tensors.append(split_tensor.reshape(new_shape[:-2] + (-1, )))
        start_offset = end_offset

    return split_tensors


def general_device_put(tensor: jax.Array,
                       sharding,
                       *,
                       layout=None,
                       source_mesh=None) -> jax.Array:
    """
    Put a tensor onto devices with the given sharding.
    This method handles both single-host and multi-host cases.

    `source_mesh` specifies the mesh on which the input tensor is currently located.
    """

    def _put(t):
        multihost_backend = envs.TPU_MULTIHOST_BACKEND
        if multihost_backend != "ray":
            return jax.device_put(t, sharding)

        # NOTE: at here, num_global_devices != num_local_devices
        # meaning we are in multi-host setup. Each host will run the same process
        # and each process only need to handle the devices accessible to this host.
        shape = t.shape
        ctx = nullcontext() if source_mesh is None else jax.set_mesh(
            source_mesh)
        # `t[i]` needs to be operated in the same mesh as `t`, which is provided as
        # `source_mesh`.
        with ctx:
            x_split = [
                jax.device_put(t[i], device) for device, i in
                sharding.addressable_devices_indices_map(shape).items()
            ]
        global_array = jax.make_array_from_single_device_arrays(shape,
                                                                sharding,
                                                                x_split,
                                                                dtype=t.dtype)
        if layout is not None:
            dst_mesh = sharding.mesh
            with jax.set_mesh(dst_mesh):
                global_array = jax.device_put(global_array,
                                              Format(layout, sharding))
        return global_array

    return jax.tree_util.tree_map(_put, tensor)


def cpu_mesh() -> Mesh:
    global _cpu_mesh
    if _cpu_mesh is None:
        _cpu_mesh = Mesh(jax.devices("cpu")[:1], ("cpu", ))
    return _cpu_mesh


def cpu_mesh_context():
    """A context to enforce using CPU mesh, used for loading weights on CPU."""
    return jax.set_mesh(cpu_mesh())
