"""
Affine partitions generated by fixed sequences of matrices.
"""

# Part of: q1ss
# Copyright (C) 2023 Hashberg Ltd and 20squares UG

# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
# USA

from __future__ import annotations
from collections.abc import Sequence
from typing import Literal, Union, final
from typing_extensions import Self
import numpy as np
from typing_validation import validate

from ..binalg import binvec, binmat, AffineSubspace
from ..binalg.binvec import Bit
from .vectorized import sequence_ap_label

from .base import AP
from .generated import GenAP, APMatrixGen


@final
class SeqAPData:
    r"""
    Data used to construct a :class:`SeqAP` instance:

    - an ambient vector space dimension :math:`n`
    - an :math:`n \times n` invertible ``start`` matrix
    - two sequences ``mats0`` and ``mats1`` of :math:`m` invertible matrices,
      where the matrix at index :math:`j=0,...,m-1` in each sequence has size
      :math:`(n-1-j)\times(n-1-j)`

    The inverses to all matrices above can be supplied, or are otherwise
    computed, and are stored in ``start_inv``, ``mats0_inv``, and ``mats1_inv``.
    The subspace dimension for the affine partition is :math:`n-m-1`.

    """

    @staticmethod
    def random(
        subsp_dim: int,
        ambient_dim: int,
        *,
        rng: np.random.Generator | int | None = None,
    ) -> SeqAPData:
        """
        Creates random :class:`SeqAP` data for given subspace and ambient
        dimension, using :meth:`binmat.random_inv` to sample matrix-inverse
        pairs for all matrices involved.

        .. warning::

            For moderately large dimensions, this method is computationally
            intensive. For example, ``subsp_dim=128`` and ``ambient_dim=256``
            takes roughly 20s on a i7-1065G7. This will be optimised in future
            releases.

        """
        assert AffineSubspace._validate_dimensions(subsp_dim, ambient_dim)
        if not isinstance(rng, np.random.Generator):
            rng = np.random.default_rng(rng)
        start, start_inv = binmat.random_inv(ambient_dim, rng=rng)
        num_mats = ambient_dim - subsp_dim - 1
        mats0 = [
            binmat.random_inv(ambient_dim - 1 - i, rng=rng)
            for i in range(num_mats)
        ]
        mats1 = [
            binmat.random_inv(ambient_dim - 1 - i, rng=rng)
            for i in range(num_mats)
        ]
        data = SeqAPData(ambient_dim, (start, start_inv), mats0, mats1)
        assert data.is_complete
        return data

    _ambient_dim: int
    _subsp_dim: int
    _start: binmat
    _start_inv: binmat
    _mats0: tuple[binmat, ...]
    _mats1: tuple[binmat, ...]
    _mats0_inv: tuple[binmat, ...]
    _mats1_inv: tuple[binmat, ...]

    __slots__ = (
        "__weakref__",
        "_ambient_dim",
        "_subsp_dim",
        "_start",
        "_start_inv",
        "_mats0",
        "_mats1",
        "_mats0_inv",
        "_mats1_inv",
    )

    def __new__(
        cls,
        ambient_dim: int,
        start: binmat | tuple[binmat, binmat],
        mats0: Sequence[Union[binmat, tuple[binmat, binmat]]],
        mats1: Sequence[Union[binmat, tuple[binmat, binmat]]],
        *,
        copy: bool = False,
    ) -> Self:
        """
        Creates an empty :class:`SeqAP` data container for given ambient
        dimension.
        Each invertible matrix in ``start``, ``mats0`` and ``mats1`` can be
        passed as either a single matrix, or a pair of a matrix and its inverse:
        the inverse is automatically computed in the first case and it is
        validated in the second case.
        All matrices are made read-only as part of the constructor: if this is
        not desirable, the ``copy`` parameter can be set to ``True`` to make
        fresh copies of all matrices.

        :meta public:
        """
        binvec.validate_dim(ambient_dim, positive=True)
        instance = super().__new__(cls)
        instance._ambient_dim = ambient_dim
        instance.__set_start(start, copy=copy)
        instance.__set_mats(mats0, b=0, copy=copy)
        instance.__set_mats(mats1, b=1, copy=copy)
        return instance

    def __getnewargs__(
        self,
    ) -> tuple[
        int,
        binmat | tuple[binmat, binmat],
        Sequence[Union[binmat, tuple[binmat, binmat]]],
        Sequence[Union[binmat, tuple[binmat, binmat]]],
    ]:
        """
        Method for pickling.
        """
        return (
            self.ambient_dim,
            (self.start, self.start_inv),
            tuple(zip(self.mats0, self.mats0_inv)),
            tuple(zip(self.mats1, self.mats1_inv)),
        )

    @property
    def as_matrix_gen(self) -> SeqAPMatrixGen:
        """
        Returns the matrix generating function for this data.
        """
        return SeqAPMatrixGen(self)

    @property
    def is_complete(self) -> bool:
        """
        Whether the affine partition data has been completely specified.
        """
        return (
            hasattr(self, "_start")
            and hasattr(self, "_mats0")
            and hasattr(self, "_mats1")
        )

    @property
    def ambient_dim(self) -> int:
        """
        The ambient dimension for the affine partition.
        """
        return self._ambient_dim

    @property
    def subsp_dim(self) -> int:
        """
        The subspace dimension for the affine partition.
        """
        return self._subsp_dim

    @property
    def label_dim(self) -> int:
        """
        The label dimension for the affine partition.
        """
        return self.ambient_dim - self.subsp_dim

    @property
    def start(self) -> binmat:
        """
        The ``start`` matrix for this affine partition.
        """
        return self._start

    @property
    def start_inv(self) -> binmat:
        """
        Inverse of the ``start`` matrix for this affine partition.
        """
        return self._start_inv

    @property
    def mats0(self) -> tuple[binmat, ...]:
        """
        The ``mats0`` matrices for this affine partition.
        """
        return self._mats0

    @property
    def mats0_inv(self) -> tuple[binmat, ...]:
        """
        Inverses of the ``mats0`` matrices for this affine partition.
        """
        return self._mats0_inv

    @property
    def mats1(self) -> tuple[binmat, ...]:
        """
        The ``mats1`` matrices for this affine partition.
        """
        return self._mats1

    @property
    def mats1_inv(self) -> tuple[binmat, ...]:
        """
        Inverses of the ``mats1`` matrices for this affine partition.
        """
        return self._mats1_inv

    def __set_start(
        self, value: binmat | tuple[binmat, binmat], *, copy: bool
    ) -> None:

        if isinstance(value, binmat):
            assert self.__validate_mat(value)
            start = value
            start_inv = ~value
        else:
            assert validate(value, tuple[binmat, binmat])
            start, start_inv = value
            assert self.__validate_mat(start, start_inv)
        if copy:
            start, start_inv = start.copy(), start_inv.copy()
        binmat.make_readonly(start, start_inv)
        self._start = start
        self._start_inv = start_inv

    def __set_mats(
        self,
        values: Sequence[Union[binmat, tuple[binmat, binmat]]],
        b: Bit,
        *,
        copy: bool,
    ) -> None:
        assert self.__validate_mats_seq(values, b=b)
        mats: list[binmat] = []
        mats_inv: list[binmat] = []
        for j, value in enumerate(values):
            if isinstance(value, binmat):
                assert self.__validate_mat(value, j=j, b=b)
                mats.append(value)
                mats_inv.append(~value)
            else:
                assert validate(value, tuple[binmat, binmat])
                m, m_inv = value
                assert self.__validate_mat(m, m_inv, j=j, b=b)
                mats.append(m)
                mats_inv.append(~m)
        if copy:
            mats = [m.copy() for m in mats]
            mats_inv = [m.copy() for m in mats_inv]
        binmat.make_readonly(*mats, *mats_inv)
        if b == 0:
            self._mats0 = tuple(mats)
            self._mats0_inv = tuple(mats_inv)
        else:
            self._mats1 = tuple(mats)
            self._mats1_inv = tuple(mats_inv)
        if not hasattr(self, "_subsp_dim"):
            self._subsp_dim = self.ambient_dim - len(values) - 1

    def __validate_mat(
        self,
        mat: binmat,
        mat_inv: binmat | None = None,
        *,
        j: int = -1,
        b: Bit | None = None,
    ) -> Literal[True]:
        n = self._ambient_dim
        if b is None:
            mat_label = "start"
        else:
            mat_label = f"mats{b}[{j}]"
        k = n - j - 1
        if mat.shape != (k, k):
            raise ValueError(
                f"The '{mat_label}' matrix must have "
                f"shape {(k, k)}, found {mat.shape = }."
            )
        if mat_inv is not None:
            if mat.shape != (k, k):
                raise ValueError(
                    f"Inverse of the '{mat_label}' matrix must have "
                    f"shape {(k, k)}, found {mat.shape = }."
                )
            if not (mat @ mat_inv).is_eye:
                raise ValueError(
                    f"The '{mat_label}' matrix and the supplied inverse must "
                    "actually be inverses of each other."
                )
        return True

    def __validate_mats_seq(
        self, mats: Sequence[Union[binmat, tuple[binmat, binmat]]], b: Bit
    ) -> Literal[True]:
        validate(mats, Sequence[Union[binmat, tuple[binmat, binmat]]])
        if len(mats) >= self.ambient_dim - 1:
            raise ValueError(
                f"Expected at most {self.ambient_dim - 2} matrices, "
                f"found {len(mats)}."
            )
        return True


