import sys
import traceback
from collections.abc import Mapping
from typing import Optional

import torch
import torch.nn as nn

from torchlinops.utils import INDENT

from .nameddim import NS, isequal, NamedDimension as ND
from .namedlinop import NamedLinop


class Chain(NamedLinop):
    """A sequence or composition of linops"""

    def __init__(self, *linops, name: Optional[str] = None):
        """
        Parameters
        ----------
        *linops : list
            Linops in order of execution
            i.e. if `linops = [A, B, C]`, then mathematically, the linop in question is `CBA`

        """
        super().__init__(NS(linops[0].ishape, linops[-1].oshape), name=name)
        self.linops = nn.ModuleList(list(linops))
        self._check_inputs_outputs()

    def _check_inputs_outputs(self):
        curr_shape = self.ishape
        for i, linop in enumerate(self.linops):
            if not isequal(linop.ishape, curr_shape):
                raise ValueError(
                    f"Mismatched shape: expected {linop.ishape}, got {curr_shape} at input to {linop}. Full stack: {self}, index {i}"
                )
            curr_shape = linop.oshape

    def forward(self, x):
        for linop in self.linops:
            x = linop(x)
        return x

    @staticmethod
    def fn(chain, x: torch.Tensor, /, data_list):
        assert len(chain.linops) == len(data_list), (
            f"Length {len(data_list)} data_list does not match length {len(chain.linops)} chain linop"
        )
        for linop, data in zip(chain.linops, data_list):
            x = linop.fn(x, *data)
        return x

    @staticmethod
    def adj_fn(chain, x: torch.Tensor, /, data_list):
        assert len(chain.linops) == len(data_list), (
            f"Length {len(data_list)} data_list does not match length {len(chain.linops)} chain adjoint linop"
        )
        for linop, data in zip(reversed(chain.linops), reversed(data_list)):
            x = linop.adj_fn(x, data)
        return x

    @staticmethod
    def normal_fn(chain, x: torch.Tensor, /, data_list):
        # fn does the reversing so it's unnecessary to do it here
        # If the normal hasn't been explicitly formed with`.N`, do things the naive way
        return chain.adj_fn(chain.fn(x, data_list), data_list)

    def split_forward(self, ibatches, obatches):
        """ibatches, obatches specified according to the shape of the
        forward op
        """
        linops = [
            linop.split_forward(ibatch, obatch)
            for linop, ibatch, obatch in zip(self.linops, ibatches, obatches)
        ]
        return type(self)(*linops, name=self._name)

    def size(self, dim):
        out = None
        for linop in self.linops:
            tmp = linop.size(dim)
            if tmp is not None:
                if out is None:
                    out = tmp
                elif out != tmp:
                    raise ValueError(
                        f"Conflicting linop sizes found: {out} and {tmp} for dim {dim} in linop {linop} out of all linops {self.linops}"
                    )
        return out

    def size_fn(self, dim, data):
        for linop, data in zip(self.linops, data):
            out = linop.size_fn(dim, data)
            if out is not None:
                return out
        return None

    @property
    def dims(self):
        """Get the dims that appear anywhere in this linop chain."""
        return set().union(*[linop.dims for linop in self.linops])

    def adjoint(self):
        linops = list(linop.adjoint() for linop in reversed(self.linops))
        return type(self)(*linops, name=self._name)

    def normal(self, inner=None):
        for linop in reversed(self.linops):
            inner = linop.normal(inner)
        return inner

    @staticmethod
    def split(chain, tile: Mapping[ND | str, slice]):
        """Split a linop into sub-linops.

        Parameters
        ----------
        chain : Chain
            The chain linop to split.
        tile : Mapping[ND | str, slice]
            Dictionary specifying how to slice the linop dimensions
        """
        ibatches = [
            [tile.get(dim, slice(None)) for dim in linop.ishape]
            for linop in chain.linops
        ]
        obatches = [
            [tile.get(dim, slice(None)) for dim in linop.oshape]
            for linop in chain.linops
        ]
        return chain.split_forward(ibatches, obatches)

    @staticmethod
    def adj_split(chain, tile: Mapping[ND | str, slice]):
        """Split an adjoint linop into sub-linops.

        Parameters
        ----------
        chain : Chain
            The chain linop to split.
        tile : Mapping[ND | str, slice]
            Dictionary specifying how to slice the linop dimensions
        """
        ibatches = [
            [tile.get(dim, slice(None)) for dim in linop.ishape]
            for linop in chain.linops
        ]
        obatches = [
            [tile.get(dim, slice(None)) for dim in linop.oshape]
            for linop in chain.linops
        ]
        return chain.H.split_forward(obatches, ibatches).H

    # DEPRECATED/TODO: Remove
    # @staticmethod
    # def split_fn(chain, *iobatchesdata):
    #     """Return split versions of the data that can be passed
    #     into fn and adj_fn to produce split versions
    #     """
    #     ibatches = iobatchesdata[: len(iobatchesdata) // 3]
    #     obatches = iobatchesdata[len(iobatchesdata) // 3 : len(iobatchesdata) * 2 // 3]
    #     data = iobatchesdata[len(iobatchesdata) * 2 // 3 :]
    #     return chain.split_forward_fn(ibatches, obatches, data)

    # @staticmethod
    # def adj_split_fn(chain, *iobatchesdata):
    #     ibatches = iobatchesdata[: len(iobatchesdata) // 3]
    #     obatches = iobatchesdata[len(iobatchesdata) // 3 : len(iobatchesdata) * 2 // 3]
    #     data = iobatchesdata[len(iobatchesdata) * 2 // 3]
    #     return chain.split_forward_fn(obatches, ibatches, data)

    @property
    def shape(self):
        return NS(self.linops[0].ishape, self.linops[-1].oshape)

    @shape.setter
    def shape(self, val):
        self.ishape = val.ishape
        self.oshape = val.oshape

    @property
    def ishape(self):
        return self.linops[0].ishape

    @ishape.setter
    def ishape(self, val):
        self.linops[0].ishape = val

    @property
    def oshape(self):
        return self.linops[-1].oshape

    @oshape.setter
    def oshape(self, val):
        self.linops[-1].oshape = val

    def flatten(self):
        return list(self.linops)

    def __getitem__(self, idx):
        linops = self.linops[idx]
        if isinstance(linops, NamedLinop):
            return linops
        return Chain(*linops, name=self._name)

    def __len__(self):
        return len(self.linops)

    def __repr__(self):
        output = ""
        output += INDENT.indent(self.repr_name + "(\n")
        with INDENT:
            for linop in self.linops:
                output += repr(linop) + "\n"
        output += INDENT.indent(")")
        return output
