from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Generic, Sequence, TypeVar

import numpy as np
import scipy.sparse
import sparse
import torch

from ._tensor_schema import TensorSchema
from .types import TensorKind

try:
    nested_tensor = torch.nested.nested_tensor
except AttributeError:
    nested_tensor = getattr(torch, "nested_tensor", None)


T = TypeVar("T")


class Collator(ABC, Generic[T]):
    """Mapper of instances of type `T` to `torch.Tensor`s"""

    @abstractmethod
    def convert(self, value: T) -> torch.Tensor:
        """Convert a tensor-like object of type `T` into a `torch.Tensor`."""

    @abstractmethod
    def collate(self, batch: Sequence[T]) -> torch.Tensor:
        """
        Collate a batch of tensor-like objects of type `T` into a `torch.Tensor`
        with an additional outer dimension - batch size.
        """

    @classmethod
    def from_schemas(cls, *schemas: TensorSchema[Any]) -> Collator[Any]:
        """
        Return an appropriate Collator for collating instances generated by the
        `iter_tensors()` method of the given schema(s).
        """
        if len(schemas) > 1:
            return RowCollator(tuple(map(cls.from_schemas, schemas)))

        schema = schemas[0]
        collator: Collator[Any]
        if schema.kind is TensorKind.DENSE:
            collator = NumpyArrayCollator()
        elif schema.kind is TensorKind.RAGGED:
            collator = NumpyArrayCollator(to_nested=True)
        else:
            assert schema.kind in (TensorKind.SPARSE_COO, TensorKind.SPARSE_CSR)
            to_csr = schema.kind is TensorKind.SPARSE_CSR
            # The sparse arrays to be converted/collated are scipy.sparse.csr_matrix if
            # (and only if) the schema ndim == 2
            if len(schema.shape) == 2:
                collator = ScipySparseCSRCollator(to_csr)
            else:
                collator = SparseCOOCollator(to_csr)

        num_fields = schema.num_fields
        return RowCollator((collator,) * num_fields) if num_fields > 1 else collator


@dataclass(frozen=True)
class RowCollator(Collator[Sequence[Any]]):
    """
    Collator of "row" tuples.

    - All rows must have the same length.
    - The i-th element of every row (i.e. the i-th "column") must have the same type.
    - The i-th column values are collated by the i-th collator given in the constructor.
    """

    column_collators: Sequence[Collator[Any]]

    def convert(self, value: Sequence[Any]) -> Sequence[torch.Tensor]:
        assert len(value) == len(self.column_collators)
        return tuple(
            collator.convert(item)
            for collator, item in zip(self.column_collators, value)
        )

    def collate(self, batch: Sequence[Sequence[Any]]) -> Sequence[torch.Tensor]:
        columns = tuple(zip(*batch))
        assert len(columns) == len(self.column_collators)
        return tuple(
            collator.collate(column)
            for collator, column in zip(self.column_collators, columns)
        )


@dataclass(frozen=True)
class NumpyArrayCollator(Collator[np.ndarray]):
    """Collator of Numpy arrays

    to_nested: If true, collate 1D Numpy arrays with possibly different length to nested
        `torch.Tensor`s. Otherwise, all the arrays to be collated must have the same shape
    """

    to_nested: bool = False

    def convert(self, value: np.ndarray) -> torch.Tensor:
        return torch.from_numpy(value)

    def collate(self, batch: Sequence[np.ndarray]) -> torch.Tensor:
        if self.to_nested:
            return nested_tensor(list(map(torch.from_numpy, batch)))
        else:
            return torch.from_numpy(np.stack(batch))


@dataclass(frozen=True)
class SparseCOOCollator(Collator[sparse.COO]):
    """Collator of sparse.COO instances

    to_csr: Collate to `torch.Tensor`s with `sparse_csr` layout if true and `sparse_coo`
        layout if False
    """

    to_csr: bool = False

    def convert(self, value: sparse.COO) -> torch.Tensor:
        if self.to_csr:
            csr = value.tocsr()
            return torch.sparse_csr_tensor(csr.indptr, csr.indices, csr.data, csr.shape)
        else:
            return torch.sparse_coo_tensor(value.coords, value.data, value.shape)

    def collate(self, batch: Sequence[sparse.COO]) -> torch.Tensor:
        return self.convert(sparse.stack(batch))


@dataclass(frozen=True)
class ScipySparseCSRCollator(Collator[scipy.sparse.csr_matrix]):
    """Collator of `scipy.sparse.csr_matrix` instances

    to_csr: Collate to `torch.Tensor`s with `sparse_csr` layout if true and `sparse_coo`
        layout if False
    """

    to_csr: bool = False

    def convert(
        self, value: scipy.sparse.csr_matrix, _collating: bool = False
    ) -> torch.Tensor:
        if self.to_csr:
            return torch.sparse_csr_tensor(
                value.indptr, value.indices, value.data, value.shape
            )

        # if not collating, convert (1, N) sparse arrays to 1D sparse COO tensors
        if not _collating and value.shape[0] == 1:
            return torch.sparse_coo_tensor(
                value.indices.reshape(1, -1), value.data, value.shape[1:]
            )

        coo = value.tocoo()
        coords = np.stack((coo.row, coo.col))
        return torch.sparse_coo_tensor(coords, coo.data, coo.shape)

    def collate(self, batch: Sequence[scipy.sparse.csr_matrix]) -> torch.Tensor:
        return self.convert(scipy.sparse.vstack(batch), _collating=True)
