<reponame>darrenreger/zEpid
import warnings
import numpy as np
import pandas as pd
from scipy.stats import norm
import statsmodels.api as sm
import statsmodels.formula.api as smf
from statsmodels.genmod.families import links
from statsmodels.nonparametric.smoothers_lowess import lowess
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.ticker as mticker


class EffectMeasurePlot:
    """Used to generate effect measure (AKA forest) plots. Estimates and confidence intervals are plotted in a diagram
    on the left and a table of the corresponding estimates is provided in the same plot. See the Graphics page on
    ReadTheDocs examples of the plots

    Parameters
    --------------
    label : list
        List of labels to use for y-axis
    effect_measure : list
        List of numbers for point estimates to plot. If point estimate has trailing zeroes,
        input as a character object rather than a float
    lcl : list
        List of numbers for upper confidence limits to plot. If point estimate has trailing
        zeroes, input as a character object rather than a float
    ucl : list
        List of numbers for upper confidence limits to plot. If point estimate has
        trailing zeroes, input as a character object rather than a float

    Examples
    -------------
    Setting up the data to plot

    >>> from matplotlib.pyplot as plt
    >>> from zepid.graphics import EffectMeasurePlot
    >>> lab = ['One','Two']
    >>> emm = [1.01,1.31]
    >>> lcl = ['0.90',1.01]  # String allows for trailing zeroes in the table
    >>> ucl = [1.11,1.53]

    Setting up the plot, measure labels, and point colors

    >>> x = EffectMeasurePlot(lab, emm, lcl, ucl)
    >>> x.labels(effectmeasure='RR')  # Changing label of measure
    >>> x.colors(pointcolor='r')  # Changing color of the points

    Generating matplotlib axes object of forest plot

    >>> x.plot(t_adjuster=0.13)
    >>> plt.show()
    """
    def __init__(self, label, effect_measure, lcl, ucl):
        self.df = pd.DataFrame()
        self.df['study'] = label
        self.df['OR'] = effect_measure
        self.df['LCL'] = lcl
        self.df['UCL'] = ucl
        self.df['OR2'] = self.df['OR'].astype(str).astype(float)
        if (all(isinstance(item, float) for item in lcl)) & (all(isinstance(item, float) for item in effect_measure)):
            self.df['LCL_dif'] = self.df['OR'] - self.df['LCL']
        else:
            self.df['LCL_dif'] = (pd.to_numeric(self.df['OR'])) - (pd.to_numeric(self.df['LCL']))
        if (all(isinstance(item, float) for item in ucl)) & (all(isinstance(item, float) for item in effect_measure)):
            self.df['UCL_dif'] = self.df['UCL'] - self.df['OR']
        else:
            self.df['UCL_dif'] = (pd.to_numeric(self.df['UCL'])) - (pd.to_numeric(self.df['OR']))
        self.em = 'OR'
        self.ci = '95% CI'
        self.scale = 'linear'
        self.center = 1
        self.errc = 'dimgrey'
        self.shape = 'd'
        self.pc = 'k'
        self.linec = 'gray'

    def labels(self, **kwargs):
        """Function to change the labels of the outputted table. Additionally, the scale and reference
        value can be changed.

        Parameters
        -------------
        effectmeasure : string, optional
            Changes the effect measure label
        conf_int : string, optional
            Changes the confidence interval label
        scale : string, optional
            Changes the scale to either log or linear
        center : float, integer, optional
            Changes the reference line for the center
        """
        if 'effectmeasure' in kwargs:
            self.em = kwargs['effectmeasure']
        if 'conf_int' in kwargs:
            self.ci = kwargs['conf_int']
        if 'scale' in kwargs:
            self.scale = kwargs['scale']
        if 'center' in kwargs:
            self.center = kwargs['center']

    def colors(self, **kwargs):
        """Function to change colors and shapes.

        Parameters
        ---------------
        errorbarcolor : string, optional
            Changes the error bar colors
        linecolor : string, optional
            Changes the color of the reference line
        pointcolor : string, optional
            Changes the color of the points
        pointshape : string, optional
            Changes the shape of points
        """
        if 'errorbarcolor' in kwargs:
            self.errc = kwargs['errorbarcolor']
        if 'pointshape' in kwargs:
            self.shape = kwargs['pointshape']
        if 'linecolor' in kwargs:
            self.linec = kwargs['linecolor']
        if 'pointcolor' in kwargs:
            self.pc = kwargs['pointcolor']

    def plot(self, figsize=(3, 3), t_adjuster=0.01, decimal=3, size=3, max_value=None, min_value=None):
        """Generates the matplotlib effect measure plot with the default or specified attributes.
        The following variables can be used to further fine-tune the effect measure plot

        Parameters
        -----------------
        figsize : tuple, optional
            Adjust the size of the figure. Syntax is same as matplotlib `figsize`
        t_adjuster : float, optional
            Used to refine alignment of the table with the line graphs. When generate plots, trial and error for this
            value are usually necessary. I haven't come up with an algorithm to determine this yet...
        decimal : integer, optional
            Number of decimal places to display in the table
        size : integer,
            Option to adjust the size of the lines and points in the plot
        max_value : float, optional
            Maximum value of x-axis scale. Default is None, which automatically determines max value
        min_value : float, optional
            Minimum value of x-axis scale. Default is None, which automatically determines min value

        Returns
        ---------
        matplotlib axes
        """
        tval = []
        ytick = []
        for i in range(len(self.df)):
            if not np.isnan(self.df['OR2'][i]):
                if ((isinstance(self.df['OR'][i], float)) & (isinstance(self.df['LCL'][i], float)) &
                        (isinstance(self.df['UCL'][i], float))):
                    tval.append([round(self.df['OR2'][i], decimal), (
                                '(' + str(round(self.df['LCL'][i], decimal)) + ', ' +
                                str(round(self.df['UCL'][i], decimal)) + ')')])
                else:
                    tval.append(
                        [self.df['OR'][i], ('(' + str(self.df['LCL'][i]) + ', ' + str(self.df['UCL'][i]) + ')')])
                ytick.append(i)
            else:
                tval.append([' ', ' '])
                ytick.append(i)
        if max_value is None:
            if pd.to_numeric(self.df['UCL']).max() < 1:
                maxi = round(((pd.to_numeric(self.df['UCL'])).max() + 0.05),
                             2)  # setting x-axis maximum for UCL less than 1
            if (pd.to_numeric(self.df['UCL']).max() < 9) and (pd.to_numeric(self.df['UCL']).max() >= 1):
                maxi = round(((pd.to_numeric(self.df['UCL'])).max() + 1),
                             0)  # setting x-axis maximum for UCL less than 10
            if pd.to_numeric(self.df['UCL']).max() > 9:
                maxi = round(((pd.to_numeric(self.df['UCL'])).max() + 10),
                             0)  # setting x-axis maximum for UCL less than 100
        else:
            maxi = max_value
        if min_value is None:
            if pd.to_numeric(self.df['LCL']).min() > 0:
                mini = round(((pd.to_numeric(self.df['LCL'])).min() - 0.1), 1)  # setting x-axis minimum
            if pd.to_numeric(self.df['LCL']).min() < 0:
                mini = round(((pd.to_numeric(self.df['LCL'])).min() - 0.05), 2)  # setting x-axis minimum
        else:
            mini = min_value
        plt.figure(figsize=figsize)  # blank figure
        gspec = gridspec.GridSpec(1, 6)  # sets up grid
        plot = plt.subplot(gspec[0, 0:4])  # plot of data
        tabl = plt.subplot(gspec[0, 4:])  # table of OR & CI
        plot.set_ylim(-1, (len(self.df)))  # spacing out y-axis properly
        if self.scale == 'log':
            try:
                plot.set_xscale('log')
            except:
                raise ValueError('For the log scale, all values must be positive')
        plot.axvline(self.center, color=self.linec, zorder=1)
        plot.errorbar(self.df.OR2, self.df.index, xerr=[self.df.LCL_dif, self.df.UCL_dif], marker='None', zorder=2,
                      ecolor=self.errc, elinewidth=(size / size), linewidth=0)
        plot.scatter(self.df.OR2, self.df.index, c=self.pc, s=(size * 25), marker=self.shape, zorder=3,
                     edgecolors='None')
        plot.xaxis.set_ticks_position('bottom')
        plot.yaxis.set_ticks_position('left')
        plot.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())
        plot.get_xaxis().set_minor_formatter(matplotlib.ticker.NullFormatter())
        plot.set_yticks(ytick)
        plot.set_xlim([mini, maxi])
        plot.set_xticks([mini, self.center, maxi])
        plot.set_xticklabels([mini, self.center, maxi])
        plot.set_yticklabels(self.df.study)
        plot.yaxis.set_ticks_position('none')
        plot.invert_yaxis()  # invert y-axis to align values properly with table
        tb = tabl.table(cellText=tval, cellLoc='center', loc='right', colLabels=[self.em, self.ci],
                        bbox=[0, t_adjuster, 1, 1])
        tabl.axis('off')
        tb.auto_set_font_size(False)
        tb.set_fontsize(12)
        for key, cell in tb.get_celld().items():
            cell.set_linewidth(0)
        return plot


