import numpy as np
import pandas as pd
import warnings
import itertools

def ibd_scores(ibd_table, rank_threshold:int=100):
    """
    Genetic distance between offspring and parents for pairs of plausible 
    candidate parents.

    This starts with the premise that crossed individuals can be thought of as 
    a mosaic of the two parents, such that in any single window they ought to be
    identical to one, but not the other parent.
    
    Beginning with a table of genetic distances between a test individual and
    each genotype in a reference panel, this function compares all plausible 
    pairs of candidates and pulls out the *maximum* genetic similarity to the 
    offspring in each window. It then takes the average over windows, excluding
    windows with NA values.
    
    Because there can be a very large number of pairs to consider, candidates 
    are first filtered to retain the N candidates with the lowest mean genetic
    distance from the offspring. N is controlled by the argument
    `rank_threshold`.

    Parameters
    ==========
    ibd_table: pd.DataFrame
        DataFrame with a row for each window in the genome and a column for each 
        sample in the reference panel. Elements show genetic distance between the 
        test individual and each reference individual in a single window.
        This is generated by ibd_table().
    rank_threshold: int, optional
        Number of most-plausible to consider as pairs. Defaults to 20. Setting 
        this number to something large will make the output very large!

    Returns
    =======
    A DataFrame with a row for each pair of candidate parents, giving candidate
    IDs and the average mimimum genetic distance from either parent to the
    offspring, averaged over windows. Values closer to zero indicate that the 
    pair is more likely to be a match.
    """
    # Coerce missing data to NaN for correct column means.
    ibd_table = ibd_table.replace(-9,np.nan)
    
    # Get column-mean IBD for each candidate, allowing for missing data
    # Note the index from 1 onwards to skip the window column
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        ibd_scores_for_each_candidate = np.array(
            [ np.nanmean(ibd_table[col]) for col in ibd_table.keys()[1:] ]
        )
    # Get the indices of the top N values in ibd_scores_for_each_candidate
    # where N is rank_threshold
    if rank_threshold < ibd_table.shape[1]:
        score_ix = np.argpartition(ibd_scores_for_each_candidate, -rank_threshold)[-rank_threshold:]
    else:
        score_ix = np.arange(len(ibd_scores_for_each_candidate))
    
       
    # Reference sample names, excluding the column for window
    all_names = ibd_table.keys()[1:] # All names in the reference panel
    ref_names = all_names[score_ix] # Top candidates
    # For calculating scores for pairs we don't want combinations with 'Expected_F1'
    pair_names = list(ref_names)
    if 'Expected_F1' in pair_names:
        pair_names.remove("Expected_F1")
    
    # Dataframe giving the scores if parents are selfed.
    score_table = pd.DataFrame({
        'parent1'    : ref_names,
        'parent2'    : ref_names, # Under selfing Parent1 and Parent2 are the same. 
        'similarity' : ibd_scores_for_each_candidate[score_ix]
    })
    
    # Get scores for all pairs of top parents.
    # This is the maxmimum similarity score for either parent in a pair in each
    # window.
    # No pairs should contain "Expected_F1" as a candidate.
    if rank_threshold > 1:
        scores_for_pairs = []
        for i,j in itertools.combinations(pair_names,2):
            # The maximum similarity at each window between the test individual and 
            # *either* reference genotype in this pair
            max_similarity_pair = ibd_table[[i,j]].max(axis=1, skipna=False)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                # List containing the ID of candidates 1 and 2 and the average of
                # the best similarities at each window.
                scores_for_pairs.append(
                    [ i,j, np.nanmean(max_similarity_pair) ]
                    )
        scores_for_pairs = pd.DataFrame(scores_for_pairs, columns = ['parent1', 'parent2', 'similarity'])
    
        # A single dataframe combining scores for selfed and pairs.
        score_table = pd.concat([score_table, scores_for_pairs])
        
    # Sort so the best matches are at the top
    score_table = score_table.sort_values('similarity', ascending=False)
    
    return score_table