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

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

# %% ../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 5
def simulate_gametes(genome, parent_genomes, rate =1 , shape=1):
    """
    Simulate the formation of gametes for multiple parents using vectorized operations.

    Args:
        genome (Genome): The Genome instance containing the genetic map and other parameters.
        parent_genomes (torch.Tensor): Genomes of the parents. 
                                       Shape: (num_individuals, ploidy, num_chromosomes, num_loci)

    Returns:
        torch.Tensor: The resultant gametes. 
                      Shape: (num_individuals, ploidy//2, num_chromosomes, num_loci)
    """
    device = genome.device
    genetic_map = genome.genetic_map  # torch.Size([num_chromosomes, num_loci])
    num_individuals, ploidy, num_chromosomes, num_loci = parent_genomes.shape

    chromosome_lengths = genetic_map.max(dim=1).values

    # Simulate crossover positions for all chromosomes at once
#     all_crossovers = gamma_interference_model(chromosome_lengths, rate, shape, device)
    all_crossovers = poisson_crossing_over(chromosome_lengths)
    
    
    
    # Initialize gametes tensor
    gametes = torch.zeros(num_individuals, ploidy // 2, num_chromosomes, num_loci, device=device, dtype=parent_genomes.dtype)
    
    for chr_idx in range(num_chromosomes):
        crossovers = all_crossovers[chr_idx]

        if len(crossovers) > 0:
            crossover_mask = torch.zeros(num_loci, device=device, dtype=torch.bool)
            positions_idx = torch.searchsorted(genetic_map[chr_idx], crossovers)
            crossover_mask[positions_idx] = True

            parent_genome_1 = parent_genomes[:, ::2, chr_idx]
            parent_genome_2 = parent_genomes[:, 1::2, chr_idx]

            for ploid_idx in range(ploidy // 2):
                gametes[:, ploid_idx, chr_idx] = torch.where(crossover_mask.unsqueeze(0), 
                                                             parent_genome_1[:, ploid_idx], 
                                                             parent_genome_2[:, ploid_idx])
        else:
            gametes[:, :, chr_idx] = parent_genomes[:, ::2, chr_idx]

    return gametes

# Define your Genome class or struct here if needed, ensuring it includes 'device' and 'genetic_map'
