# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/20_plotting.ipynb (unless otherwise specified).

__all__ = ['custom_params', 'icc_comparison_plot']

# Cell
import seaborn as sns
import pandas as pd
from matplotlib import pyplot as plt
custom_params = {"axes.spines.right": False, "axes.spines.top": False}
sns.set_theme(style="white", font_scale = 1.3, rc=custom_params)

# Cell
def icc_comparison_plot(iccs,
                        icc_type = "ICC1",
                       method_labels = {'mean':'Based on mean method',
                      'trial_sep':'Based on predicted values (separate models)',
                      'trial_joint':'Based on predicted values (joint model)',
                      'model_calc':'Based on model-calculated variances'
                      },
                        qualitative_labels = False,
                        x_offset = 3.5,
                        font_scale = 1.3,
                        **catplot_kwgs):
    custom_params = {"axes.spines.right": False, "axes.spines.top": False}
    sns.set_theme(style="white", font_scale = font_scale, rc=custom_params)
    icc_type = "ICC1"
    plot_df = pd.DataFrame({})
    for k, v in iccs.items():
        if 'iccs' in v.keys():
            icc = v['iccs'][0].query('type=="%s"'%icc_type)['ICC'].iloc[0]
        else:
            icc = v['icc']
        plot_df.loc[k,'icc'] = icc

    if 'aspect' not in catplot_kwgs.keys():
        catplot_kwgs['aspect'] = 1.3
    if 'height' not in catplot_kwgs.keys():
        catplot_kwgs['height'] = 6

    plot_df = plot_df.reset_index()
    plot_df.rename(columns = {'index':'category','icc':"Test-retest reliability (%s)"%icc_type}, inplace = True)
    plot_df['trial_type'] = plot_df.category.apply(lambda x: x.split('_')[-1])
    plot_df['Method'] = plot_df.category.apply(lambda x: '_'.join(x.split('_')[:-1]))
    plot_df.Method.replace(method_labels, inplace = True)
    g = sns.catplot(x = "trial_type",       # x variable name
                y = "Test-retest reliability (%s)"%icc_type,       # y variable name
                hue = "Method",  # group variable name
                data = plot_df,     # dataframe to plot
                kind = "bar",
                    legend = False,
               ci = None,
                **catplot_kwgs)
    g.set(ylim=(0, 1))
    ax = g.axes[0][0]

    if qualitative_labels:
        y_offset = -.07
        x_offset = x_offset
        alignment = 'left'
        plt.axhline(.5, color = 'grey', linestyle = '--',xmax = 1)
        plt.axhline(.75, color = 'grey', linestyle = '--',xmax = 1)
        plt.annotate("poor reliability", (0+x_offset,.5+y_offset),horizontalalignment = alignment)
        plt.annotate("moderate reliability", (0+ x_offset,.67+y_offset),horizontalalignment = alignment)
        plt.annotate("good reliability", (0+x_offset,.86+y_offset),horizontalalignment = alignment)
    plt.legend(loc='upper left', bbox_to_anchor=(1, 1.1))

    return g