import warnings
from typing import Optional

import matplotlib.axes
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import pearsonr

from starfish import IntensityTable
from starfish.types import Features
from .util import equalize_axes_extents


def compare_copy_number(
        intensities: IntensityTable,
        other: pd.Series,
        ax: Optional[matplotlib.axes.Axes]=None,
        equalize_axes: bool=False,
        **kwargs,
) -> matplotlib.axes.Axes:
    """
    Plot the copy number of gene expression across all cells against a reference generated by
    another computational approach. Report the $r^2$ value.

    Parameters
    ----------
    intensities : IntensityTable
        decoded IntensityTable
    other : pd.Series
        A pandas series whose index contains gene names and whose values are expression abundances
    ax : Optional[matplotlib.axes.Axes]
        The axis in which to plot the figure. Defaults to the current axis if not provided.
    equalize_axes : bool
        If true, equalize the size of the two axes such that the slope of the comparison can be
        easily interpreted
    kwargs : Dict
        Additional keyword arguments to pass to sns.regplot

    Returns
    -------
    matplotlib.axes.Axes :
        The axis containing the plot.

    """

    ax = ax if ax is not None else plt.gca()
    targets, counts = np.unique(intensities[Features.TARGET].values, return_counts=True)
    this = pd.Series(counts, index=targets)
    tmp = pd.concat([this, other], join='outer', axis=1, sort=True)
    tmp.columns = ['result', 'comparison']

    # this seaborn function does not accept an axis, but we can make sure it draws to the right
    # place by setting the current axis
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', FutureWarning)
        sns.regplot('comparison', 'result', data=tmp, ax=ax, **kwargs)

    # calculate the correlation coefficient
    # mask nans
    tmp = tmp.dropna()
    r, _ = pearsonr(tmp['result'], tmp['comparison'])

    plt.text(0.1, 0.85, f'r = {r:.3}', transform=ax.transAxes, fontsize=14)

    if equalize_axes:
        ax = equalize_axes_extents(ax)
    return ax