class SeqAPMatrixGen(APMatrixGen):
    """
    Matrix generator based on :class:`SeqAPData`.
    """

    _data: SeqAPData
    __slots__ = ("_data",)

    def __new__(cls, data: SeqAPData) -> Self:
        """
        Creates a new matrix generator using the given :class:`SeqAPData`

        :meta public:
        """
        assert validate(data, SeqAPData)
        ambient_dim = data.ambient_dim
        max_label_dim = data.label_dim
        instance = super().__new__(cls, ambient_dim, max_label_dim)
        instance._data = data
        return instance

    @property
    def data(self) -> SeqAPData:
        """
        The data used to generate matrices.
        """
        return self._data

    def __getnewargs__(self) -> tuple[SeqAPData]:
        """
        Method for pickling.
        """
        return (self._data,)

    def _get_matrix(self, label: binvec) -> tuple[binmat, binmat]:
        """
        Generates a (matrix, inverse) pair based on the last bit of the label.

        - If ``label`` is empty, returns the pair of ``start`` and its inverse.
        - If the last bit of ``label`` is 0, returns the pair of ``mats0[k-1]``
          and its inverse, where ``k=len(label)``.
        - If the last bit of ``label`` is 1, returns the pair of ``mats1[k-1]``
          and its inverse, where ``k=len(label)``.

        :meta public:
        """
        data = self._data
        k = len(label)
        if k == 0:
            return (data.start, data.start_inv)
        b = label[-1]
        if b:
            return (data.mats1[k - 1], data.mats1_inv[k - 1])
        return (data.mats0[k - 1], data.mats0_inv[k - 1])


