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

# %% auto 0
__all__ = ['gamma_interference_model', 'simulate_meiosis', 'simulate_gametes']

# %% ../nbs/03_meiosis.ipynb 4
from .core import *
import torch
def gamma_interference_model(length, rate, shape):
    """
    Simulate crossover events using a gamma interference model.
    
    Parameters:
    length (float): Length of the chromosome.
    rate (float): Rate of crossover events.
    shape (float): Shape parameter for the gamma distribution.
    
    Returns:
    torch.Tensor: Positions of crossover events.
    """
    num_crossovers = torch.poisson(torch.tensor([rate * length]))
    intervals = torch.distributions.gamma.Gamma(shape, rate).sample((int(num_crossovers.item()),))
    crossover_positions = torch.cumsum(intervals, dim=0)
    crossover_positions = crossover_positions[crossover_positions < length]
    return crossover_positions


import torch

def simulate_meiosis(num_chromosomes, map_length, num_individuals, num_crossovers):
    """
    This function simulates random crossover events across chromosomes.
    """
    return [torch.sort(torch.rand((num_individuals, num_crossovers)) * map_length, dim=-1)[0] for _ in range(num_chromosomes)]

def simulate_gametes(genetic_map, parent_genome):
    """
    Simulate the formation of a single gamete given crossover positions, genetic map, and parent genomes.

    Parameters:
    genetic_map (list of torch.Tensor): List of positions of genetic markers on the chromosomes.
    parent_genome (torch.Tensor): Genomes of the parents. SHAPE:  (ploidy, num_chromosomes, num_loci)

    Returns:
    torch.Tensor: The resultant single gamete. SHAPE: (ploidy//2, num_chromosomes, num_loci)
    """
    ploidy, num_chromosomes, num_loci = parent_genome.shape
    gamete_genome = torch.zeros((ploidy // 2, num_chromosomes, num_loci), dtype=parent_genome.dtype)
    crossover_positions = simulate_meiosis(num_chromosomes, 100.0, 1, 1)

    for chrom in range(num_chromosomes):
        # Create a crossover mask for each chromosome
        crossover_mask = torch.zeros(num_loci, dtype=torch.bool)
        crossover_sites = crossover_positions[chrom].flatten()
        for position in crossover_sites:
            # Find the nearest marker index for each crossover position
            index = torch.argmin(torch.abs(genetic_map[chrom] - position))
            crossover_mask[index:] = ~crossover_mask[index:]  # Flip values at and after the crossover index

        # Use the mask to select genes from parent 1 or parent 2
        gamete_genome[0, chrom] = torch.where(crossover_mask, parent_genome[1, chrom], parent_genome[0, chrom])

    return gamete_genome


