# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/charts.ipynb (unless otherwise specified).

__all__ = ['alt', 'Chart', 'confusion_matrix', 'CM']

# Cell
# export
import altair as alt
from altair import Chart
alt = alt
Chart = Chart
from .loader import *

# Cell
def confusion_matrix(df=None, truth=None, pred=None, mapping=None):
    if df is None:
        df = pd.DataFrame({'truth': truth, 'pred': pred})
        truth = 'truth'
        pred = 'pred'
    threshold = len(df)
    if mapping:
        assert isinstance(mapping, dict), 'mapping should be a dictionary'
        df[truth] = df[truth].map(lambda x: mapping[x])
        df[pred] = df[pred].map(lambda x: mapping[x])

    sz = 450 if len(df[truth].unique()) > 4 else 250
    base = Chart(df, height=sz, width=sz).transform_aggregate(
        num_vals='count()',
        groupby=[truth, pred]
    ).transform_calculate(
        rev_num_vals='-(datum.num_vals) + max(datum.num_vals)',
    ).encode(
        alt.Y(f'{truth}:O', scale=alt.Scale(paddingInner=0)),
        alt.X(f'{pred}:O', scale=alt.Scale(paddingInner=0)),
    )

    hm = base.mark_rect().encode(
        color=alt.Color(
            'num_vals:Q', scale=alt.Scale(scheme="lightorange"),
            legend=None
        )
    )

    tx = base.mark_text(baseline='middle').encode(
        text='num_vals:Q',
#         color=alt.Color(alt.value('gray'))
#         color='rev_num_vals:Q'
#         color=alt.Color(
#             'num_vals:Q', scale=alt.Scale(scheme="redyellowgreen"),
#         )
        color=alt.condition(
            alt.datum.num_vals > threshold,
            alt.value('black'),
            alt.value('black'))
    )

    try:
        from sklearn.metrics import classification_report
        print(classification_report(df[truth], df[pred]))
    except:
        logger.info('Skipping Report')
    return hm + tx

CM = confusion_matrix