from copy import deepcopy
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,
    default_to,
    dict_product,
)

from ..nameddim import NamedDimension as ND
from .add import Add
from .chain import Chain
from .concat import Concat
from .device import ToDevice
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 = deepcopy(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()
