# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_meiosis.ipynb.

# %% auto 0
__all__ = ['poisson_crossing_over', 'Meiosis']

# %% ../nbs/03_meiosis.ipynb 3
import torch
from .core import *
from .trait import *

def poisson_crossing_over(chrom_lengths: torch.Tensor) -> list:
    """
    Generate crossing over locations for each chromosome given its length.
    
    Args:
    chrom_lengths (torch.Tensor): A 1D tensor containing the lengths of the chromosomes in centiMorgans.
    
    Returns:
    list of tensors: A list of tensors, each containing the crossing over locations for a chromosome.
    """
    crossing_over_locations = []
    for length in chrom_lengths:
        # Generate a random number of crossing over points for each chromosome
        num_crossovers = torch.randint(1, 5, (1,)).item()  # Random number between 1 and 5
        # Generate random crossing over locations within the chromosome length
        crossovers = torch.rand(num_crossovers) * length
        # Sort the crossing over locations
        crossovers, _ = torch.sort(crossovers)
        crossing_over_locations.append(crossovers)
    
    return crossing_over_locations

# %% ../nbs/03_meiosis.ipynb 7
import torch
from typing import List

from .core import Genome, Individual, Population
class Meiosis:
    """Class to handle meiosis and gamete formation."""

    @staticmethod
    def simulate_gametes(parent: Individual, min_crossovers: int = 1, 
                         max_crossovers: int = 5) -> List[Individual]:
        """
        Simulates gamete formation for a single parent with random crossovers.

        Args:
            parent (Individual): The parent individual.
            min_crossovers (int): Minimum number of crossovers per chromosome. Defaults to 1.
            max_crossovers (int): Maximum number of crossovers per chromosome. Defaults to 5.

        Returns:
            List[Individual]: A list of gametes (haploid Individuals).
        """
        genome = parent.genome  # Access genome from the parent Individual
        ploidy, n_chromosomes, n_loci = genome.shape()
        chromosome_lengths = genome.genetic_map.max(dim=1).values
        gametes = []

        for chrom_pair in range(ploidy // 2):  # Iterate over chromosome pairs
            haplotype = torch.zeros((1, n_chromosomes, n_loci), 
                                   device=genome.device, dtype=torch.int)
            
            for chr_idx in range(n_chromosomes):
                # Sample number of crossovers for this chromosome
                n_crossovers = torch.randint(min_crossovers, max_crossovers + 1, (1,)).item()

                # Sample crossover locations within chromosome length
                crossover_positions = torch.sort(
                    torch.rand(n_crossovers, device=genome.device) * chromosome_lengths[chr_idx]
                ).values

                # Efficiently recombine haplotypes using masks and indexing
                haplotype[0, chr_idx] = Meiosis._recombine_chromosome_pytorch(
                    parent.haplotypes[2*chrom_pair, chr_idx],
                    parent.haplotypes[2*chrom_pair + 1, chr_idx],
                    crossover_positions,
                    genome.genetic_map[chr_idx]
                )

            gametes.append(Individual(genome, haplotype))

        return gametes

    @staticmethod
    def _recombine_chromosome_pytorch(haplotype1: torch.Tensor, 
                                        haplotype2: torch.Tensor, 
                                        crossover_positions: torch.Tensor,
                                        genetic_map: torch.Tensor) -> torch.Tensor:
        """
        Recombines two haplotypes using PyTorch operations for efficiency.

        Args:
            haplotype1 (torch.Tensor): First haplotype (n_loci).
            haplotype2 (torch.Tensor): Second haplotype (n_loci).
            crossover_positions (torch.Tensor): Sorted crossover positions (n_crossovers).
            genetic_map (torch.Tensor): Genetic map of the chromosome (n_loci).

        Returns:
            torch.Tensor: Recombined haplotype (n_loci).
        """
        recombined = haplotype1.clone()  
        switch_mask = torch.zeros_like(genetic_map, dtype=torch.bool)

        # Create a mask to switch between haplotypes at crossover points
        for co_pos in crossover_positions:
            switch_mask |= genetic_map >= co_pos 

        # Use the mask to select from the appropriate haplotype
        recombined[switch_mask] = haplotype2[switch_mask]

        return recombined
