"""Utility functions for manipulating JAX embedding tables and inputs."""

import collections
import dataclasses
import typing
from typing import Any, Mapping, NamedTuple, Sequence, TypeAlias, TypeVar

import jax
import numpy as np
from jax import numpy as jnp
from jax_tpu_embedding.sparsecore.lib.nn import embedding
from jax_tpu_embedding.sparsecore.lib.nn.embedding_spec import FeatureSpec
from jax_tpu_embedding.sparsecore.lib.nn.embedding_spec import StackedTableSpec
from jax_tpu_embedding.sparsecore.lib.nn.embedding_spec import TableSpec

from keras_rs.src.types import Nested

T = TypeVar("T")

# Any to support tf.Ragged without needing an explicit TF dependency.
ArrayLike: TypeAlias = jax.Array | np.ndarray | Any  # type: ignore
Shape: TypeAlias = tuple[int, ...]


class FeatureSamples(NamedTuple):
    tokens: ArrayLike
    weights: ArrayLike


class ShardedCooMatrix(NamedTuple):
    shard_starts: ArrayLike
    shard_ends: ArrayLike
    col_ids: ArrayLike
    row_ids: ArrayLike
    values: ArrayLike


class InputStatsPerTable(NamedTuple):
    max_ids_per_partition: int
    max_unique_ids_per_partition: int
    required_buffer_size_per_device: int


