import traceback
from collections.abc import Callable
from pprint import pformat
from typing import Optional, Union

import torch
import torch.nn as nn
from tqdm import tqdm
from typing_extensions import Self

import torchlinops.config as config
from torchlinops.utils import INDENT

from ..nameddim import NamedShape as NS, Shape
from .identity import ShapeSpec
from .namedlinop import NamedLinop
from .split import split_linop

__all__ = ["Batch"]


class Batch(NamedLinop):
    """Wrapper that applies a linop in batches over named dimensions.

    Splits the input along specified batch dimensions, applies the wrapped
    linop to each chunk, and concatenates the results.
    """

    def __init__(
        self,
        linop: NamedLinop,
        input_device: torch.device,
        output_device: torch.device,
        input_dtype: Union[str, torch.dtype],
        output_dtype: Union[str, torch.dtype],
        input_shape: Optional[Shape] = None,
        output_shape: Optional[Shape] = None,
        pbar: bool = False,
        name: Optional[str] = None,
        post_batch_hook: Optional[Callable] = None,
        **batch_sizes,
    ):
        """
        post_batch_hook : Callable, optional
            Function that takes in the newly-created batch object and does stuff

        Deprecated, will be removed in 0.6.0.
        """
        # TODO: Should batch even have a shape???
        super().__init__(NS(linop.ishape, linop.oshape), name=name)

        self.linop = linop
        if input_shape is not None:
            self.linop = self.linop @ ShapeSpec(input_shape)
        if output_shape is not None:
            self.linop = ShapeSpec(output_shape) @ self.linop
        self.input_device = input_device
        self.output_device = output_device
        self.input_dtype = input_dtype
        self.output_dtype = output_dtype
        self.pbar = pbar
        # self.name = name if name is not None else ""
        self.batch_sizes = batch_sizes
        self.post_batch_hook = post_batch_hook
        self.sizes = {dim: self.linop.size(dim) for dim in self.linop.dims}
        self.setup_batching()

    def setup_batching(self, hook: Optional[Callable] = None):
        _linops, _input_batches, _output_batches = split_linop(
            self.linop,
            self.batch_sizes,
        )
        self._linops = nn.ModuleList(_linops.flatten().tolist())
        self._input_batches = _input_batches.flatten().tolist()
        self._output_batches = _output_batches.flatten().tolist()
        self._shape = NS(self.linop.ishape, self.linop.oshape)
        super().reset_adjoint_and_normal()
        if self.post_batch_hook is not None:
            self.post_batch_hook(self)

    def to(self, device: torch.device | str) -> Self:
        self.input_device = device
        self.output_device = device
        # self.linop.to(device)
        # Avoid copying the _linops by doing super() first
        super().to(device)
        self.setup_batching()
        return self

    def forward(self, x: torch.Tensor):
        # Complete the size specifications
        for dim, total in zip(self.ishape, x.shape):
            self.sizes[dim] = total

        y = torch.zeros(
            tuple(self.sizes[dim] for dim in self.oshape),
            dtype=self.output_dtype,
            device=self.output_device,
        )
        for linop, in_batch, out_batch in tqdm(
            zip(self._linops, self._input_batches, self._output_batches),
            total=len(self._linops),
            desc=f"Batch({self.name}: {self.batch_sizes})",
            disable=(not self.pbar),
        ):
            try:
                xbatch = x[tuple(in_batch)]
                ybatch = linop(xbatch)
                y[tuple(out_batch)] += ybatch
            except RuntimeError:
                print(
                    f"linop: {linop}, in_batch: {in_batch}, out_batch: {out_batch}, self.batch_sizes: {self.batch_sizes}"
                )
                raise
        return y

    @property
    def H(self):
        if config.cache_adjoint_normal:
            config._warn_if_caching_enabled()
            if self._adjoint is None:
                try:
                    _adjoint = self.adjoint()
                    _adjoint._adjoint = [self]
                    self._adjoint = [_adjoint]
                except AttributeError as e:
                    traceback.print_exc()
                    raise e
            return self._adjoint[0]
        return self.adjoint()

    def adjoint(self):
        batch_sizes = {str(k): v for k, v in self.batch_sizes.items()}
        adj = type(self)(
            linop=self.linop.H,
            input_device=self.output_device,
            output_device=self.input_device,
            input_dtype=self.output_dtype,
            output_dtype=self.input_dtype,
            name=self.name + ".H",
            pbar=self.pbar,
            post_batch_hook=self.post_batch_hook,
            **batch_sizes,
        )
        return adj

    @property
    def N(self):
        if config.cache_adjoint_normal:
            config._warn_if_caching_enabled()
            if self._normal is None:
                try:
                    _normal = self.normal()
                    self._normal = [_normal]
                except AttributeError as e:
                    traceback.print_exc()
                    raise e
            return self._normal[0]
        return self.normal()

    def normal(self, inner=None):
        normal_linop = self.linop.N
        # Collect shape updates from computing the normal
        shape_updates = getattr(normal_linop, "_shape_updates", {})
        new_batch_sizes = self.batch_sizes.copy()
        for d, nd in shape_updates.items():
            if d in self.batch_sizes:
                new_batch_sizes[shape_updates[d]] = self.batch_sizes[d]
            elif str(d) in self.batch_sizes:
                new_batch_sizes[str(shape_updates[d])] = self.batch_sizes[str(d)]
        batch_size_kwargs = {str(k): v for k, v in new_batch_sizes.items()}
        normal = type(self)(
            linop=self.linop.N,
            input_device=self.input_device,
            output_device=self.input_device,
            input_dtype=self.input_dtype,
            output_dtype=self.input_dtype,
            name=self.name + ".N",
            pbar=self.pbar,
            post_batch_hook=self.post_batch_hook,
            **batch_size_kwargs,
        )

        return normal

    # @staticmethod
    # def fn(self, x, /, data):
    #     """TODO: Functional interface
    #     Specify data as a tuple of data entries, one for each linop in linops"""
    #     raise NotImplementedError(f"Batched functional interface not available yet.")
    #     sizes = {}
    #     for dim in self.linop.dims:
    #         sizes[dim] = self.linop.size_fn(dim, data)
    #     for dim, total in zip(self.ishape, x.shape):
    #         sizes[dim] = total
    #     batch_iterators = self._make_batch_iterators(sizes, self.batch_sizes)
    #     ishapes = [linop.ishape for linop in self.linop.flatten()]
    #     oshapes = [linop.oshape for linop in self.linop.flatten()]

    #     y = torch.zeros(
    #         tuple(sizes[dim] for dim in self.oshape),
    #         dtype=self.output_dtype,
    #         device=self.output_device,
    #     )
    #     for tile in tqdm(
    #         dict_product(batch_iterators),
    #         desc=f"Batch({self.batch_sizes})",
    #         disable=(not self.pbar),
    #     ):
    #         ibatches = [
    #             [tile.get(dim, slice(None)) for dim in ishape] for ishape in ishapes
    #         ]
    #         obatches = [
    #             [tile.get(dim, slice(None)) for dim in oshape] for oshape in oshapes
    #         ]
    #         split_data = self.linop.split_fn(*ibatches, *obatches, *data)
    #         xbatch = x[ibatches[0]].to(self.input_device)
    #         ybatch = self.linop.fn(xbatch, split_data)
    #         y[obatches[-0]] += ybatch
    #     return y

    @staticmethod
    def adj_fn(self, x, /, data):
        raise NotImplementedError("Batched linop has no adjoint (yet).")

    def size(self, dim):
        return self.linop.size(dim)

    def __repr__(self):
        """Helps prevent recursion error caused by .H and .N"""
        output = ""
        output += INDENT.indent(self.name + self._suffix) + "(\n"
        with INDENT:
            output += repr(self.linop)
            output += ", "
            output += repr(pformat(self.batch_sizes)).strip("'")
            output += "\n"
        output += INDENT.indent(")")
        return output