def functional_form_plot(df, outcome, var, f_form=None, outcome_type='binary', discrete=False, link_dist=None,
                         loess=True, loess_value=0.25, legend=True, model_results=True, points=False):
    """Creates a functional form plot to aid in functional form assessment for continuous/discrete variables. Plots can
    be created for binary and continuous outcomes. Default options are set to create a functional form plot for a
    binary outcome. To convert to a continuous outcome, `outcome_type` needs to be changed, in addition to
    the `link_dist`

    Parameters
    ------------
    df : DataFrame
        Pandas dataframe that contains the variables of interest
    outcome : string
        Column name of the outcome variable of interest
    var : string
        Column name of the variable of interest for the functional form assessment
    f_form : string, optional
        Regression equation of the functional form to assess. Default is None, which will produce a linear functional
        form. Input the regression equation following the `patsy` syntax. For example, 'var + var_sq'
    outcome_type : string, optional
        Variable type of the outcome variable. Currently, only binary and continuous variables are supported. Default
        is 'binary'
    link_dist : optional
        Link and distribution for the GLM regression equation. Change this to any valid link and distributions
        supported by `statsmodels`. Default is `None`, which defaults to logistic regression
    loess_value : float, optional
        Fraction of observations to use to fit the LOESS curve. This may need to be changed iteratively to determine
        which percent works best for the data. Default is 0.25
    legend : bool, optional
        Turn the legend on or off. Default is True, displaying the legend in the graph
    model_results : bool, optional
        Whether to produce the model results. Default is True, which provides model results
    loess : bool, optional
        Whether to plot the LOESS curve along with the functional form. Default is True
    points : bool, optional
        Whether to plot the data points, where size is relative to the number of observations. Default is False
    discrete : bool, optional
        If your data is truly continuous, leave setting to auto bin the dat. Will automatically bin observations into
        categories for fitting a model with a disjoint indicator. If data is discrete, you can set this to `True` to
        use the actual values for the disjoint indicator. If you get a perfect `SeparationError` from `statsmodels`, it
        means you might have to reshift your categories.

    Returns
    -----------
    matplotlib axes
        Returns a matplotlib graph with a LOESS line (dashed red-line), regression line (sold blue-line), and
        confidence interval (shaded blue)

    Examples
    ------------
    Setting up the environment

    >>> from zepid import load_sample_data
    >>> from zepid.graphics import functional_form_plot
    >>> import matplotlib.pyplot as plt
    >>> df = load_sample_data(timevary=False)
    >>> df['cd4_sq'] = df['cd4']**2

    Creating a functional form plot for a linear functional form

    >>> functional_form_plot(df, outcome='dead', var='cd4')
    >>> plt.show()

    Functional form assessment for a quadractic functional form

    >>> functional_form_plot(df, outcome='dead', var='cd4', f_form='cd4 + cd4_sq')
    >>> plt.show()

    Varying the LOESS value (increased LOESS value to smooth LOESS curve further)

    >>> functional_form_plot(df, outcome='dead', var='cd4', loess_value=0.5)
    >>> plt.show()

    Removing the LOESS curve and the legend from the plot

    >>> functional_form_plot(df, outcome='dead', var='cd4', loess=False, legend=False)
    >>> plt.show()

    Adding summary points to the plot. Points are grouped together and their size reflects their relative n

    >>> functional_form_plot(df, outcome='dead', var='cd4', loess=False, legend=False, points=True)
    >>> plt.show()

    Functional form assessment for a discrete variable (age)

    >>> functional_form_plot(df, outcome='dead', var='age0', discrete=True)
    >>> plt.show()
    """
    # Copying out the dataframe to a new object we will manipulate a bit
    rf = df.copy()
    rf = rf.dropna(subset=[var, outcome]).sort_values(by=[var, outcome]).reset_index()
    if rf.shape[0] != df.shape[0]:
        warnings.warn('Warning: missing observations of model variables are dropped. ' +
                      str(int(df.shape[0] - rf.shape[0])) +
                      ' observations were dropped from the functional form assessment')

    # Functional form for the model
    if f_form is None:
        f_form = var
    else:
        pass

    # Generating Models
    if outcome_type == 'binary':
        if link_dist is None:
            link_dist = sm.families.family.Binomial()
        else:
            pass
    elif outcome_type == 'continuous':
        if link_dist is None:
            link_dist = sm.families.family.Gaussian(sm.families.links.identity)
        else:
            pass
    else:
        raise ValueError('Only binary or continuous outcomes are currently supported')

    # Generating LOESS or points if requested
    ax = plt.gca()
    if loess or points:
        if outcome_type == 'binary':
            if discrete is False:
                # Binning continuous variable into categories to get "General" functional form
                categories = int((np.max(rf[var]) - np.min(rf[var])) / 5)
                if model_results:
                    print('''A total of ''' + str(categories) + ''' categories were created. If you would like to
                            influence  the number of categories the spline is fit to, do the following\n\tIncrease:
                            multiply by constant >1\n\tDecrease: multiply by contast <1 and >0''')
                rf['vbin'] = pd.qcut(rf[var], q=categories, duplicates='drop').cat.codes
                djm = smf.glm(outcome + '~ C(vbin)', rf, family=link_dist).fit()
            else:
                djm = smf.glm(outcome + '~ C(' + var + ')', rf, family=link_dist).fit()
            djf = djm.get_prediction(rf).summary_frame()
            dj = pd.concat([rf, djf], axis=1)
            dj.sort_values(var, inplace=True)
            if points:
                pf = dj.groupby(by=[var, 'mean']).count().reset_index()
                ax.scatter(pf[var], pf['mean'], s=[100 * (n / np.max(pf[var])) for n in pf[var]],
                           color='gray', label='Data point')
            if loess:
                yl = lowess(list(dj['mean']), list(dj[var]), frac=loess_value)
                lowess_x = list(zip(*yl))[0]
                lowess_y = list(zip(*yl))[1]
                ax.plot(lowess_x, lowess_y, '--', color='red', linewidth=1, label='LOESS')
        if outcome_type == 'continuous':
            if points:
                pf = rf.groupby(by=[var, outcome]).count().reset_index()
                ax.scatter(pf[var], pf[outcome], color='gray', label='Data point')
            if loess:
                yl = lowess(list(rf[outcome]), list(rf[var]), frac=loess_value)
                lowess_x = list(zip(*yl))[0]
                lowess_y = list(zip(*yl))[1]
                ax.plot(lowess_x, lowess_y, '--', color='red', linewidth=1, label='LOESS')

    # Functional form model fitting
    ffm = smf.glm(outcome + ' ~ ' + f_form, rf, family=link_dist).fit()
    if model_results is True:
        print(ffm.summary())
        print('AIC: ', ffm.aic)
        print('BIC: ', ffm.bic)
    fff = ffm.get_prediction(rf).summary_frame()
    ff = pd.concat([rf, fff], axis=1)
    ff.sort_values(var, inplace=True)

    # Generating plot for functional form
    ax.fill_between(ff[var], ff['mean_ci_upper'], ff['mean_ci_lower'], alpha=0.1, color='blue', label='95% CI')
    ax.plot(ff[var], ff['mean'], '-', color='blue', label='Regression')
    ax.set_xlabel(var)
    ax.set_ylabel('Outcome')
    if legend is True:
        ax.legend()
    return ax


def pvalue_plot(point, sd, color='b', fill=True, null=0, alpha=None):
    """Creates a plot of the p-value distribution based on a point estimate and standard deviation.
    I find this plot to be useful to explain p-values and how much evidence weight you have in a
    specific value. I think it is useful to explain what exactly a p-value tells you. Note that this
    plot only works for measures on a linear scale (i.e. it will plot exp(log(RR)) incorrectly). It also
    helps to understand what exactly confidence intervals are telling you. These  plots are based on
    Rothman's Epidemiology 2nd Edition pg 152-153 and explained more fully within.

    Parameters
    -------------
    point : float
        Point estimate. Must be on a linear scale (RD / log(RR))
    sd : float
        Standard error of the estimate. Must for linear scale (SE(RD) / SE(log(RR)))
    color : str, optional
        Change color of p-value plot
    fill : bool, optional
        Hhether to fill the curve under the p-value distribution. Setting to False prevents fill
    null : float, integer, optional
        The main value to compare to. The default is zero
    alpha : float, optional
        Whether to draw a line designating significance level area. Default is None, which does not draw this line.
        Generally, would be set to 0.05 to correspond to the widely used alpha of 0.05

    Returns
    -----------
    matplotlib axes

    Examples
    -----------
    Setting up the environment

    >>> from zepid.graphics import pvalue_plot
    >>> import matplotlib.pyplot as plt

    Basic P-value plot

    >>> pvalue_plot(point=-0.1, sd=0.061, color='r')
    >>> plt.show()

    P-value plot with significance line drawn at 'alpha'

    >>> pvalue_plot(point=-0.1, sd=0.061, color='r', alpha=0.025)
    >>> plt.show()

    P-value plot with different comparison value

    >>> pvalue_plot(point=-0.1, sd=0.061, color='r', null=0.1)
    >>> plt.show()

    References
    ----------
    <NAME>. (2012). Epidemiology: an introduction. Oxford university press.
    """
    if point <= null:
        lower = (point - 3 * sd)
        if (point + 3 * sd) < 0:
            upper = point + 3 * sd
        else:
            upper = null + 3 * sd
    if point > null:
        upper = (point + 3 * sd)
        if (point - 3 * sd) > 0:
            lower = null - 3 * sd
        else:
            lower = point - 3 * sd
    ax = plt.gca()
    x1 = np.linspace(lower, point, 100)
    x2 = np.linspace(point, upper, 100)
    ax.plot(x2, 2 * (1 - norm.cdf(x2, loc=point, scale=sd)), c=color)
    ax.plot(x1, 2 * norm.cdf(x1, loc=point, scale=sd), c=color)
    if fill == True:
        ax.fill_between(x2, 2 * (1 - norm.cdf(x2, loc=point, scale=sd)), color=color, alpha=0.2)
        ax.fill_between(x1, 2 * norm.cdf(x1, loc=point, scale=sd), color=color, alpha=0.2)
    ax.vlines(null, 0, 1, colors='k')
    ax.set_xlim([lower, upper])
    ax.set_ylim([0, 1])
    ax.set_ylabel('P-value')
    if alpha is not None:
        ax.hlines(alpha, lower, upper)
    return ax


def spaghetti_plot(df, idvar, variable, time):
    """Create a spaghetti plot by an ID variable. A spaghetti plot can be useful for visualizing
    trends or looking at longitudinal data patterns for individuals all at once.

    Parameters
    ------------
    df : DataFrame
        Pandas dataframe containing variables of interest
    idvar : str
        ID variable for observations. This should indicate the group or individual followed over the time variable
    variable : str
        Variable of interest to see how it varies over time
    time : str
        Time or other variable in which the variable variation occurs

    Returns
    -----------
    matplotlib axes

    Examples
    -----------
    Setting up the environment

    >>> from zepid import load_sample_data
    >>> from zepid.graphics import spaghetti_plot
    >>> df = load_sample_data(timevary=True)

    Generating spaghetti plot for changing CD4 count

    >>> spaghetti_plot(df, idvar='id', variable='cd4', time='enter')
    >>> plt.show()
    """
    ax = plt.gca()
    for i in df[idvar].unique():
        s = df.loc[df[idvar] == i].copy()
        s.sort_values(time, ascending=False)
        ax.plot(s[time], s[variable])
    ax.set_xlabel(time)
    ax.set_ylabel(variable)
    return ax


