from __future__ import annotations

from typing import overload, Sequence, Any, Literal, TYPE_CHECKING

from .base import CodecBase

from radiate.radiate import PyBitCodec
from radiate.genome import Genotype, Gene, Chromosome, GeneType
from radiate._bridge.wrapper import RsObject
from radiate._typing import (
    AtLeastOne,
)

if TYPE_CHECKING:
    from radiate._dependancies import numpy as np


class BitCodec[D](CodecBase[bool, D], RsObject):
    """BitCodec for bit-based chromosomes. Encodes/decodes to bit strings."""

    gene_type = GeneType.BIT

    @overload
    def __new__(
        cls,
        shape: int,
        *,
        use_numpy: Literal[False] = ...,
    ) -> "BitCodec[list[bool]]": ...

    @overload
    def __new__(
        cls,
        shape: int,
        *,
        use_numpy: Literal[True] = ...,
    ) -> "BitCodec[np.ndarray]": ...

    @overload
    def __new__(
        cls,
        shape: Sequence[int],
        *,
        use_numpy: Literal[False] = ...,
    ) -> "BitCodec[list[list[bool]]]": ...

    @overload
    def __new__(
        cls,
        shape: Sequence[int],
        *,
        use_numpy: Literal[True] = ...,
    ) -> "BitCodec[list[np.ndarray]]": ...

    def __new__(cls, *args: Any, **kwargs: Any) -> "BitCodec[Any]":
        return super().__new__(cls)

    def __init__(
        self,
        shape: AtLeastOne[int] | None = None,
        use_numpy: bool = False,
    ):
        """
        Initialize the bit codec with number of chromosomes and value bounds.
        :param chromosomes: Number of chromosomes with the number of genes in each chromosome.
        """
        if shape is not None:
            if isinstance(shape, int):
                self._pyobj = self.__vector(length=shape, use_numpy=use_numpy)
            elif isinstance(shape, (tuple, list)):
                self._pyobj = self.__matrix(shape=shape, use_numpy=use_numpy)
            else:
                raise ValueError(
                    "Shape must be an int, tuple of ints, or list of ints."
                )
        else:
            raise ValueError("Shape must be provided.")

    def encode(self) -> Genotype[bool]:
        """
        Encode the codec into a Genotype.
        :return: A Genotype instance.
        """
        return Genotype.from_rust(self.__backend__().encode_py())

    def decode(self, genotype: Genotype[bool]) -> D:
        """
        Decode a Genotype into its bit representation.
        :param genotype: A Genotype instance to decode.
        :return: The decoded bit representation of the Genotype.
        """
        if not isinstance(genotype, Genotype):
            raise TypeError("genotype must be an instance of Genotype.")
        return self.__backend__().decode_py(genotype.__backend__())

    @staticmethod
    def __matrix(shape: AtLeastOne[int], use_numpy: bool = False) -> BitCodec[Any]:
        """
        Initialize the bit codec with a matrix of chromosomes.
        Args:
            shape: A list of integers specifying the shape of the matrix.
        Returns:
            A new BitCodec instance with matrix configuration.

        Example
        --------
        >>> rd.BitCodec.matrix(chromosome_lengths=[5, 5])
        BitCodec(...)
        """
        if isinstance(shape, tuple):
            if len(shape) != 2:
                raise ValueError("Shape must be a tuple of (rows, cols).")
            rows, cols = shape
            if rows < 1 or cols < 1:
                raise ValueError("Rows and columns must be at least 1.")
            shape = [cols for _ in range(rows)]
        elif isinstance(shape, list):
            if not all(isinstance(x, int) and x > 0 for x in shape):
                raise ValueError("Shape must be a list of positive integers.")

        return PyBitCodec.matrix(chromosome_lengths=shape, use_numpy=use_numpy)

    @staticmethod
    def __vector(length: int = 8, use_numpy: bool = False) -> BitCodec[Any]:
        """
        Initialize the bit codec with a single chromosome of specified length.
        Args:
            length: Length of the chromosome.
        Returns:
            A new BitCodec instance with vector configuration.

        Example
        --------
        >>> rd.BitCodec.vector(length=5)
        BitCodec(...)
        """
        return PyBitCodec.vector(chromosome_length=length, use_numpy=use_numpy)

    @staticmethod
    def __from_genes(
        genes: AtLeastOne[Gene[bool]], use_numpy: bool = False
    ) -> PyBitCodec:
        if not isinstance(genes, (list, tuple)):
            raise TypeError("genes must be a list or tuple of Gene instances.")

        return PyBitCodec.from_genes(
            list(map(lambda g: g.__backend__(), genes)), use_numpy=use_numpy
        )

    @staticmethod
    def __from_chromosomes(
        chromosomes: AtLeastOne[Chromosome[bool]],
        use_numpy: bool,
    ) -> PyBitCodec:
        from radiate.genome import GeneType

        if isinstance(chromosomes, Chromosome):
            chromosomes = [chromosomes]
        if not all(isinstance(c, Chromosome) for c in chromosomes):
            raise TypeError(
                "chromosomes must be a list or tuple of Chromosome instances."
            )
        if not all(g.gene_type() == GeneType.BIT for c in chromosomes for g in c):
            raise TypeError("All chromosomes must be of type 'bool'.")

        return PyBitCodec.from_chromosomes(
            list(map(lambda c: c.__backend__(), chromosomes)), use_numpy=use_numpy
        )
