from dataclasses import dataclass, field
from math import ceil
from typing import Optional
from warnings import warn

import numpy as np
import torch
from torch.cuda import Stream

from torchlinops.utils import (
    ModuleMemoryMap,
    RepeatedEvent,
    batch_iterator,
    dict_product,
    default_to,
)

from .add import Add
from .chain import Chain
from .concat import Concat
from .device import ToDevice
from ..nameddim import NamedDimension as ND
from .namedlinop import NamedLinop

__all__ = ["split_linop", "create_batched_linop", "BatchSpec"]

Batch = tuple[int, slice]
# Represents a single batch at index 0 over the full extent
# Could convert to a full class
DEFAULT_BATCH = (0, slice(None))
Tile = dict[ND | str, Batch]


@dataclass
class BatchSpec:
    """Specification for splitting and distributing a linop across devices.

    Parameters
    ----------
    batch_sizes : dict[ND | str, int]
        Mapping from dimension names to chunk sizes for tiling.
    device_matrix : np.ndarray or list, optional
        Array of ``torch.device`` objects specifying target devices for each
        tile. Broadcast to match the tile grid shape.
    base_device : torch.device, optional
        The device where input/output data resides. Default is CPU.
    """

    batch_sizes: dict[ND | str, int]
    device_matrix: np.ndarray | None = None
    base_device: torch.device | None = None

    def __post_init__(self):
        if not isinstance(self.batch_sizes, dict):
            warn(
                f"Got {self.batch_sizes} of type {type(self.batch_sizes).__name__} for batch_sizes instead of dict."
            )
        # Ensure ndarray
        if isinstance(self.device_matrix, list | tuple):
            self.device_matrix = np.array(self.device_matrix)

    def broadcast_device_matrix(self, linop):
        # Compute the number of tiles along each batched axis/dimension
        batch_dims = list(self.batch_sizes.keys())
        sizes = {dim: linop.size(dim) for dim in linop.dims}
        tiled_shape = tuple(
            ceil(sizes[dim] / self.batch_sizes[dim]) for dim in batch_dims
        )

        # Broadcast device_matrix over requested tiles.
        # Each tile should receive a single device.
        device_matrix = fuzzy_broadcast_to(self.device_matrix, tiled_shape)
        return device_matrix


def create_batched_linop(
    linop,
    batch_specs: BatchSpec | list[BatchSpec],
    default_device: torch.device = None,
    _mmap=None,
):
    """Split and distribute a linop across devices according to batch specs.

    Recursively processes a list of ``BatchSpec`` objects: the first spec
    splits the linop into tiles, optionally places each tile on a target
    device, then passes remaining specs to each tile recursively. Tiles are
    reassembled via ``Concat`` (for partitioned dimensions) or ``Add`` (for
    reduced dimensions).

    Parameters
    ----------
    linop : NamedLinop
        The operator to split and distribute.
    batch_specs : BatchSpec or list[BatchSpec]
        One or more batch specifications to apply (processed in order).
    _mmap : ModuleMemoryMap, optional
        Internal memory map for efficient device transfers. Created
        automatically on the first call. Probably don't set this manually.
    _default_device : torch.device, optional
        The default device to use if no device info is provided in the batch spec.

    Returns
    -------
    NamedLinop
        A composite linop (tree of ``Concat``/``Add``/``ToDevice`` operators)
        that is functionally equivalent to the original but distributed
        according to the batch specs.
    """
    if default_device is None:
        default_device = torch.device("cpu")
    if isinstance(batch_specs, BatchSpec):
        # Ensure list
        batch_specs = [batch_specs]
    if _mmap is None:
        _mmap = ModuleMemoryMap()
        _mmap.register_module(linop)
    if len(batch_specs) == 0:
        # Recursive ending
        return linop
    batch_spec = batch_specs[0]
    # Set defaults
    batch_spec.base_device = default_to(default_device, batch_spec.base_device)
    batch_spec.device_matrix = default_to(
        np.array([default_device]), batch_spec.device_matrix
    )

    # Split linop into tiles and broadcast device spec to the tile array.
    linops, ibatches, obatches = split_linop(linop, batch_spec.batch_sizes)
    device_matrix = batch_spec.broadcast_device_matrix(linop)
    if device_matrix.shape != linops.shape:
        raise ValueError(
            f"device_matrix and linops should have same shape after broadcasting, but got device_matrix: {device_matrix.shape} and linops: {linops.shape}"
        )

    # Create event to trigger all tiles in the linop.
    source_device = batch_spec.base_device
    # Allocate output
    for idx in np.ndindex(linops.shape):
        linop, target_device = linops[idx], device_matrix[idx]

        # Recursive call to batch the tile
        tiled_linop = create_batched_linop(
            linop, batch_specs[1:], default_device=target_device, _mmap=_mmap
        )

        # Move linop to device
        tiled_linop = _mmap.memory_aware_to(tiled_linop, target_device)

        # Wrap with device movement linops
        if source_device != target_device:
            tiled_linop = Chain(
                ToDevice(
                    source_device,
                    target_device,
                    ioshape=tiled_linop.ishape,
                ),
                tiled_linop,
                ToDevice(
                    target_device,
                    source_device,
                    ioshape=tiled_linop.oshape,
                ),
            )

        # Overwrite entry in linops
        linops[idx] = tiled_linop

    for dim in reversed(batch_spec.batch_sizes):
        # Manual axis reduction because I made Concat and Add too nice
        flat_linops = linops.reshape(-1, linops.shape[-1])
        new_linops = np.empty(flat_linops.shape[0], dtype=object)
        for i, linop_arr in enumerate(flat_linops):
            linop = linop_arr[0]
            if dim in linop.ishape and dim in linop.oshape:
                new_linop = Concat(*linop_arr, idim=dim, odim=dim)
            elif dim not in linop.ishape and dim in linop.oshape:
                new_linop = Concat(*linop_arr, odim=dim)
            elif dim in linop.ishape and dim not in linop.oshape:
                new_linop = Concat(*linop_arr, idim=dim)
            else:
                new_linop = Add(*linop_arr)
            new_linops[i] = new_linop
        linops = new_linops.reshape(linops.shape[:-1])
    linop = linops.item()
    return linop