@final
class SeqAP(AP):
    """
    Class for ordered balanced affine partitions of ``n``-bitstrings defined
    by a fixed sequence of invertible matrices for partial labels.
    """

    @staticmethod
    def random(
        subsp_dim: int,
        ambient_dim: int,
        *,
        rng: np.random.Generator | int | None = None,
    ) -> SeqAP:
        """
        Creates a random :class:`SeqAP` for given subspace and ambient
        dimension, using :meth:`SeqAPData.random` to create random data.
        """
        data = SeqAPData.random(subsp_dim, ambient_dim, rng=rng)
        return SeqAP(data)

    _data: SeqAPData
    __slots__ = ("_data",)

    def __new__(cls, data: SeqAPData) -> Self:
        """
        Constructs a new affine partition from the given sequences of matrices.

        :meta public:
        """
        assert SeqAP.__validate_data(data)
        instance = super().__new__(cls, data.subsp_dim, data.ambient_dim)
        instance._data = data
        return instance

    def __getnewargs__(self) -> tuple[SeqAPData]:
        """
        Method for pickling.
        """
        return (self._data,)

    @property
    def as_generated(self) -> GenAP:
        """
        Returns the :class:`GenAP` instance corresponding to this
        sequence-based affine partition.
        """
        return GenAP(self.subsp_dim, self.data.as_matrix_gen)

    @property
    def data(self) -> SeqAPData:
        """
        The data used to construct this affine partition.
        """
        return self._data

    def _get_subspace(self, label: binvec) -> AffineSubspace:
        r"""
        Returns the affine subspace corresponding to the given label,
        providing the underlying logic for :meth:`~AP.__getitem__`.

        Let ``k`` be the :attr:`~AP.label_dim` and ``n`` be the
        :attr:`~AP.ambient_dim` for this generator, so that ``n-k`` is the
        :attr:`~AP.subsp_dim`.

        1. Start from the standard affine subspace spanned by the first
           ``n-k`` standard basis vectors. The basepoint is given by a vector
           which is zero on the first ``n-k`` components, and is the reverse of
           the label on the remaining ``k`` components.
        2. Selects a sequence of ``k-2`` inverse matrices from the :attr:`data`,
           selecting from ``data.mats0_inv`` and ``data.mats1_inv`` based on the
           reverse bits of the label, starting from the penultimate and ending
           at the first. Then appends ``data.start_inv`` as the final matrix
           in the sequence, which has ``k-1`` matrices in total.
        3. Transform the standard subspace using the inverse matrices in the
           sequence, where matrix at index ``i`` acts on the subspace spanned
           by the first ``n-k+1+i`` std. basis vectors (``n-k+1`` -> ``n``).

        Uses the :meth:`AffineSubspace.transform` method with ``partial=True``.

        :meta public:
        """
        subsp_dim, ambient_dim = self.subsp_dim, self.ambient_dim
        data = self.data
        basepoint = binvec.zeros(subsp_dim) | label[::-1]
        subsp = AffineSubspace.std(subsp_dim, ambient_dim, basepoint)
        mats_inv = (data.mats0_inv, data.mats1_inv)
        matrices = [
            mats_inv[label[i]][i] for i in range(len(label) - 2, -1, -1)
        ] + [data.start_inv]
        return subsp.transform(matrices, partial=True)

    def _label(self, vec: binvec) -> binvec:
        """
        Returns the label of the given vector in this affine partition,
        providing the underlying logic for :meth:`~AP.label`.

        Let ``k`` be the :attr:`~AP.label_dim` and ``n`` be the
        :attr:`~AP.ambient_dim` for this generator.

        1. Start from the given vector.
        2. Applies the ``data.start`` matrix to the vector (unless ``k`` is 0,
           in which case the vector is returned unchanged by this method).
        3. For each ``i in range(k-1)``, generates a matrix using the bit at
           position ``k-i-1`` in the vector, selects ``mats0[i]`` or
           ``mats1[i]`` based on its value, and applies the matrix to the
           first ``n-i-1`` bits of the vector.
        4. Returns the last ``k`` bits of the resulting vector, read in reverse.

        Note that each matrix application keeps fixed the bits of the vector
        upon which the matrix itself depended.

        :meta public:
        """
        subsp_dim, ambient_dim = self.subsp_dim, self.ambient_dim
        data = self.data
        if subsp_dim == ambient_dim:
            return binvec([])
        elif subsp_dim == ambient_dim - 1:
            return binvec([(data.start @ vec)[-1] % 2])
        return binvec(
            sequence_ap_label(
                ambient_dim,
                self.label_dim,
                data.start._data,
                [m._data for m in data.mats0],
                [m._data for m in data.mats1],
                vec._data,
            )
        )
        ## Below is the original non-vectorized version:
        # start, mats0, mats1 = data.start, data.mats0, data.mats1
        # n, k = self.ambient_dim, self.label_dim
        # vec = start @ vec
        # for i in range(k - 1):
        #     if vec[-1 - i]:
        #         vec[: n - i - 1] = mats1[i] @ vec[: n - i - 1]
        #     else:
        #         vec[: n - i - 1] = mats0[i] @ vec[: n - i - 1]
        # return vec[-1: -k - 1: -1]

    @staticmethod
    def __validate_data(data: SeqAPData) -> Literal[True]:
        validate(data, SeqAPData)
        if not data.is_complete:
            raise ValueError("Data for affine partition is not complete.")
        return True