def roc(df, true, threshold, youden_index=True):
    """Generate a Receiver Operator Curve from true values and predicted probabilities. Youden's Index can also be
    calculated. Youden's index is calculated as

    .. math::

        P_{Yi} = max(Se_i + Sp_i - 1)

    Parameters
    ------------
    df : DataFrame
        Pandas dataframe containing variables of interest
    true : str
        True designation of the outcome (1, 0)
    threshold : str
        Predicted probabilities for the outcome
    youden_index : bool, optional
        Whether to calculate Youden's index. Youden's index maximizes both sensitivity and specificity. The formula
        finds the maximum of (sensitivity + specificity - 1)

    Returns
    -----------
    matplotlib axes

    Examples
    --------
    Creating a dataframe with true disease status (`'d'`) and predicted probability of the outcome (`'p'`)

    >>> import pandas as pd
    >>> import matplotlib.pyplot as plt
    >>> from zepid.graphics import roc
    >>> df = pd.DataFrame()
    >>> df['d'] = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
    >>> df['p'] = [0.1, 0.15, 0.1, 0.7, 0.5, 0.9, 0.95, 0.5, 0.4, 0.8, 0.99, 0.99, 0.89, 0.95]

    Creating ROC curve

    >>> roc(df, true='d', threshold='p', youden_index=False)
    >>> plt.show()
    """
    sens = []
    fpr = []
    thresh = []
    tf = df[[threshold, true]].copy()
    if tf.isnull().values.sum() != 0:
        raise ValueError('ROC curve cannot handle missing data for probability or true values')

    # Getting all possible cutpoints
    values = (list(np.unique(tf[threshold])))
    values = [float(np.min(tf[threshold]) - 0.001)] + values + [float(np.max(tf[threshold]) + 0.001)]
    # Going through all the cutpoints and calculating Sensitivity and 1-Specificity
    for v in list(reversed(values)):
        thresh.append(v)
        prediction = np.where(tf[threshold] >= v, 1, 0)
        se = prediction[tf[true] == 1].mean()
        sens.append(se)
        sp = prediction[tf[true] == 0].mean()
        fpr.append(sp)

    # If requested, calculate Youden's Index
    if youden_index is True:
        spec = [1 - i for i in fpr]
        youdens = []
        for i, j in zip(sens, spec):
            youdens.append(i + j - 1)
        ind = np.argmax(youdens)
        print('----------------------------------------------------------------------')
        print("Youden's Index: ", thresh[ind])
        print("Predictive values at Youden's Index")
        print("\tSensitivity: ", sens[ind])
        print("\tSpecificity: ", spec[ind])
        print('----------------------------------------------------------------------')

    # Creating ROC plot
    ax = plt.gca()
    ax.plot(fpr, sens, color='blue')
    ax.plot([0, 1], [0, 1], color='gray', linestyle='--')
    if youden_index is True:
        ax.text(0.65, 0.35, "Youden's Index:\n      " + str(round(thresh[ind], 5)))
    ax.set_xlim([-0.01, 1.01])
    ax.set_ylim([-0.01, 1.01])
    ax.set_ylabel('Sensitivity')
    ax.set_xlabel('1 -Specificity')
    return ax