def _round_up_to_multiple(value: int, multiple: int) -> int:
    return ((value + multiple - 1) // multiple) * multiple


def _default_stacked_table_spec(
    table_spec: TableSpec, num_shards: int, batch_size: int
) -> StackedTableSpec:
    return StackedTableSpec(
        stack_name=table_spec.name,
        stack_vocab_size=_round_up_to_multiple(
            table_spec.vocabulary_size, 8 * num_shards
        ),
        stack_embedding_dim=_round_up_to_multiple(table_spec.embedding_dim, 8),
        optimizer=table_spec.optimizer,
        combiner=table_spec.combiner,
        total_sample_count=batch_size,
        max_ids_per_partition=table_spec.max_ids_per_partition,
        max_unique_ids_per_partition=table_spec.max_unique_ids_per_partition,
    )


def _get_stacked_table_spec(
    table_spec: TableSpec, num_shards: int, batch_size: int = 0
) -> StackedTableSpec:
    return table_spec.stacked_table_spec or _default_stacked_table_spec(
        table_spec, num_shards, batch_size
    )


def pad_table(
    table_spec: TableSpec,
    table_values: jax.Array,
    num_shards: int,
    pad_value: jnp.float32 = jnp.nan,
) -> jax.Array:
    """Adds appropriate padding to a table to prepare for stacking.

    Args:
        table_spec: Table specification describing the table to pad.
        table_values: Table values array to pad.
        num_shards: Number of shards in the table (typically
            `global_device_count * num_sc_per_device`).
        pad_value: Value to use for padding.

    Returns:
        Padded table values.
    """
    vocabulary_size = table_spec.vocabulary_size
    embedding_dim = table_spec.embedding_dim
    padded_vocabulary_size = _round_up_to_multiple(
        vocabulary_size, 8 * num_shards
    )
    stack_embedding_dim = _get_stacked_table_spec(
        table_spec, num_shards
    ).stack_embedding_dim
    return jnp.pad(
        table_values,
        (
            (0, padded_vocabulary_size - vocabulary_size),
            (0, stack_embedding_dim - embedding_dim),
        ),
        constant_values=pad_value,
    )


def _stack_and_shard_table(
    stacked_table: jax.Array,
    table_spec: TableSpec,
    table: jax.Array,
    num_shards: int,
    pad_value: jnp.float32,
) -> jax.Array:
    """Stacks and shards a single table for use in sparsecore lookups."""
    padded_values = pad_table(table_spec, table, num_shards, pad_value)
    sharded_padded_vocabulary_size = padded_values.shape[0] // num_shards
    stack_embedding_dim = stacked_table.shape[-1]

    # Mod-shard vocabulary across devices.
    sharded_values = jnp.swapaxes(
        padded_values.reshape(-1, num_shards, stack_embedding_dim),
        0,
        1,
    )

    # Rotate shards.
    setting_in_stack = table_spec.setting_in_stack
    rotated_values = jnp.roll(
        sharded_values, setting_in_stack.shard_rotation, axis=0
    )

    # Insert table into the stack.
    table_row = setting_in_stack.row_offset_in_shard
    stacked_table = stacked_table.at[
        :, table_row : (table_row + sharded_padded_vocabulary_size), :
    ].set(rotated_values)

    return stacked_table


def stack_and_shard_tables(
    table_specs: Nested[TableSpec],
    tables: Nested[ArrayLike],
    num_shards: int,
    pad_value: jnp.float32 = jnp.nan,
) -> dict[str, Nested[jax.Array]]:
    """Stacks and shards tables for use in sparsecore lookups.

    Args:
        table_specs: Nested collection of unstacked table specifications.
        tables: Table values corresponding to the table_specs.
        num_shards: Number of shards in the table (typically
            `global_device_count * num_sc_per_device`).
        pad_value: Value to use for padding.

    Returns:
        A mapping of stacked table names to stacked table values.
    """

    # Gather stacked table information.
    stacked_table_map: dict[
        str,
        tuple[StackedTableSpec, list[TableSpec]],
    ] = {}

    def collect_stacked_tables(table_spec: TableSpec) -> None:
        stacked_table_spec = _get_stacked_table_spec(table_spec, num_shards)
        stacked_table_name = stacked_table_spec.stack_name
        if stacked_table_name not in stacked_table_map:
            stacked_table_map[stacked_table_name] = (stacked_table_spec, [])
        stacked_table_map[stacked_table_name][1].append(table_spec)

    _ = jax.tree.map(collect_stacked_tables, table_specs)

    table_map: dict[str, Nested[jax.Array]] = {}

    def collect_tables(table_spec: TableSpec, table: Nested[jax.Array]) -> None:
        table_map[table_spec.name] = table

    _ = jax.tree.map(collect_tables, table_specs, tables)

    stacked_tables: dict[str, Nested[jax.Array]] = {}
    for (
        stacked_table_spec,
        table_specs,
    ) in stacked_table_map.values():
        stack_vocab_size = stacked_table_spec.stack_vocab_size
        sharded_vocab_size = stack_vocab_size // num_shards
        stack_embedding_dim = stacked_table_spec.stack_embedding_dim

        # Allocate initial buffer.  The stacked table will be divided among
        # shards by splitting the vocabulary dimension:
        #   [ v, e ] -> [s, v/s, e]
        stacked_table_tree = jax.tree.map(
            lambda _: jnp.zeros(
                # pylint: disable-next=cell-var-from-loop, used only in loop body.
                shape=(num_shards, sharded_vocab_size, stack_embedding_dim),
                dtype=jnp.float32,
            ),
            table_map[table_specs[0].name],
        )

        for table_spec in table_specs:
            table_tree = table_map[table_spec.name]
            stacked_table_tree = jax.tree.map(
                lambda stacked_table, table: _stack_and_shard_table(
                    # pylint: disable-next=cell-var-from-loop, used only in loop body.
                    stacked_table,
                    # pylint: disable-next=cell-var-from-loop, used only in loop body.
                    table_spec,
                    table,
                    num_shards,
                    pad_value,
                ),
                stacked_table_tree,
                table_tree,
            )

        stacked_tables[stacked_table_spec.stack_name] = stacked_table_tree

    return stacked_tables


def _unshard_and_unstack_table(
    table_spec: TableSpec,
    stacked_table_tree: Nested[jax.Array],
    num_shards: int,
) -> Nested[jax.Array]:
    """Unshards and unstacks a single table."""
    vocabulary_size = table_spec.vocabulary_size
    embedding_dim = table_spec.embedding_dim

    def _unshard_and_unstack_single_table(
        table_spec: TableSpec, stacked_table: jax.Array
    ) -> jax.Array:
        stack_embedding_dim = stacked_table.shape[-1]

        # Maybe re-shape in case it was flattened.
        stacked_table = stacked_table.reshape(
            num_shards, -1, stack_embedding_dim
        )
        sharded_vocabulary_size = (
            _round_up_to_multiple(vocabulary_size, 8 * num_shards) // num_shards
        )

        # Extract padded values from the stacked table.
        setting_in_stack = table_spec.setting_in_stack
        row = setting_in_stack.row_offset_in_shard
        padded_values = stacked_table[
            :, row : (row + sharded_vocabulary_size), :
        ]

        # Un-rotate shards.
        padded_values = jnp.roll(
            padded_values, -setting_in_stack.shard_rotation, axis=0
        )

        # Un-mod-shard.
        padded_values = jnp.swapaxes(padded_values, 0, 1).reshape(
            -1, stack_embedding_dim
        )

        # Un-pad.
        return padded_values[:vocabulary_size, :embedding_dim]

    output: Nested[jax.Array] = jax.tree.map(
        lambda stacked_table: _unshard_and_unstack_single_table(
            table_spec, stacked_table
        ),
        stacked_table_tree,
    )
    return output


def unshard_and_unstack_tables(
    table_specs: Nested[TableSpec],
    stacked_tables: Mapping[str, Nested[jax.Array]],
    num_shards: int,
) -> Nested[jax.Array]:
    """Unshards and unstacks a collection of tables.

    Args:
        table_specs: Nested collection of unstacked table specifications.
        stacked_tables: Mapping of stacked table names to stacked table values.
        num_shards: Number of shards in the table (typically
            `global_device_count * num_sc_per_device`).

    Returns:
        A mapping of table names to unstacked table values.
    """
    output: Nested[jax.Array] = jax.tree.map(
        lambda table_spec: _unshard_and_unstack_table(
            table_spec,
            stacked_tables[
                _get_stacked_table_spec(table_spec, num_shards=1).stack_name
            ],
            num_shards,
        ),
        table_specs,
    )
    return output


def get_table_specs(feature_specs: Nested[FeatureSpec]) -> dict[str, TableSpec]:
    table_spec_map: dict[str, TableSpec] = {}
    flat_feature_specs, _ = jax.tree.flatten(feature_specs)
    for feature_spec in flat_feature_specs:
        table_spec = feature_spec.table_spec
        table_spec_map[table_spec.name] = table_spec
    return table_spec_map


def get_table_stacks(
    table_specs: Nested[TableSpec],
) -> dict[str, list[TableSpec]]:
    """Extracts lists of tables that are stacked together.

    Args:
        table_specs: Nested collection of table specifications.

    Returns:
        A mapping of stacked table names to lists of table specifications for
        each stack.
    """
    stacked_table_specs: dict[str, list[TableSpec]] = collections.defaultdict(
        list
    )
    flat_table_specs, _ = jax.tree.flatten(table_specs)
    for table_spec in flat_table_specs:
        table_spec = typing.cast(TableSpec, table_spec)
        stacked_table_spec = table_spec.stacked_table_spec
        if stacked_table_spec is not None:
            stacked_table_specs[stacked_table_spec.stack_name].append(
                table_spec
            )
        else:
            stacked_table_specs[table_spec.name].append(table_spec)

    return stacked_table_specs


def get_stacked_table_stats(
    feature_specs: Nested[FeatureSpec],
) -> dict[str, InputStatsPerTable]:
    """Extracts the stacked-table input statistics from the feature specs.

    Args:
        feature_specs: Feature specs from which to extracts the statistics.

    Returns:
        A mapping of stacked table names to input statistics per table.
    """
    stacked_table_specs: dict[str, StackedTableSpec] = {}
    for feature_spec in jax.tree.flatten(feature_specs)[0]:
        feature_spec = typing.cast(FeatureSpec, feature_spec)
        stacked_table_spec = typing.cast(
            StackedTableSpec, feature_spec.table_spec.stacked_table_spec
        )
        stacked_table_specs[stacked_table_spec.stack_name] = stacked_table_spec

    stats: dict[str, InputStatsPerTable] = {}
    for stacked_table_spec in stacked_table_specs.values():
        buffer_size = stacked_table_spec.suggested_coo_buffer_size_per_device
        buffer_size = buffer_size or 0
        stats[stacked_table_spec.stack_name] = InputStatsPerTable(
            max_ids_per_partition=stacked_table_spec.max_ids_per_partition,
            max_unique_ids_per_partition=stacked_table_spec.max_unique_ids_per_partition,
            required_buffer_size_per_device=buffer_size,
        )

    return stats


def update_stacked_table_stats(
    feature_specs: Nested[FeatureSpec],
    stats: Mapping[str, InputStatsPerTable],
) -> None:
    """Updates stacked-table input properties in the supplied feature specs.

    Args:
        feature_specs: Feature specs to update in-place.
        stats: Per-stacked-table input statistics.
    """
    # Collect table specs and stacked table specs.
    table_specs: dict[str, TableSpec] = {}
    for feature_spec in jax.tree.flatten(feature_specs)[0]:
        feature_spec = typing.cast(FeatureSpec, feature_spec)
        table_specs[feature_spec.table_spec.name] = feature_spec.table_spec

    stacked_table_specs: dict[str, StackedTableSpec] = {}
    for table_spec in table_specs.values():
        stacked_table_spec = typing.cast(
            StackedTableSpec, table_spec.stacked_table_spec
        )
        stacked_table_specs[stacked_table_spec.stack_name] = stacked_table_spec

    # Replace fields in the stacked_table_specs.
    stack_names = stacked_table_specs.keys()
    for stack_name in stack_names:
        stack_stats = stats[stack_name]
        stacked_table_spec = stacked_table_specs[stack_name]
        buffer_size = stack_stats.required_buffer_size_per_device or None
        stacked_table_specs[stack_name] = dataclasses.replace(
            stacked_table_spec,
            max_ids_per_partition=stack_stats.max_ids_per_partition,
            max_unique_ids_per_partition=stack_stats.max_unique_ids_per_partition,
            suggested_coo_buffer_size_per_device=buffer_size,
        )

    # Insert new stacked tables into tables.
    for table_spec in table_specs.values():
        stacked_table_spec = typing.cast(
            StackedTableSpec, table_spec.stacked_table_spec
        )
        table_spec.stacked_table_spec = stacked_table_specs[
            stacked_table_spec.stack_name
        ]


def convert_to_numpy(
    ragged_or_dense: np.ndarray[Any, Any] | Sequence[Sequence[Any]] | Any,
    dtype: Any,
) -> np.ndarray[Any, Any]:
    """Converts a ragged or dense list of inputs to a ragged/dense numpy array.

    The output is adjusted to be 2D.

    Args:
        ragged_or_dense: Input that is either already a numpy array, or nested
            sequence.
        dtype: Numpy dtype of output array.

    Returns:
        Corresponding numpy array.
    """
    if hasattr(ragged_or_dense, "numpy"):
        # Support tf.RaggedTensor and other TF input dtypes.
        if callable(getattr(ragged_or_dense, "numpy")):
            ragged_or_dense = ragged_or_dense.numpy()

    if isinstance(ragged_or_dense, jax.Array):
        ragged_or_dense = np.asarray(ragged_or_dense)

    if isinstance(ragged_or_dense, np.ndarray):
        # Convert 1D to 2D.
        if ragged_or_dense.dtype != np.ndarray and ragged_or_dense.ndim == 1:
            return ragged_or_dense.reshape(-1, 1).astype(dtype)

        # If dense, return converted dense type.
        if ragged_or_dense.dtype != np.ndarray:
            return ragged_or_dense.astype(dtype)

        # Ragged numpy array.
        return ragged_or_dense

    # Handle 1D sequence input.
    if not isinstance(ragged_or_dense[0], collections.abc.Sequence):
        return np.asarray(ragged_or_dense, dtype=dtype).reshape(-1, 1)

    # Assemble elements into an nd-array.
    counts = [len(vals) for vals in ragged_or_dense]
    if all([count == counts[0] for count in counts]):
        # Dense input.
        return np.asarray(ragged_or_dense, dtype=dtype)
    else:
        # Ragged input, convert to ragged numpy arrays.
        return np.array(
            [np.array(row, dtype=dtype) for row in ragged_or_dense],
            dtype=np.ndarray,
        )


def ones_like(
    ragged_or_dense: np.ndarray[Any, Any], dtype: Any = None
) -> np.ndarray[Any, Any]:
    """Creates an array of ones the same as as the input.

    This differs from traditional numpy in that a ragged input will lead to
    a resulting ragged array of ones, whereas np.ones_like(...) will instead
    only consider the outer array and return a 1D dense array of ones.

    Args:
        ragged_or_dense: The ragged or dense input whose shape and data-type
              define these same attributes of the returned array.
        dtype: The data-type of the returned array.

    Returns:
        An array of ones with the same shape as the input, and specified data
        type.
    """
    dtype = dtype or ragged_or_dense.dtype
    if ragged_or_dense.dtype == np.ndarray:
        # Ragged.
        return np.array(
            [np.ones_like(row, dtype=dtype) for row in ragged_or_dense],
            dtype=np.ndarray,
        )
    else:
        # Dense.
        return np.ones_like(ragged_or_dense, dtype=dtype)


def create_feature_samples(
    feature_structure: Nested[T],
    feature_ids: Nested[ArrayLike | Sequence[int] | Sequence[Sequence[int]]],
    feature_weights: None
    | (Nested[ArrayLike | Sequence[float] | Sequence[Sequence[float]]]),
) -> Nested[FeatureSamples]:
    """Constructs a collection of sample tuples from provided IDs and weights.

    Args:
        feature_structure: The nested structure of the inputs (typically
            `FeatureSpec`s).
        feature_ids: The feature IDs to use for the samples.
        feature_weights: The feature weights to use for the samples.  Defaults
            to ones if not provided.

    Returns:
        A nested collection of `FeatureSamples` corresponding to the input IDs
        and weights, for use in embedding lookups.
    """
    # Create numpy arrays from inputs.
    feature_ids = jax.tree.map(
        lambda _, ids: convert_to_numpy(ids, np.int32),
        feature_structure,
        feature_ids,
    )

    if feature_weights is None:
        # Make ragged or dense ones_like.
        feature_weights = jax.tree.map(
            lambda _, ids: ones_like(ids, np.float32),
            feature_structure,
            feature_ids,
        )
    else:
        feature_weights = jax.tree.map(
            lambda _, wgts: convert_to_numpy(wgts, np.float32),
            feature_structure,
            feature_weights,
        )

    # Assemble.
    def _create_feature_samples(
        sample_ids: np.ndarray[Any, Any],
        sample_weights: np.ndarray[Any, Any],
    ) -> FeatureSamples:
        return FeatureSamples(sample_ids, sample_weights)

    output: Nested[FeatureSamples] = jax.tree.map(
        lambda _, sample_ids, sample_weights: _create_feature_samples(
            sample_ids, sample_weights
        ),
        feature_structure,
        feature_ids,
        feature_weights,
    )
    return output


def stack_and_shard_samples(
    feature_specs: Nested[FeatureSpec],
    feature_samples: Nested[FeatureSamples],
    local_device_count: int,
    global_device_count: int,
    num_sc_per_device: int,
    static_buffer_size: int | Mapping[str, int] | None = None,
) -> tuple[dict[str, ShardedCooMatrix], dict[str, InputStatsPerTable]]:
    """Prepares input samples for use in embedding lookups.

    Args:
        feature_specs: Nested collection of feature specifications.
        feature_samples: Nested collection of feature samples.
        local_device_count: Number of local JAX devices.
        global_device_count: Number of global JAX devices.
        num_sc_per_device: Number of sparsecores per device.
        static_buffer_size: The static buffer size to use for the samples.
          Defaults to None, in which case an upper-bound for the buffer size
          will be automatically determined.

    Returns:
        The preprocessed inputs, and statistics useful for updating FeatureSpecs
        based on the provided input data.
    """
    del static_buffer_size  # Currently ignored.
    flat_feature_specs, _ = jax.tree.flatten(feature_specs)

    feature_tokens = []
    feature_weights = []

    def collect_tokens_and_weights(
        feature_spec: FeatureSpec, samples: FeatureSamples
    ) -> None:
        del feature_spec
        feature_tokens.append(samples.tokens)
        feature_weights.append(samples.weights)

    jax.tree.map(collect_tokens_and_weights, feature_specs, feature_samples)

    preprocessed_inputs, stats = embedding.preprocess_sparse_dense_matmul_input(
        feature_tokens,
        feature_weights,
        flat_feature_specs,
        local_device_count=local_device_count,
        global_device_count=global_device_count,
        num_sc_per_device=num_sc_per_device,
        sharding_strategy="MOD",
        has_leading_dimension=False,
        allow_id_dropping=True,
    )

    out: dict[str, ShardedCooMatrix] = {}
    out_stats: dict[str, InputStatsPerTable] = {}
    tables_names = preprocessed_inputs.lhs_row_pointers.keys()
    for table_name in tables_names:
        shard_ends = preprocessed_inputs.lhs_row_pointers[table_name]
        shard_starts = np.concatenate(
            [np.asarray([0]), _round_up_to_multiple(shard_ends[:-1], 8)]
        )
        out[table_name] = ShardedCooMatrix(
            shard_starts=shard_starts,
            shard_ends=shard_ends,
            col_ids=preprocessed_inputs.lhs_embedding_ids[table_name],
            row_ids=preprocessed_inputs.lhs_sample_ids[table_name],
            values=preprocessed_inputs.lhs_gains[table_name],
        )
        out_stats[table_name] = InputStatsPerTable(
            max_ids_per_partition=np.max(
                stats.max_ids_per_partition[table_name]
            ),
            max_unique_ids_per_partition=np.max(
                stats.max_unique_ids_per_partition[table_name]
            ),
            required_buffer_size_per_device=np.max(
                stats.required_buffer_size_per_sc[table_name]
            )
            * num_sc_per_device,
        )

    return out, out_stats
