# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_cross.ipynb.

# %% auto 0
__all__ = ['Cross']

# %% ../nbs/04_cross.ipynb 3
from .core import *
from .trait import *
from .meiosis import *
from typing import Tuple, Optional, List, Union

import torch

# %% ../nbs/04_cross.ipynb 4
class Cross:
    """Class to handle crossing between individuals in a population."""

    @staticmethod
    def random_crosses(parent_population: Population, n_crosses: int) -> Population:
        """
        Generates a new population by randomly crossing individuals from a parent population.

        Args:
            parent_population (Population): Population of parent individuals.
            n_crosses (int): Number of crosses to generate.

        Returns:
            Population: A new population of progeny individuals. 
        """
        genome = parent_population.individuals[0].genome
        device = genome.device
        n_parents = parent_population.size()

        # Efficiently sample parent indices for all crosses at once
        female_indices = torch.randint(0, n_parents, (n_crosses,), device=device)
        male_indices = torch.randint(0, n_parents, (n_crosses,), device=device)

        # Get parent individuals from indices
        female_parents = [parent_population.individuals[i] for i in female_indices]
        male_parents = [parent_population.individuals[i] for i in male_indices]

        # Simulate gametes for all parents individually
        female_gametes = [Meiosis.simulate_gametes(parent)[0] for parent in female_parents] # Get the first gamete for each parent
        male_gametes = [Meiosis.simulate_gametes(parent)[0] for parent in male_parents] # Get the first gamete for each parent

        progeny_individuals = []
        for i in range(n_crosses):
            # Combine gametes to form progeny haplotypes
            progeny_haplotypes = torch.cat([female_gametes[i].haplotypes, male_gametes[i].haplotypes], dim=0)
            
            # Create progeny individual
            progeny_individual = Individual(genome, progeny_haplotypes)
            progeny_individuals.append(progeny_individual)

        return Population(progeny_individuals)
    
    
    @staticmethod
    def dh_crosses(parent_population: Population, n_family: int) -> List[Population]:
        """
        Generates doubled haploid progeny for each parent in a population.

        Args:
            parent_population (Population): Population of parent individuals.
            n_family (int): Number of DH lines to produce per parent.

        Returns:
            List[Population]: A list of populations, where each population 
                             represents the DH progeny of a single parent.
        """
        families = []
        for parent in parent_population.individuals:
            dh_individuals = []
            for _ in range(n_family):
                gamete = Meiosis.simulate_gametes(parent)[0]  # Get one gamete
                dh_haplotypes = gamete.haplotypes.repeat(2, 1, 1) # Double the haplotype
                dh_individual = Individual(parent.genome, dh_haplotypes)
                dh_individuals.append(dh_individual)
            families.append(Population(dh_individuals)) 

        return families