def split_linop(linop: NamedLinop, batch_sizes: dict[ND | str, int]):
    """Split a linop into an nd-array of sub-linops according to batch sizes.

    Parameters
    ----------
    linop : NamedLinop
        The linop to be split.
    batch_sizes : dict[ND | str, int]
        Dictionary mapping dimension names to chunk sizes.

    Returns
    -------
    linops : np.ndarray
        Array of sub-linops with shape determined by the number of tiles
        per dimension.
    input_batches : np.ndarray
        Corresponding input slices for each tile.
    output_batches : np.ndarray
        Corresponding output slices for each tile.
    """
    # Precompute sizes and shapes
    batch_sizes = {ND.infer(k): v for k, v in batch_sizes.items()}
    sizes = {dim: linop.size(dim) for dim in linop.dims}

    # Make tiles. Each tile is a dictionary mapping a dimension to an integer
    # index of the tile and a slice over that dimension.
    batch_iterators = make_batch_iterators(sizes, batch_sizes)
    tiles: list[dict[ND, Batch]] = list(dict_product(batch_iterators))

    # Allocate outputs
    batch_dims = list(batch_sizes.keys())
    tiled_shape = tuple(ceil(sizes[dim] / batch_sizes[dim]) for dim in batch_dims)
    linops = np.ndarray(tiled_shape, dtype=object)
    input_batches = np.ndarray(tiled_shape, dtype=object)
    output_batches = np.ndarray(tiled_shape, dtype=object)

    for tile in tiles:
        idx = _tile_get_idx(tile, batch_dims)
        linop_tile = _split_linop_with_tile(linop, tile)
        linop_flat = linop_tile.flatten()
        first_linop, last_linop = linop_flat[0], linop_flat[-1]
        linops[idx] = linop_tile
        input_batches[idx] = [
            tile.get(dim, DEFAULT_BATCH)[1] for dim in first_linop.ishape
        ]
        output_batches[idx] = [
            tile.get(dim, DEFAULT_BATCH)[1] for dim in last_linop.oshape
        ]
    return linops, input_batches, output_batches


def fuzzy_broadcast_to(arr: np.ndarray, target_shape):
    """Broadcast an array to a target shape, truncating or repeating if necessary.

    Examples
    --------
    >>> import numpy as np
    >>> arr = np.array([1, 2, 3])
    >>> fuzzy_broadcast_to(arr, (4,))
    array([1, 2, 3, 1])
    >>> fuzzy_broadcast_to(arr, (3, 1))
    array([[1, 2, 3],
           [1, 2, 3],
           [1, 2, 3]])
    >>> arr2 = np.array([1, 2])
    >>> fuzzy_broadcast_to(arr2, (3, 5))
    array([[1, 2, 1, 2, 1],
           [1, 2, 1, 2, 1],
           [1, 2, 1, 2, 1]])
    """
    # Ensure target shape and arr.shape have same length
    if len(target_shape) < arr.ndim:
        target_shape = (1,) * (arr.ndim - len(target_shape)) + tuple(target_shape)
    while len(target_shape) > arr.ndim:
        arr = np.expand_dims(arr, 0)

    if not is_broadcastable(arr.shape, target_shape):
        # Identify offending dimensions and repeat
        repeats = []
        for source_dim, target_dim in zip(arr.shape, target_shape):
            if source_dim == target_dim or source_dim == 1:
                repeats.append(1)
            elif source_dim < target_dim:
                repeats.append(int(ceil(target_dim / source_dim)))
            else:
                repeats.append(1)
        arr = tile_along_axes(arr, repeats)
        # Trim the excess
        slices = tuple(slice(0, dim) for dim in target_shape)
        arr = arr[slices]
    target_shape = np.broadcast_shapes(target_shape, arr.shape)
    return np.broadcast_to(arr, target_shape)


