from __future__ import annotations

from collections.abc import Iterable
from radiate.radiate import PyGenotype
from radiate._bridge.wrapper import RsObject

from .chromosome import Chromosome
from .gene import GeneType


class Genotype[T](RsObject):
    """
    Represents a genotype in a genome.
    """

    def __init__(
        self,
        chromosomes: Iterable[Chromosome[T]] | Chromosome[T] | None = None,
    ):
        super().__init__()

        if isinstance(chromosomes, Chromosome):
            self._pyobj = PyGenotype([chromosomes.__backend__()])
        elif isinstance(chromosomes, Iterable):
            self._pyobj = PyGenotype(list(map(lambda c: c.__backend__(), chromosomes)))

    def __repr__(self):
        return self._pyobj.__repr__()

    def __len__(self):
        """
        Returns the length of the genotype.
        :return: Length of the genotype.
        """
        return len(self._pyobj)

    def __getitem__(self, index: int | slice) -> Chromosome[T]:
        """
        Returns the chromosome at the specified index.
        :param index: Index of the chromosome to retrieve.
        :return: Chromosome instance at the specified index.
        """
        return self.chromosomes()[index]

    def __iter__(self):
        """
        Returns an iterator over the chromosomes in the genotype.
        :return: An iterator over the chromosomes in the genotype.
        """
        for chromosome in self.chromosomes():
            yield chromosome

    def gene_type(self) -> "GeneType":
        """
        Returns the type of the genes in the genotype.
        :return: The gene type as a string.
        """
        from . import GeneType

        return GeneType.from_str(self._pyobj.gene_type())

    def chromosomes(self) -> list[Chromosome[T]]:
        """
        Get the chromosomes of the genotype.
        :return: The chromosomes of the genotype.
        """
        return self.try_get_cache(
            "chromosomes_cache",
            lambda: [Chromosome.from_rust(c) for c in self.__backend__().chromosomes()],
        )