def dynamic_risk_plot(risk_exposed, risk_unexposed, measure='RD', loess=True, loess_value=0.25, point_color='darkblue',
                      line_color='b', scale='linear'):
    """Creates a plot of how the risk difference or risk ratio changes over time with survival data. See the
    references for an example of this plot. Input data should be pandas Series indexed by 'timeline' where 'timeline'
    is the time corresponding to the risk estimate

    Parameters
    --------------
    risk_exposed : Series
        Pandas Series with the probability of the outcome among the exposed group. Index by 'timeline' where 'timeline'
        is the time. If you directly output the ``1 - survival_function_`` from lifelines.KaplanMeierFitter(), this
        should create a valid input
    risk_unexposed : Series
        Pandas Series with the probability of the outcome among the exposed group. Index by 'timeline' where 'timeline'
        is the time
    measure : str, optional
        Whether to generate the risk difference (RD) or risk ratio (RR). Default is 'RD'
    loess : bool, optional
        Whether to generate LOESS curve fit to the calculated points. Default is True
    loess_value : float, optional
        Fraction of values to fit LOESS curve to. Default is 0.25
    point_color : str, optional
        Color of the points
    line_color : str, optional
        Color of the LOESS line generated and plotted
    scale : str, optional
        Change the y-axis scale. Options are 'linear' (default), 'log', 'log-transform'. 'log' and 'log-transform' is
        only a valid option for Risk Ratio plots

    Returns
    -----------
    matplotlib axes

    Examples
    --------
    See graphics documentation or causal documentation for a detailed example.

    >>> import matplotlib.pyplot as plt
    >>> from zepid.graphics import dynamic_risk_plot

    >>> dynamic_risk_plot(a, b, loess=True)
    >>> plt.show()

    References
    ----------
    <NAME>, et al. (2014). Estimation of the standardized risk difference and ratio in a competing risks framework:
    application to injection drug use and progression to AIDS after initiation of antiretroviral therapy.
    AJE, 181(4), 238-245.
    """
    re = risk_exposed.drop_duplicates(keep='first').iloc[:, 0].rename('exposed').reset_index()
    ru = risk_unexposed.drop_duplicates(keep='first').iloc[:, 0].rename('unexposed').reset_index()
    re.timeline = np.round(re.timeline * 100000).astype(int) # This avoids a merge issue on floats
    ru.timeline = np.round(ru.timeline * 100000).astype(int)
    r = pd.merge(re, ru, how='outer', left_on='timeline', right_on='timeline').sort_values(by='timeline')
    r.timeline /= 100000
    r.ffill(inplace=True)
    if measure == 'RD':
        r['m'] = r['exposed'] - r['unexposed']
    elif measure == 'RR':
        r['m'] = r['exposed'] / r['unexposed']
        if scale == 'log-transform':
            r['m'] = np.log(r['m'])
    else:
        raise ValueError('Only "RD" and "RR" are currently supported')

    # Generating the plot
    ax = plt.gca()
    ax.plot(r['timeline'], r['m'], 'o', c=point_color)
    if loess is True:
        l = lowess(list(r['m']), list(r['timeline']), frac=loess_value)
        lowess_x = list(zip(*l))[0]
        lowess_y = list(zip(*l))[1]
        ax.plot(lowess_x, lowess_y, '-', c=line_color, linewidth=4)
    if measure == 'RD':
        ax.hlines(0, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
        ax.set_ylabel('Risk Difference')
    if measure == 'RR':
        if scale == 'log-transform':
            ax.hlines(0, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
            ax.set_ylabel('ln(Risk Ratio)')
        elif scale == 'log':
            ax.set_ylabel('Risk Ratio')
            ax.set_yscale('log')
            ax.yaxis.set_major_formatter(mticker.ScalarFormatter())
            ax.yaxis.get_major_formatter().set_scientific(False)
            ax.yaxis.get_major_formatter().set_useOffset(False)
            ax.hlines(1, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
        else:
            ax.hlines(1, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
    ax.set_xlabel('Time')
    ax.set_xlim([0, np.max(r['timeline']) + 0.5])
    return ax


def labbe_plot(r1=None, r0=None, scale='both', additive_tuner=12, multiplicative_tuner=12, figsize=(7, 4), **plot_kwargs):
    """L'Abbe plots are useful for summarizing measure modification on the difference or ratio scale. Primarily
    invented for meta-analysis usage, these plots display risk differences (or ratios) by their individual risks
    by an exposure. I find them most useful for a visualization of why if there is an association and there is no
    modfication on one scale (additive or multiplicative), then there must be modification on the other scale.

    Parameters
    ----------
    r1 : float, list, optional
        Single probability or a list of probabilities when exposure is 1. Default is None, which does not display points
    r0 : float, list, optional
        Single probability or a list of probabilities when exposure is 0. Default is None, which does not display points
    scale : str, optional
        Which scale to plot. The default is 'both', which generates side-by-side plots of additive scale and
        multiplicative scale. Other options are; 'additive' to display the additive plot, and 'multiplicative' to
        display the multiplicative plot
    additive_tuner : int, optional
        Optional parameter to change the number of lines displayed in the additive L'Abbe plot. Higher integer
        produces more reference lines
    multiplicative_tuner : int, optional
        Optional parameter to change the number of lines displayed in the multiplicative L'Abbe plot. Higher integer
        produces more reference lines
    figsize : set, optional
        Optional parameter to change the L'Abbe plot size. Only changes the plot size when scale='both'
    **plot_kwargs : optional
        Optional keyword arguments for matplotlib. kwargs will pass matplotlib.pyploy.plot kwargs are accepted. See
        matplotlib 'plot()' function documentation for further details

    Returns
    -------
    matplotlib axes

    Examples
    --------
    Setting up environment

    >>> import matplotlib.pyplot as plt
    >>> from zepid.graphics import labbe_plot

    Creating a blank plot

    >>> labbe_plot()
    >>> plt.show()

    Adding customized points to the plot

    >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='additive', color='r', marker='D', markersize=10, linestyle='')
    >>> plt.show()

    Only returning the additive plot

    >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='additive', markersize=10)
    >>> plt.show()

    Only returning the multiplicative plot

    >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='multiplicative', markersize=10)
    >>> plt.show()
    """
    if r1 is not None or r0 is not None:
        if len(list(r1)) != len(list(r0)):
            raise ValueError('The length of `r1` must be the same as `r0`')

    ya0 = np.linspace(0.0001, 0.9999, 12)

    if scale == 'both':
        fig, ax = plt.subplots(1, 2, figsize=figsize)
        ax[0].plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-1, 1, additive_tuner):
            ax[0].plot(ya0, ya0 + i, color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax[0].plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax[0].plot(r0, r1, 'o', **plot_kwargs)
        ax[0].set_xlim([0, 1])
        ax[0].set_ylim([0, 1])
        ax[0].set_yticks([0, 1])
        ax[0].set_xticks([0, 1])
        ax[0].set_xlabel("$P(Y|A=0)$")
        ax[0].set_ylabel("$P(Y|A=1)$")
        ax[0].set_title("Additive")

        ax[1].plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-3.5, 3.5, multiplicative_tuner):
            ax[1].plot(ya0, np.exp(np.log(ya0) + i), color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax[1].plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax[1].plot(r0, r1, 'o', **plot_kwargs)
        ax[1].set_xlim([0, 1])
        ax[1].set_ylim([0, 1])
        ax[1].set_yticks([])
        ax[1].set_xticks([0, 1])
        ax[1].set_xlabel("$P(Y|A=0)$")
        ax[1].set_title("Multiplicative")

    elif scale == 'additive':
        ax = plt.gca()
        ax.plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-1, 1, additive_tuner):
            ax.plot(ya0, ya0 + i, color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax.plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax.plot(r0, r1, 'o', **plot_kwargs)
        ax.set_xlim([0, 1])
        ax.set_ylim([0, 1])
        ax.set_yticks([0, 1])
        ax.set_xticks([0, 1])
        ax.set_xlabel("$P(Y|A=0)$")
        ax.set_ylabel("$P(Y|A=1)$")
        ax.set_title("Additive")

    elif scale == 'multiplicative':
        ax = plt.gca()
        ax.plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-3.2, 3.2, multiplicative_tuner):
            plt.plot(ya0, np.exp(np.log(ya0) + i), color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax.plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax.plot(r0, r1, 'o', **plot_kwargs)
        ax.set_xlim([0, 1])
        ax.set_ylim([0, 1])
        ax.set_yticks([0, 1])
        ax.set_xticks([0, 1])
        ax.set_xlabel("$P(Y|A=0)$")
        ax.set_ylabel("$P(Y|A=1)$")
        ax.set_title("Multiplicative")

    else:
        raise ValueError("`scale` must be either 'additive', 'multplicative', or 'both'")

    return ax

import warnings
import numpy as np
import pandas as pd
from scipy.stats import norm
import statsmodels.api as sm
import statsmodels.formula.api as smf
from statsmodels.genmod.families import links
from statsmodels.nonparametric.smoothers_lowess import lowess
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.ticker as mticker


class EffectMeasurePlot:
    """Used to generate effect measure (AKA forest) plots. Estimates and confidence intervals are plotted in a diagram
    on the left and a table of the corresponding estimates is provided in the same plot. See the Graphics page on
    ReadTheDocs examples of the plots

    Parameters
    --------------
    label : list
        List of labels to use for y-axis
    effect_measure : list
        List of numbers for point estimates to plot. If point estimate has trailing zeroes,
        input as a character object rather than a float
    lcl : list
        List of numbers for upper confidence limits to plot. If point estimate has trailing
        zeroes, input as a character object rather than a float
    ucl : list
        List of numbers for upper confidence limits to plot. If point estimate has
        trailing zeroes, input as a character object rather than a float

    Examples
    -------------
    Setting up the data to plot

    >>> from matplotlib.pyplot as plt
    >>> from zepid.graphics import EffectMeasurePlot
    >>> lab = ['One','Two']
    >>> emm = [1.01,1.31]
    >>> lcl = ['0.90',1.01]  # String allows for trailing zeroes in the table
    >>> ucl = [1.11,1.53]

    Setting up the plot, measure labels, and point colors

    >>> x = EffectMeasurePlot(lab, emm, lcl, ucl)
    >>> x.labels(effectmeasure='RR')  # Changing label of measure
    >>> x.colors(pointcolor='r')  # Changing color of the points

    Generating matplotlib axes object of forest plot

    >>> x.plot(t_adjuster=0.13)
    >>> plt.show()
    """
    def __init__(self, label, effect_measure, lcl, ucl):
        self.df = pd.DataFrame()
        self.df['study'] = label
        self.df['OR'] = effect_measure
        self.df['LCL'] = lcl
        self.df['UCL'] = ucl
        self.df['OR2'] = self.df['OR'].astype(str).astype(float)
        if (all(isinstance(item, float) for item in lcl)) & (all(isinstance(item, float) for item in effect_measure)):
            self.df['LCL_dif'] = self.df['OR'] - self.df['LCL']
        else:
            self.df['LCL_dif'] = (pd.to_numeric(self.df['OR'])) - (pd.to_numeric(self.df['LCL']))
        if (all(isinstance(item, float) for item in ucl)) & (all(isinstance(item, float) for item in effect_measure)):
            self.df['UCL_dif'] = self.df['UCL'] - self.df['OR']
        else:
            self.df['UCL_dif'] = (pd.to_numeric(self.df['UCL'])) - (pd.to_numeric(self.df['OR']))
        self.em = 'OR'
        self.ci = '95% CI'
        self.scale = 'linear'
        self.center = 1
        self.errc = 'dimgrey'
        self.shape = 'd'
        self.pc = 'k'
        self.linec = 'gray'

    def labels(self, **kwargs):
        """Function to change the labels of the outputted table. Additionally, the scale and reference
        value can be changed.

        Parameters
        -------------
        effectmeasure : string, optional
            Changes the effect measure label
        conf_int : string, optional
            Changes the confidence interval label
        scale : string, optional
            Changes the scale to either log or linear
        center : float, integer, optional
            Changes the reference line for the center
        """
        if 'effectmeasure' in kwargs:
            self.em = kwargs['effectmeasure']
        if 'conf_int' in kwargs:
            self.ci = kwargs['conf_int']
        if 'scale' in kwargs:
            self.scale = kwargs['scale']
        if 'center' in kwargs:
            self.center = kwargs['center']

    def colors(self, **kwargs):
        """Function to change colors and shapes.

        Parameters
        ---------------
        errorbarcolor : string, optional
            Changes the error bar colors
        linecolor : string, optional
            Changes the color of the reference line
        pointcolor : string, optional
            Changes the color of the points
        pointshape : string, optional
            Changes the shape of points
        """
        if 'errorbarcolor' in kwargs:
            self.errc = kwargs['errorbarcolor']
        if 'pointshape' in kwargs:
            self.shape = kwargs['pointshape']
        if 'linecolor' in kwargs:
            self.linec = kwargs['linecolor']
        if 'pointcolor' in kwargs:
            self.pc = kwargs['pointcolor']

    def plot(self, figsize=(3, 3), t_adjuster=0.01, decimal=3, size=3, max_value=None, min_value=None):
        """Generates the matplotlib effect measure plot with the default or specified attributes.
        The following variables can be used to further fine-tune the effect measure plot

        Parameters
        -----------------
        figsize : tuple, optional
            Adjust the size of the figure. Syntax is same as matplotlib `figsize`
        t_adjuster : float, optional
            Used to refine alignment of the table with the line graphs. When generate plots, trial and error for this
            value are usually necessary. I haven't come up with an algorithm to determine this yet...
        decimal : integer, optional
            Number of decimal places to display in the table
        size : integer,
            Option to adjust the size of the lines and points in the plot
        max_value : float, optional
            Maximum value of x-axis scale. Default is None, which automatically determines max value
        min_value : float, optional
            Minimum value of x-axis scale. Default is None, which automatically determines min value

        Returns
        ---------
        matplotlib axes
        """
        tval = []
        ytick = []
        for i in range(len(self.df)):
            if not np.isnan(self.df['OR2'][i]):
                if ((isinstance(self.df['OR'][i], float)) & (isinstance(self.df['LCL'][i], float)) &
                        (isinstance(self.df['UCL'][i], float))):
                    tval.append([round(self.df['OR2'][i], decimal), (
                                '(' + str(round(self.df['LCL'][i], decimal)) + ', ' +
                                str(round(self.df['UCL'][i], decimal)) + ')')])
                else:
                    tval.append(
                        [self.df['OR'][i], ('(' + str(self.df['LCL'][i]) + ', ' + str(self.df['UCL'][i]) + ')')])
                ytick.append(i)
            else:
                tval.append([' ', ' '])
                ytick.append(i)
        if max_value is None:
            if pd.to_numeric(self.df['UCL']).max() < 1:
                maxi = round(((pd.to_numeric(self.df['UCL'])).max() + 0.05),
                             2)  # setting x-axis maximum for UCL less than 1
            if (pd.to_numeric(self.df['UCL']).max() < 9) and (pd.to_numeric(self.df['UCL']).max() >= 1):
                maxi = round(((pd.to_numeric(self.df['UCL'])).max() + 1),
                             0)  # setting x-axis maximum for UCL less than 10
            if pd.to_numeric(self.df['UCL']).max() > 9:
                maxi = round(((pd.to_numeric(self.df['UCL'])).max() + 10),
                             0)  # setting x-axis maximum for UCL less than 100
        else:
            maxi = max_value
        if min_value is None:
            if pd.to_numeric(self.df['LCL']).min() > 0:
                mini = round(((pd.to_numeric(self.df['LCL'])).min() - 0.1), 1)  # setting x-axis minimum
            if pd.to_numeric(self.df['LCL']).min() < 0:
                mini = round(((pd.to_numeric(self.df['LCL'])).min() - 0.05), 2)  # setting x-axis minimum
        else:
            mini = min_value
        plt.figure(figsize=figsize)  # blank figure
        gspec = gridspec.GridSpec(1, 6)  # sets up grid
        plot = plt.subplot(gspec[0, 0:4])  # plot of data
        tabl = plt.subplot(gspec[0, 4:])  # table of OR & CI
        plot.set_ylim(-1, (len(self.df)))  # spacing out y-axis properly
        if self.scale == 'log':
            try:
                plot.set_xscale('log')
            except:
                raise ValueError('For the log scale, all values must be positive')
        plot.axvline(self.center, color=self.linec, zorder=1)
        plot.errorbar(self.df.OR2, self.df.index, xerr=[self.df.LCL_dif, self.df.UCL_dif], marker='None', zorder=2,
                      ecolor=self.errc, elinewidth=(size / size), linewidth=0)
        plot.scatter(self.df.OR2, self.df.index, c=self.pc, s=(size * 25), marker=self.shape, zorder=3,
                     edgecolors='None')
        plot.xaxis.set_ticks_position('bottom')
        plot.yaxis.set_ticks_position('left')
        plot.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())
        plot.get_xaxis().set_minor_formatter(matplotlib.ticker.NullFormatter())
        plot.set_yticks(ytick)
        plot.set_xlim([mini, maxi])
        plot.set_xticks([mini, self.center, maxi])
        plot.set_xticklabels([mini, self.center, maxi])
        plot.set_yticklabels(self.df.study)
        plot.yaxis.set_ticks_position('none')
        plot.invert_yaxis()  # invert y-axis to align values properly with table
        tb = tabl.table(cellText=tval, cellLoc='center', loc='right', colLabels=[self.em, self.ci],
                        bbox=[0, t_adjuster, 1, 1])
        tabl.axis('off')
        tb.auto_set_font_size(False)
        tb.set_fontsize(12)
        for key, cell in tb.get_celld().items():
            cell.set_linewidth(0)
        return plot


def functional_form_plot(df, outcome, var, f_form=None, outcome_type='binary', discrete=False, link_dist=None,
                         loess=True, loess_value=0.25, legend=True, model_results=True, points=False):
    """Creates a functional form plot to aid in functional form assessment for continuous/discrete variables. Plots can
    be created for binary and continuous outcomes. Default options are set to create a functional form plot for a
    binary outcome. To convert to a continuous outcome, `outcome_type` needs to be changed, in addition to
    the `link_dist`

    Parameters
    ------------
    df : DataFrame
        Pandas dataframe that contains the variables of interest
    outcome : string
        Column name of the outcome variable of interest
    var : string
        Column name of the variable of interest for the functional form assessment
    f_form : string, optional
        Regression equation of the functional form to assess. Default is None, which will produce a linear functional
        form. Input the regression equation following the `patsy` syntax. For example, 'var + var_sq'
    outcome_type : string, optional
        Variable type of the outcome variable. Currently, only binary and continuous variables are supported. Default
        is 'binary'
    link_dist : optional
        Link and distribution for the GLM regression equation. Change this to any valid link and distributions
        supported by `statsmodels`. Default is `None`, which defaults to logistic regression
    loess_value : float, optional
        Fraction of observations to use to fit the LOESS curve. This may need to be changed iteratively to determine
        which percent works best for the data. Default is 0.25
    legend : bool, optional
        Turn the legend on or off. Default is True, displaying the legend in the graph
    model_results : bool, optional
        Whether to produce the model results. Default is True, which provides model results
    loess : bool, optional
        Whether to plot the LOESS curve along with the functional form. Default is True
    points : bool, optional
        Whether to plot the data points, where size is relative to the number of observations. Default is False
    discrete : bool, optional
        If your data is truly continuous, leave setting to auto bin the dat. Will automatically bin observations into
        categories for fitting a model with a disjoint indicator. If data is discrete, you can set this to `True` to
        use the actual values for the disjoint indicator. If you get a perfect `SeparationError` from `statsmodels`, it
        means you might have to reshift your categories.

    Returns
    -----------
    matplotlib axes
        Returns a matplotlib graph with a LOESS line (dashed red-line), regression line (sold blue-line), and
        confidence interval (shaded blue)

    Examples
    ------------
    Setting up the environment

    >>> from zepid import load_sample_data
    >>> from zepid.graphics import functional_form_plot
    >>> import matplotlib.pyplot as plt
    >>> df = load_sample_data(timevary=False)
    >>> df['cd4_sq'] = df['cd4']**2

    Creating a functional form plot for a linear functional form

    >>> functional_form_plot(df, outcome='dead', var='cd4')
    >>> plt.show()

    Functional form assessment for a quadractic functional form

    >>> functional_form_plot(df, outcome='dead', var='cd4', f_form='cd4 + cd4_sq')
    >>> plt.show()

    Varying the LOESS value (increased LOESS value to smooth LOESS curve further)

    >>> functional_form_plot(df, outcome='dead', var='cd4', loess_value=0.5)
    >>> plt.show()

    Removing the LOESS curve and the legend from the plot

    >>> functional_form_plot(df, outcome='dead', var='cd4', loess=False, legend=False)
    >>> plt.show()

    Adding summary points to the plot. Points are grouped together and their size reflects their relative n

    >>> functional_form_plot(df, outcome='dead', var='cd4', loess=False, legend=False, points=True)
    >>> plt.show()

    Functional form assessment for a discrete variable (age)

    >>> functional_form_plot(df, outcome='dead', var='age0', discrete=True)
    >>> plt.show()
    """
    # Copying out the dataframe to a new object we will manipulate a bit
    rf = df.copy()
    rf = rf.dropna(subset=[var, outcome]).sort_values(by=[var, outcome]).reset_index()
    if rf.shape[0] != df.shape[0]:
        warnings.warn('Warning: missing observations of model variables are dropped. ' +
                      str(int(df.shape[0] - rf.shape[0])) +
                      ' observations were dropped from the functional form assessment')

    # Functional form for the model
    if f_form is None:
        f_form = var
    else:
        pass

    # Generating Models
    if outcome_type == 'binary':
        if link_dist is None:
            link_dist = sm.families.family.Binomial()
        else:
            pass
    elif outcome_type == 'continuous':
        if link_dist is None:
            link_dist = sm.families.family.Gaussian(sm.families.links.identity)
        else:
            pass
    else:
        raise ValueError('Only binary or continuous outcomes are currently supported')

    # Generating LOESS or points if requested
    ax = plt.gca()
    if loess or points:
        if outcome_type == 'binary':
            if discrete is False:
                # Binning continuous variable into categories to get "General" functional form
                categories = int((np.max(rf[var]) - np.min(rf[var])) / 5)
                if model_results:
                    print('''A total of ''' + str(categories) + ''' categories were created. If you would like to
                            influence  the number of categories the spline is fit to, do the following\n\tIncrease:
                            multiply by constant >1\n\tDecrease: multiply by contast <1 and >0''')
                rf['vbin'] = pd.qcut(rf[var], q=categories, duplicates='drop').cat.codes
                djm = smf.glm(outcome + '~ C(vbin)', rf, family=link_dist).fit()
            else:
                djm = smf.glm(outcome + '~ C(' + var + ')', rf, family=link_dist).fit()
            djf = djm.get_prediction(rf).summary_frame()
            dj = pd.concat([rf, djf], axis=1)
            dj.sort_values(var, inplace=True)
            if points:
                pf = dj.groupby(by=[var, 'mean']).count().reset_index()
                ax.scatter(pf[var], pf['mean'], s=[100 * (n / np.max(pf[var])) for n in pf[var]],
                           color='gray', label='Data point')
            if loess:
                yl = lowess(list(dj['mean']), list(dj[var]), frac=loess_value)
                lowess_x = list(zip(*yl))[0]
                lowess_y = list(zip(*yl))[1]
                ax.plot(lowess_x, lowess_y, '--', color='red', linewidth=1, label='LOESS')
        if outcome_type == 'continuous':
            if points:
                pf = rf.groupby(by=[var, outcome]).count().reset_index()
                ax.scatter(pf[var], pf[outcome], color='gray', label='Data point')
            if loess:
                yl = lowess(list(rf[outcome]), list(rf[var]), frac=loess_value)
                lowess_x = list(zip(*yl))[0]
                lowess_y = list(zip(*yl))[1]
                ax.plot(lowess_x, lowess_y, '--', color='red', linewidth=1, label='LOESS')

    # Functional form model fitting
    ffm = smf.glm(outcome + ' ~ ' + f_form, rf, family=link_dist).fit()
    if model_results is True:
        print(ffm.summary())
        print('AIC: ', ffm.aic)
        print('BIC: ', ffm.bic)
    fff = ffm.get_prediction(rf).summary_frame()
    ff = pd.concat([rf, fff], axis=1)
    ff.sort_values(var, inplace=True)

    # Generating plot for functional form
    ax.fill_between(ff[var], ff['mean_ci_upper'], ff['mean_ci_lower'], alpha=0.1, color='blue', label='95% CI')
    ax.plot(ff[var], ff['mean'], '-', color='blue', label='Regression')
    ax.set_xlabel(var)
    ax.set_ylabel('Outcome')
    if legend is True:
        ax.legend()
    return ax


def pvalue_plot(point, sd, color='b', fill=True, null=0, alpha=None):
    """Creates a plot of the p-value distribution based on a point estimate and standard deviation.
    I find this plot to be useful to explain p-values and how much evidence weight you have in a
    specific value. I think it is useful to explain what exactly a p-value tells you. Note that this
    plot only works for measures on a linear scale (i.e. it will plot exp(log(RR)) incorrectly). It also
    helps to understand what exactly confidence intervals are telling you. These  plots are based on
    Rothman's Epidemiology 2nd Edition pg 152-153 and explained more fully within.

    Parameters
    -------------
    point : float
        Point estimate. Must be on a linear scale (RD / log(RR))
    sd : float
        Standard error of the estimate. Must for linear scale (SE(RD) / SE(log(RR)))
    color : str, optional
        Change color of p-value plot
    fill : bool, optional
        Hhether to fill the curve under the p-value distribution. Setting to False prevents fill
    null : float, integer, optional
        The main value to compare to. The default is zero
    alpha : float, optional
        Whether to draw a line designating significance level area. Default is None, which does not draw this line.
        Generally, would be set to 0.05 to correspond to the widely used alpha of 0.05

    Returns
    -----------
    matplotlib axes

    Examples
    -----------
    Setting up the environment

    >>> from zepid.graphics import pvalue_plot
    >>> import matplotlib.pyplot as plt

    Basic P-value plot

    >>> pvalue_plot(point=-0.1, sd=0.061, color='r')
    >>> plt.show()

    P-value plot with significance line drawn at 'alpha'

    >>> pvalue_plot(point=-0.1, sd=0.061, color='r', alpha=0.025)
    >>> plt.show()

    P-value plot with different comparison value

    >>> pvalue_plot(point=-0.1, sd=0.061, color='r', null=0.1)
    >>> plt.show()

    References
    ----------
    <NAME>. (2012). Epidemiology: an introduction. Oxford university press.
    """
    if point <= null:
        lower = (point - 3 * sd)
        if (point + 3 * sd) < 0:
            upper = point + 3 * sd
        else:
            upper = null + 3 * sd
    if point > null:
        upper = (point + 3 * sd)
        if (point - 3 * sd) > 0:
            lower = null - 3 * sd
        else:
            lower = point - 3 * sd
    ax = plt.gca()
    x1 = np.linspace(lower, point, 100)
    x2 = np.linspace(point, upper, 100)
    ax.plot(x2, 2 * (1 - norm.cdf(x2, loc=point, scale=sd)), c=color)
    ax.plot(x1, 2 * norm.cdf(x1, loc=point, scale=sd), c=color)
    if fill == True:
        ax.fill_between(x2, 2 * (1 - norm.cdf(x2, loc=point, scale=sd)), color=color, alpha=0.2)
        ax.fill_between(x1, 2 * norm.cdf(x1, loc=point, scale=sd), color=color, alpha=0.2)
    ax.vlines(null, 0, 1, colors='k')
    ax.set_xlim([lower, upper])
    ax.set_ylim([0, 1])
    ax.set_ylabel('P-value')
    if alpha is not None:
        ax.hlines(alpha, lower, upper)
    return ax


def spaghetti_plot(df, idvar, variable, time):
    """Create a spaghetti plot by an ID variable. A spaghetti plot can be useful for visualizing
    trends or looking at longitudinal data patterns for individuals all at once.

    Parameters
    ------------
    df : DataFrame
        Pandas dataframe containing variables of interest
    idvar : str
        ID variable for observations. This should indicate the group or individual followed over the time variable
    variable : str
        Variable of interest to see how it varies over time
    time : str
        Time or other variable in which the variable variation occurs

    Returns
    -----------
    matplotlib axes

    Examples
    -----------
    Setting up the environment

    >>> from zepid import load_sample_data
    >>> from zepid.graphics import spaghetti_plot
    >>> df = load_sample_data(timevary=True)

    Generating spaghetti plot for changing CD4 count

    >>> spaghetti_plot(df, idvar='id', variable='cd4', time='enter')
    >>> plt.show()
    """
    ax = plt.gca()
    for i in df[idvar].unique():
        s = df.loc[df[idvar] == i].copy()
        s.sort_values(time, ascending=False)
        ax.plot(s[time], s[variable])
    ax.set_xlabel(time)
    ax.set_ylabel(variable)
    return ax


def roc(df, true, threshold, youden_index=True):
    """Generate a Receiver Operator Curve from true values and predicted probabilities. Youden's Index can also be
    calculated. Youden's index is calculated as

    .. math::

        P_{Yi} = max(Se_i + Sp_i - 1)

    Parameters
    ------------
    df : DataFrame
        Pandas dataframe containing variables of interest
    true : str
        True designation of the outcome (1, 0)
    threshold : str
        Predicted probabilities for the outcome
    youden_index : bool, optional
        Whether to calculate Youden's index. Youden's index maximizes both sensitivity and specificity. The formula
        finds the maximum of (sensitivity + specificity - 1)

    Returns
    -----------
    matplotlib axes

    Examples
    --------
    Creating a dataframe with true disease status (`'d'`) and predicted probability of the outcome (`'p'`)

    >>> import pandas as pd
    >>> import matplotlib.pyplot as plt
    >>> from zepid.graphics import roc
    >>> df = pd.DataFrame()
    >>> df['d'] = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
    >>> df['p'] = [0.1, 0.15, 0.1, 0.7, 0.5, 0.9, 0.95, 0.5, 0.4, 0.8, 0.99, 0.99, 0.89, 0.95]

    Creating ROC curve

    >>> roc(df, true='d', threshold='p', youden_index=False)
    >>> plt.show()
    """
    sens = []
    fpr = []
    thresh = []
    tf = df[[threshold, true]].copy()
    if tf.isnull().values.sum() != 0:
        raise ValueError('ROC curve cannot handle missing data for probability or true values')

    # Getting all possible cutpoints
    values = (list(np.unique(tf[threshold])))
    values = [float(np.min(tf[threshold]) - 0.001)] + values + [float(np.max(tf[threshold]) + 0.001)]
    # Going through all the cutpoints and calculating Sensitivity and 1-Specificity
    for v in list(reversed(values)):
        thresh.append(v)
        prediction = np.where(tf[threshold] >= v, 1, 0)
        se = prediction[tf[true] == 1].mean()
        sens.append(se)
        sp = prediction[tf[true] == 0].mean()
        fpr.append(sp)

    # If requested, calculate Youden's Index
    if youden_index is True:
        spec = [1 - i for i in fpr]
        youdens = []
        for i, j in zip(sens, spec):
            youdens.append(i + j - 1)
        ind = np.argmax(youdens)
        print('----------------------------------------------------------------------')
        print("Youden's Index: ", thresh[ind])
        print("Predictive values at Youden's Index")
        print("\tSensitivity: ", sens[ind])
        print("\tSpecificity: ", spec[ind])
        print('----------------------------------------------------------------------')

    # Creating ROC plot
    ax = plt.gca()
    ax.plot(fpr, sens, color='blue')
    ax.plot([0, 1], [0, 1], color='gray', linestyle='--')
    if youden_index is True:
        ax.text(0.65, 0.35, "Youden's Index:\n      " + str(round(thresh[ind], 5)))
    ax.set_xlim([-0.01, 1.01])
    ax.set_ylim([-0.01, 1.01])
    ax.set_ylabel('Sensitivity')
    ax.set_xlabel('1 -Specificity')
    return ax


def dynamic_risk_plot(risk_exposed, risk_unexposed, measure='RD', loess=True, loess_value=0.25, point_color='darkblue',
                      line_color='b', scale='linear'):
    """Creates a plot of how the risk difference or risk ratio changes over time with survival data. See the
    references for an example of this plot. Input data should be pandas Series indexed by 'timeline' where 'timeline'
    is the time corresponding to the risk estimate

    Parameters
    --------------
    risk_exposed : Series
        Pandas Series with the probability of the outcome among the exposed group. Index by 'timeline' where 'timeline'
        is the time. If you directly output the ``1 - survival_function_`` from lifelines.KaplanMeierFitter(), this
        should create a valid input
    risk_unexposed : Series
        Pandas Series with the probability of the outcome among the exposed group. Index by 'timeline' where 'timeline'
        is the time
    measure : str, optional
        Whether to generate the risk difference (RD) or risk ratio (RR). Default is 'RD'
    loess : bool, optional
        Whether to generate LOESS curve fit to the calculated points. Default is True
    loess_value : float, optional
        Fraction of values to fit LOESS curve to. Default is 0.25
    point_color : str, optional
        Color of the points
    line_color : str, optional
        Color of the LOESS line generated and plotted
    scale : str, optional
        Change the y-axis scale. Options are 'linear' (default), 'log', 'log-transform'. 'log' and 'log-transform' is
        only a valid option for Risk Ratio plots

    Returns
    -----------
    matplotlib axes

    Examples
    --------
    See graphics documentation or causal documentation for a detailed example.

    >>> import matplotlib.pyplot as plt
    >>> from zepid.graphics import dynamic_risk_plot

    >>> dynamic_risk_plot(a, b, loess=True)
    >>> plt.show()

    References
    ----------
    <NAME>, et al. (2014). Estimation of the standardized risk difference and ratio in a competing risks framework:
    application to injection drug use and progression to AIDS after initiation of antiretroviral therapy.
    AJE, 181(4), 238-245.
    """
    re = risk_exposed.drop_duplicates(keep='first').iloc[:, 0].rename('exposed').reset_index()
    ru = risk_unexposed.drop_duplicates(keep='first').iloc[:, 0].rename('unexposed').reset_index()
    re.timeline = np.round(re.timeline * 100000).astype(int) # This avoids a merge issue on floats
    ru.timeline = np.round(ru.timeline * 100000).astype(int)
    r = pd.merge(re, ru, how='outer', left_on='timeline', right_on='timeline').sort_values(by='timeline')
    r.timeline /= 100000
    r.ffill(inplace=True)
    if measure == 'RD':
        r['m'] = r['exposed'] - r['unexposed']
    elif measure == 'RR':
        r['m'] = r['exposed'] / r['unexposed']
        if scale == 'log-transform':
            r['m'] = np.log(r['m'])
    else:
        raise ValueError('Only "RD" and "RR" are currently supported')

    # Generating the plot
    ax = plt.gca()
    ax.plot(r['timeline'], r['m'], 'o', c=point_color)
    if loess is True:
        l = lowess(list(r['m']), list(r['timeline']), frac=loess_value)
        lowess_x = list(zip(*l))[0]
        lowess_y = list(zip(*l))[1]
        ax.plot(lowess_x, lowess_y, '-', c=line_color, linewidth=4)
    if measure == 'RD':
        ax.hlines(0, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
        ax.set_ylabel('Risk Difference')
    if measure == 'RR':
        if scale == 'log-transform':
            ax.hlines(0, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
            ax.set_ylabel('ln(Risk Ratio)')
        elif scale == 'log':
            ax.set_ylabel('Risk Ratio')
            ax.set_yscale('log')
            ax.yaxis.set_major_formatter(mticker.ScalarFormatter())
            ax.yaxis.get_major_formatter().set_scientific(False)
            ax.yaxis.get_major_formatter().set_useOffset(False)
            ax.hlines(1, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
        else:
            ax.hlines(1, 0, np.max(r['timeline'] + 0.5), linewidth=1.5)
    ax.set_xlabel('Time')
    ax.set_xlim([0, np.max(r['timeline']) + 0.5])
    return ax


def labbe_plot(r1=None, r0=None, scale='both', additive_tuner=12, multiplicative_tuner=12, figsize=(7, 4), **plot_kwargs):
    """L'Abbe plots are useful for summarizing measure modification on the difference or ratio scale. Primarily
    invented for meta-analysis usage, these plots display risk differences (or ratios) by their individual risks
    by an exposure. I find them most useful for a visualization of why if there is an association and there is no
    modfication on one scale (additive or multiplicative), then there must be modification on the other scale.

    Parameters
    ----------
    r1 : float, list, optional
        Single probability or a list of probabilities when exposure is 1. Default is None, which does not display points
    r0 : float, list, optional
        Single probability or a list of probabilities when exposure is 0. Default is None, which does not display points
    scale : str, optional
        Which scale to plot. The default is 'both', which generates side-by-side plots of additive scale and
        multiplicative scale. Other options are; 'additive' to display the additive plot, and 'multiplicative' to
        display the multiplicative plot
    additive_tuner : int, optional
        Optional parameter to change the number of lines displayed in the additive L'Abbe plot. Higher integer
        produces more reference lines
    multiplicative_tuner : int, optional
        Optional parameter to change the number of lines displayed in the multiplicative L'Abbe plot. Higher integer
        produces more reference lines
    figsize : set, optional
        Optional parameter to change the L'Abbe plot size. Only changes the plot size when scale='both'
    **plot_kwargs : optional
        Optional keyword arguments for matplotlib. kwargs will pass matplotlib.pyploy.plot kwargs are accepted. See
        matplotlib 'plot()' function documentation for further details

    Returns
    -------
    matplotlib axes

    Examples
    --------
    Setting up environment

    >>> import matplotlib.pyplot as plt
    >>> from zepid.graphics import labbe_plot

    Creating a blank plot

    >>> labbe_plot()
    >>> plt.show()

    Adding customized points to the plot

    >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='additive', color='r', marker='D', markersize=10, linestyle='')
    >>> plt.show()

    Only returning the additive plot

    >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='additive', markersize=10)
    >>> plt.show()

    Only returning the multiplicative plot

    >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='multiplicative', markersize=10)
    >>> plt.show()
    """
    if r1 is not None or r0 is not None:
        if len(list(r1)) != len(list(r0)):
            raise ValueError('The length of `r1` must be the same as `r0`')

    ya0 = np.linspace(0.0001, 0.9999, 12)

    if scale == 'both':
        fig, ax = plt.subplots(1, 2, figsize=figsize)
        ax[0].plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-1, 1, additive_tuner):
            ax[0].plot(ya0, ya0 + i, color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax[0].plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax[0].plot(r0, r1, 'o', **plot_kwargs)
        ax[0].set_xlim([0, 1])
        ax[0].set_ylim([0, 1])
        ax[0].set_yticks([0, 1])
        ax[0].set_xticks([0, 1])
        ax[0].set_xlabel("$P(Y|A=0)$")
        ax[0].set_ylabel("$P(Y|A=1)$")
        ax[0].set_title("Additive")

        ax[1].plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-3.5, 3.5, multiplicative_tuner):
            ax[1].plot(ya0, np.exp(np.log(ya0) + i), color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax[1].plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax[1].plot(r0, r1, 'o', **plot_kwargs)
        ax[1].set_xlim([0, 1])
        ax[1].set_ylim([0, 1])
        ax[1].set_yticks([])
        ax[1].set_xticks([0, 1])
        ax[1].set_xlabel("$P(Y|A=0)$")
        ax[1].set_title("Multiplicative")

    elif scale == 'additive':
        ax = plt.gca()
        ax.plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-1, 1, additive_tuner):
            ax.plot(ya0, ya0 + i, color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax.plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax.plot(r0, r1, 'o', **plot_kwargs)
        ax.set_xlim([0, 1])
        ax.set_ylim([0, 1])
        ax.set_yticks([0, 1])
        ax.set_xticks([0, 1])
        ax.set_xlabel("$P(Y|A=0)$")
        ax.set_ylabel("$P(Y|A=1)$")
        ax.set_title("Additive")

    elif scale == 'multiplicative':
        ax = plt.gca()
        ax.plot([0, 1], [0, 1], '--', color='gray', linewidth=1)
        for i in np.linspace(-3.2, 3.2, multiplicative_tuner):
            plt.plot(ya0, np.exp(np.log(ya0) + i), color='gray', linewidth=1)
        if r1 is not None and r0 is not None:
            if 'marker' in plot_kwargs:
                ax.plot(r0, r1, **plot_kwargs)
            else:  # When markers are unspecified, point estimates aren't displayed. This avoids
                ax.plot(r0, r1, 'o', **plot_kwargs)
        ax.set_xlim([0, 1])
        ax.set_ylim([0, 1])
        ax.set_yticks([0, 1])
        ax.set_xticks([0, 1])
        ax.set_xlabel("$P(Y|A=0)$")
        ax.set_ylabel("$P(Y|A=1)$")
        ax.set_title("Multiplicative")

    else:
        raise ValueError("`scale` must be either 'additive', 'multplicative', or 'both'")

    return ax
Used to generate effect measure (AKA forest) plots. Estimates and confidence intervals are plotted in a diagram
     on the left and a table of the corresponding estimates is provided in the same plot. See the Graphics page on
     ReadTheDocs examples of the plots
 
     Parameters
     --------------
     label : list
         List of labels to use for y-axis
     effect_measure : list
         List of numbers for point estimates to plot. If point estimate has trailing zeroes,
         input as a character object rather than a float
     lcl : list
         List of numbers for upper confidence limits to plot. If point estimate has trailing
         zeroes, input as a character object rather than a float
     ucl : list
         List of numbers for upper confidence limits to plot. If point estimate has
         trailing zeroes, input as a character object rather than a float
 
     Examples
     -------------
     Setting up the data to plot
 
     >>> from matplotlib.pyplot as plt
     >>> from zepid.graphics import EffectMeasurePlot
     >>> lab = ['One','Two']
     >>> emm = [1.01,1.31]
     >>> lcl = ['0.90',1.01]  # String allows for trailing zeroes in the table
     >>> ucl = [1.11,1.53]
 
     Setting up the plot, measure labels, and point colors
 
     >>> x = EffectMeasurePlot(lab, emm, lcl, ucl)
     >>> x.labels(effectmeasure='RR')  # Changing label of measure
     >>> x.colors(pointcolor='r')  # Changing color of the points
 
     Generating matplotlib axes object of forest plot
 
     >>> x.plot(t_adjuster=0.13)
     >>> plt.show() Function to change the labels of the outputted table. Additionally, the scale and reference
         value can be changed.
 
         Parameters
         -------------
         effectmeasure : string, optional
             Changes the effect measure label
         conf_int : string, optional
             Changes the confidence interval label
         scale : string, optional
             Changes the scale to either log or linear
         center : float, integer, optional
             Changes the reference line for the center Function to change colors and shapes.
 
         Parameters
         ---------------
         errorbarcolor : string, optional
             Changes the error bar colors
         linecolor : string, optional
             Changes the color of the reference line
         pointcolor : string, optional
             Changes the color of the points
         pointshape : string, optional
             Changes the shape of points Generates the matplotlib effect measure plot with the default or specified attributes.
         The following variables can be used to further fine-tune the effect measure plot
 
         Parameters
         -----------------
         figsize : tuple, optional
             Adjust the size of the figure. Syntax is same as matplotlib `figsize`
         t_adjuster : float, optional
             Used to refine alignment of the table with the line graphs. When generate plots, trial and error for this
             value are usually necessary. I haven't come up with an algorithm to determine this yet...
         decimal : integer, optional
             Number of decimal places to display in the table
         size : integer,
             Option to adjust the size of the lines and points in the plot
         max_value : float, optional
             Maximum value of x-axis scale. Default is None, which automatically determines max value
         min_value : float, optional
             Minimum value of x-axis scale. Default is None, which automatically determines min value
 
         Returns
         ---------
         matplotlib axes # setting x-axis maximum for UCL less than 1 # setting x-axis maximum for UCL less than 10 # setting x-axis maximum for UCL less than 100 # setting x-axis minimum # setting x-axis minimum # blank figure # sets up grid # plot of data # table of OR & CI # spacing out y-axis properly # invert y-axis to align values properly with table Creates a functional form plot to aid in functional form assessment for continuous/discrete variables. Plots can
     be created for binary and continuous outcomes. Default options are set to create a functional form plot for a
     binary outcome. To convert to a continuous outcome, `outcome_type` needs to be changed, in addition to
     the `link_dist`
 
     Parameters
     ------------
     df : DataFrame
         Pandas dataframe that contains the variables of interest
     outcome : string
         Column name of the outcome variable of interest
     var : string
         Column name of the variable of interest for the functional form assessment
     f_form : string, optional
         Regression equation of the functional form to assess. Default is None, which will produce a linear functional
         form. Input the regression equation following the `patsy` syntax. For example, 'var + var_sq'
     outcome_type : string, optional
         Variable type of the outcome variable. Currently, only binary and continuous variables are supported. Default
         is 'binary'
     link_dist : optional
         Link and distribution for the GLM regression equation. Change this to any valid link and distributions
         supported by `statsmodels`. Default is `None`, which defaults to logistic regression
     loess_value : float, optional
         Fraction of observations to use to fit the LOESS curve. This may need to be changed iteratively to determine
         which percent works best for the data. Default is 0.25
     legend : bool, optional
         Turn the legend on or off. Default is True, displaying the legend in the graph
     model_results : bool, optional
         Whether to produce the model results. Default is True, which provides model results
     loess : bool, optional
         Whether to plot the LOESS curve along with the functional form. Default is True
     points : bool, optional
         Whether to plot the data points, where size is relative to the number of observations. Default is False
     discrete : bool, optional
         If your data is truly continuous, leave setting to auto bin the dat. Will automatically bin observations into
         categories for fitting a model with a disjoint indicator. If data is discrete, you can set this to `True` to
         use the actual values for the disjoint indicator. If you get a perfect `SeparationError` from `statsmodels`, it
         means you might have to reshift your categories.
 
     Returns
     -----------
     matplotlib axes
         Returns a matplotlib graph with a LOESS line (dashed red-line), regression line (sold blue-line), and
         confidence interval (shaded blue)
 
     Examples
     ------------
     Setting up the environment
 
     >>> from zepid import load_sample_data
     >>> from zepid.graphics import functional_form_plot
     >>> import matplotlib.pyplot as plt
     >>> df = load_sample_data(timevary=False)
     >>> df['cd4_sq'] = df['cd4']**2
 
     Creating a functional form plot for a linear functional form
 
     >>> functional_form_plot(df, outcome='dead', var='cd4')
     >>> plt.show()
 
     Functional form assessment for a quadractic functional form
 
     >>> functional_form_plot(df, outcome='dead', var='cd4', f_form='cd4 + cd4_sq')
     >>> plt.show()
 
     Varying the LOESS value (increased LOESS value to smooth LOESS curve further)
 
     >>> functional_form_plot(df, outcome='dead', var='cd4', loess_value=0.5)
     >>> plt.show()
 
     Removing the LOESS curve and the legend from the plot
 
     >>> functional_form_plot(df, outcome='dead', var='cd4', loess=False, legend=False)
     >>> plt.show()
 
     Adding summary points to the plot. Points are grouped together and their size reflects their relative n
 
     >>> functional_form_plot(df, outcome='dead', var='cd4', loess=False, legend=False, points=True)
     >>> plt.show()
 
     Functional form assessment for a discrete variable (age)
 
     >>> functional_form_plot(df, outcome='dead', var='age0', discrete=True)
     >>> plt.show() # Copying out the dataframe to a new object we will manipulate a bit # Functional form for the model # Generating Models # Generating LOESS or points if requested # Binning continuous variable into categories to get "General" functional form A total of categories were created. If you would like to
                             influence  the number of categories the spline is fit to, do the following\n\tIncrease:
                             multiply by constant >1\n\tDecrease: multiply by contast <1 and >0 # Functional form model fitting # Generating plot for functional form Creates a plot of the p-value distribution based on a point estimate and standard deviation.
     I find this plot to be useful to explain p-values and how much evidence weight you have in a
     specific value. I think it is useful to explain what exactly a p-value tells you. Note that this
     plot only works for measures on a linear scale (i.e. it will plot exp(log(RR)) incorrectly). It also
     helps to understand what exactly confidence intervals are telling you. These  plots are based on
     Rothman's Epidemiology 2nd Edition pg 152-153 and explained more fully within.
 
     Parameters
     -------------
     point : float
         Point estimate. Must be on a linear scale (RD / log(RR))
     sd : float
         Standard error of the estimate. Must for linear scale (SE(RD) / SE(log(RR)))
     color : str, optional
         Change color of p-value plot
     fill : bool, optional
         Hhether to fill the curve under the p-value distribution. Setting to False prevents fill
     null : float, integer, optional
         The main value to compare to. The default is zero
     alpha : float, optional
         Whether to draw a line designating significance level area. Default is None, which does not draw this line.
         Generally, would be set to 0.05 to correspond to the widely used alpha of 0.05
 
     Returns
     -----------
     matplotlib axes
 
     Examples
     -----------
     Setting up the environment
 
     >>> from zepid.graphics import pvalue_plot
     >>> import matplotlib.pyplot as plt
 
     Basic P-value plot
 
     >>> pvalue_plot(point=-0.1, sd=0.061, color='r')
     >>> plt.show()
 
     P-value plot with significance line drawn at 'alpha'
 
     >>> pvalue_plot(point=-0.1, sd=0.061, color='r', alpha=0.025)
     >>> plt.show()
 
     P-value plot with different comparison value
 
     >>> pvalue_plot(point=-0.1, sd=0.061, color='r', null=0.1)
     >>> plt.show()
 
     References
     ----------
     <NAME>. (2012). Epidemiology: an introduction. Oxford university press. Create a spaghetti plot by an ID variable. A spaghetti plot can be useful for visualizing
     trends or looking at longitudinal data patterns for individuals all at once.
 
     Parameters
     ------------
     df : DataFrame
         Pandas dataframe containing variables of interest
     idvar : str
         ID variable for observations. This should indicate the group or individual followed over the time variable
     variable : str
         Variable of interest to see how it varies over time
     time : str
         Time or other variable in which the variable variation occurs
 
     Returns
     -----------
     matplotlib axes
 
     Examples
     -----------
     Setting up the environment
 
     >>> from zepid import load_sample_data
     >>> from zepid.graphics import spaghetti_plot
     >>> df = load_sample_data(timevary=True)
 
     Generating spaghetti plot for changing CD4 count
 
     >>> spaghetti_plot(df, idvar='id', variable='cd4', time='enter')
     >>> plt.show() Generate a Receiver Operator Curve from true values and predicted probabilities. Youden's Index can also be
     calculated. Youden's index is calculated as
 
     .. math::
 
         P_{Yi} = max(Se_i + Sp_i - 1)
 
     Parameters
     ------------
     df : DataFrame
         Pandas dataframe containing variables of interest
     true : str
         True designation of the outcome (1, 0)
     threshold : str
         Predicted probabilities for the outcome
     youden_index : bool, optional
         Whether to calculate Youden's index. Youden's index maximizes both sensitivity and specificity. The formula
         finds the maximum of (sensitivity + specificity - 1)
 
     Returns
     -----------
     matplotlib axes
 
     Examples
     --------
     Creating a dataframe with true disease status (`'d'`) and predicted probability of the outcome (`'p'`)
 
     >>> import pandas as pd
     >>> import matplotlib.pyplot as plt
     >>> from zepid.graphics import roc
     >>> df = pd.DataFrame()
     >>> df['d'] = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
     >>> df['p'] = [0.1, 0.15, 0.1, 0.7, 0.5, 0.9, 0.95, 0.5, 0.4, 0.8, 0.99, 0.99, 0.89, 0.95]
 
     Creating ROC curve
 
     >>> roc(df, true='d', threshold='p', youden_index=False)
     >>> plt.show() # Getting all possible cutpoints # Going through all the cutpoints and calculating Sensitivity and 1-Specificity # If requested, calculate Youden's Index # Creating ROC plot Creates a plot of how the risk difference or risk ratio changes over time with survival data. See the
     references for an example of this plot. Input data should be pandas Series indexed by 'timeline' where 'timeline'
     is the time corresponding to the risk estimate
 
     Parameters
     --------------
     risk_exposed : Series
         Pandas Series with the probability of the outcome among the exposed group. Index by 'timeline' where 'timeline'
         is the time. If you directly output the ``1 - survival_function_`` from lifelines.KaplanMeierFitter(), this
         should create a valid input
     risk_unexposed : Series
         Pandas Series with the probability of the outcome among the exposed group. Index by 'timeline' where 'timeline'
         is the time
     measure : str, optional
         Whether to generate the risk difference (RD) or risk ratio (RR). Default is 'RD'
     loess : bool, optional
         Whether to generate LOESS curve fit to the calculated points. Default is True
     loess_value : float, optional
         Fraction of values to fit LOESS curve to. Default is 0.25
     point_color : str, optional
         Color of the points
     line_color : str, optional
         Color of the LOESS line generated and plotted
     scale : str, optional
         Change the y-axis scale. Options are 'linear' (default), 'log', 'log-transform'. 'log' and 'log-transform' is
         only a valid option for Risk Ratio plots
 
     Returns
     -----------
     matplotlib axes
 
     Examples
     --------
     See graphics documentation or causal documentation for a detailed example.
 
     >>> import matplotlib.pyplot as plt
     >>> from zepid.graphics import dynamic_risk_plot
 
     >>> dynamic_risk_plot(a, b, loess=True)
     >>> plt.show()
 
     References
     ----------
     <NAME>, et al. (2014). Estimation of the standardized risk difference and ratio in a competing risks framework:
     application to injection drug use and progression to AIDS after initiation of antiretroviral therapy.
     AJE, 181(4), 238-245. # This avoids a merge issue on floats # Generating the plot L'Abbe plots are useful for summarizing measure modification on the difference or ratio scale. Primarily
     invented for meta-analysis usage, these plots display risk differences (or ratios) by their individual risks
     by an exposure. I find them most useful for a visualization of why if there is an association and there is no
     modfication on one scale (additive or multiplicative), then there must be modification on the other scale.
 
     Parameters
     ----------
     r1 : float, list, optional
         Single probability or a list of probabilities when exposure is 1. Default is None, which does not display points
     r0 : float, list, optional
         Single probability or a list of probabilities when exposure is 0. Default is None, which does not display points
     scale : str, optional
         Which scale to plot. The default is 'both', which generates side-by-side plots of additive scale and
         multiplicative scale. Other options are; 'additive' to display the additive plot, and 'multiplicative' to
         display the multiplicative plot
     additive_tuner : int, optional
         Optional parameter to change the number of lines displayed in the additive L'Abbe plot. Higher integer
         produces more reference lines
     multiplicative_tuner : int, optional
         Optional parameter to change the number of lines displayed in the multiplicative L'Abbe plot. Higher integer
         produces more reference lines
     figsize : set, optional
         Optional parameter to change the L'Abbe plot size. Only changes the plot size when scale='both'
     **plot_kwargs : optional
         Optional keyword arguments for matplotlib. kwargs will pass matplotlib.pyploy.plot kwargs are accepted. See
         matplotlib 'plot()' function documentation for further details
 
     Returns
     -------
     matplotlib axes
 
     Examples
     --------
     Setting up environment
 
     >>> import matplotlib.pyplot as plt
     >>> from zepid.graphics import labbe_plot
 
     Creating a blank plot
 
     >>> labbe_plot()
     >>> plt.show()
 
     Adding customized points to the plot
 
     >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='additive', color='r', marker='D', markersize=10, linestyle='')
     >>> plt.show()
 
     Only returning the additive plot
 
     >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='additive', markersize=10)
     >>> plt.show()
 
     Only returning the multiplicative plot
 
     >>> labbe_plot(r1=[0.3, 0.5], r0=[0.2, 0.7], scale='multiplicative', markersize=10)
     >>> plt.show() # When markers are unspecified, point estimates aren't displayed. This avoids # When markers are unspecified, point estimates aren't displayed. This avoids # When markers are unspecified, point estimates aren't displayed. This avoids # When markers are unspecified, point estimates aren't displayed. This avoids