def _split_linop_with_tile(linop: NamedLinop, tile: Tile):
    """Split a linop according to batch specified in tile.

    Parameters
    ----------
    linop : NamedLinop
        The linop to split.
    tile : Tile
        Dictionary mapping dimension names to (index, slice) pairs.

    Returns
    -------
    NamedLinop
        The split sub-linop operating on the specified tile.
    """
    slice_map = {key: value[1] for key, value in tile.items()}
    linop_tile = linop.split(linop, slice_map)
    return linop_tile


def _tile_get_idx(tile: Tile, batch_dims) -> tuple[int]:
    """Get all indices from the tile"""
    return tuple(tile.get(dim, DEFAULT_BATCH)[0] for dim in batch_dims)


def make_batch_iterators(
    total_sizes: dict[str, int],
    batch_sizes: dict[str, int],
) -> Tile:
    """Build per-dimension lists of ``(index, slice)`` pairs for tiling.

    Parameters
    ----------
    total_sizes : dict[str, int]
        Total size of each dimension.
    batch_sizes : dict[str, int]
        Chunk size for dimensions that should be batched.

    Returns
    -------
    dict
        Maps each dimension to a list of ``(tile_index, slice)`` tuples.

    Notes
    -----
    If ``batch_sizes = {"D": 3}`` and the total size for ``D`` is 7::

        batch_iterators["D"] = [(0, slice(0,3)), (1, slice(3,6)), (2, slice(6,7))]

    Unbatched dimension ``E``::

        batch_iterators["E"] = [(0, slice(None))]
    """
    batch_iterators = {}
    for dim, total in total_sizes.items():
        batch_iterators[dim] = (
            [
                (i, slice(a, b))
                for i, (a, b) in enumerate(batch_iterator(total, batch_sizes[dim]))
            ]
            if dim in batch_sizes
            else [(0, slice(None))]
        )
    return batch_iterators


def flatten_recursive(nested_list, max_depth: Optional[int] = None):
    """Flatten a nested list, optionally to a maximum depth

    Examples
    -------
    >>> flatten_recursive([[1, 2], [3, 4], [[5, 6]]])
    [1, 2, 3, 4, 5, 6]

    # Setting max_depth = 1 will avoid flattening the last list completely
    >>> flatten_recursive([[1, 2], [3, 4], [[5, 6]]], max_depth=1)
    [1, 2, 3, 4, [5, 6]]

    """
    flat_list = []
    for item in nested_list:
        if isinstance(item, list):
            if max_depth is None:
                flat_list.extend(flatten_recursive(item))
            elif max_depth > 0:
                flat_list.extend(flatten_recursive(item, max_depth - 1))
            else:
                flat_list.append(item)
        else:
            flat_list.append(item)
    return flat_list


def is_broadcastable(a, b):
    try:
        np.broadcast_shapes(a, b)
        return True
    except ValueError:
        return False


def repeat_along_axes(arr, repeats):
    """Repeat a numpy array along its axes.

    Parameters
    ----------
    arr : np.ndarray
        The array to repeat.
    repeats : list or tuple of int
        One repeat count per axis. If shorter than ``arr.ndim``, remaining
        axes are not repeated. If longer, singleton dimensions are added.

    Returns
    -------
    np.ndarray
        The repeated array.
    """
    arr = np.asarray(arr)
    ndim = arr.ndim
    if len(repeats) > ndim:
        # Add singleton dimensions to match the number of repeats
        arr = arr.reshape(arr.shape + (1,) * (len(repeats) - ndim))
    elif len(repeats) < ndim:
        repeats = list(repeats) + [1] * (ndim - len(repeats))

    for axis, repeat in enumerate(repeats):
        if repeat > 1:
            arr = np.repeat(arr, repeat, axis=axis)
    return arr


def tile_along_axes(arr, tile_counts):
    """Tile a numpy array along its axes.

    Parameters
    ----------
    arr : np.ndarray
        The array to tile.
    tile_counts : list or tuple of int
        One tile count per axis. If shorter than ``arr.ndim``, missing axes
        default to 1. If longer, singleton dimensions are added.

    Returns
    -------
    np.ndarray
        The tiled array.
    """
    arr = np.asarray(arr)
    ndim = arr.ndim
    if len(tile_counts) < ndim:
        tile_counts = list(tile_counts) + [1] * (ndim - len(tile_counts))
    elif len(tile_counts) > ndim:
        arr = arr.reshape(arr.shape + (1,) * (len(tile_counts) - ndim))

    return np.tile(arr, tile_counts)


if __name__ == "__main__":
    import doctest

    doctest.testmod()
