import pandas as pd
import numpy as np

import plotly.express as px

def plot_ibd_table(ibd_table:pd.DataFrame, sample_name:str, expected_match:list=[], max_to_plot=10):
    """
    Plot allele sharing across the genome.

    Create a interactive line graph showing genetic distance from a test
    individual to each sample in a panel of reference individuals.

    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().
    sample_name: str
        Sample name for the individual to check.
        This must be present in the samples in the input VCF.
    expected_match: list
        List of sample names in the reference panel that are expected to be
        ancestors of the test individual.

    Returns
    =======
    Plotly figure object with subplots for each chromosome, showing window
    position along the x-axis and genetic distance from the test individual to
    each reference sample on the y-axis. Line colour indicates whether a sample
    is an expected parent or not. Rolling over the lines shows which sample is
    which.
    """
    # Add the expected F1 to the list of expected matches so that it is plotted in colour
    if len(expected_match) == 2:
        expected_match = expected_match + ['Expected_F1']

    # Identify the candidate names *not* among the top `max_to_plot` columns and remove
    # If `max_to_plot` is less than the number of candidates.
    if max_to_plot < ibd_table.shape[1]-1:
        # Get column-mean IBD for each candidate, allowing for missing data
        ibd_scores_for_each_candidate = np.array(
            [ np.nanmean(ibd_table[col]) for col in ibd_table.keys()[1:] ]
        )
        # Identify the candidate names *not* among the top `max_to_plot` columns
        ix = np.argpartition(-ibd_scores_for_each_candidate, max_to_plot)[max_to_plot:] # index positions
        columns_to_drop = ibd_table.keys()[ix+1].to_list() # candidate names to be removed
        # Make sure the expected parents are not among the proscribed
        [ columns_to_drop.remove(x) for x in expected_match if x in columns_to_drop ]        
        ibd_table = ibd_table.drop(columns=columns_to_drop) # drop the candidates

    # Make the table long
    ibd_table = ibd_table.melt(id_vars=['window'], var_name='candidate', value_name='similarity')
    # Column indicating which candidates should be plotted a different colour.
    ibd_table['colour'] = np.where(
        ibd_table['candidate'].isin(expected_match), ibd_table['candidate'], "Other"
        )
    # Unique list of labels for the legend, sorted to plot "Other" first, then
    # then the expected matches (including heterozygosity, if using).
    unique_legend_labels = ['Other'] + expected_match

    # Split the 'window' column up into separate columns for chromosome, start and stop positions
    ibd_table[['chr', 'window']] = ibd_table['window'].str.split(":", expand=True)
    ibd_table[['start', 'stop']] = ibd_table['window'].str.split("-", expand=True)
    # start and stop positions should be integers for sensible plotting.
    ibd_table['start'] = ibd_table['start'].astype(int)
    ibd_table['stop'] = ibd_table['stop'].astype(int)
    ibd_table['midpoint'] = (ibd_table['start'] + ibd_table['stop']) / 2
        
    # Set the palette for the lines in the plot.
    # This sets 'Others', then uses the default palette for the first N colours
    custom_colors = ['LightGrey'] + px.colors.qualitative.Set1[0:len(unique_legend_labels)]

    # Make the plot
    fig = px.line(
        ibd_table,
        x="midpoint", y="similarity", color="colour", line_group="candidate",
        title=sample_name,
        labels={
            'midpoint' : 'Position (bp)',
            'distance' : 'Similarity'        
        },
        hover_data=['candidate'],
        template="plotly_white",
        color_discrete_sequence=custom_colors,
        category_orders={'colour': unique_legend_labels},
        facet_row = "chr"
        )
    fig.for_each_annotation(lambda a: a.update(text=a.text.split("=")[-1])) # Remove "chr=" from facet labels
    fig.update_layout(legend_title_text='Candidate')
    fig.update_traces(mode="markers+lines")

    return fig
