# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/charts.ipynb.

# %% auto 0
__all__ = ['alt', 'Chart', 'CM', 'confusion_matrix']

# %% ../nbs/charts.ipynb 2
import altair as alt
from altair import Chart

alt = alt
Chart = Chart
from .loader import *

# %% ../nbs/charts.ipynb 5
def confusion_matrix(df=None, truth=None, pred=None, mapping=None, save_to=None):
    if df is None:
        df = pd.DataFrame({"truth": truth, "pred": pred})
        truth = "truth"
        pred = "pred"
    DF = df[[truth, pred]]

    try:
        from sklearn.metrics import classification_report

        print(classification_report(DF[truth], DF[pred]))
    except:
        logger.info("Skipping Report")
    df = DF.pivot_table(index=truth, columns=pred, aggfunc=len).reset_index()
    df = df.melt(id_vars=truth, var_name=pred, value_name="value")
    sz = 450 if len(DF[truth].unique()) > 4 else 250

    base = alt.Chart(df, height=sz, width=sz)
    hm = base.mark_rect().encode(
        x=f"{pred}:O",
        y=f"{truth}:O",
        color=alt.Color("value:Q", scale=alt.Scale(scheme="lightorange"), legend=None),
    )

    tx = base.mark_text(baseline="middle").encode(
        x=f"{pred}:O",
        y=f"{truth}:O",
        text="value:Q",
        color=alt.condition(
            alt.datum.value > 0,
            alt.value("black"),
            alt.value("rgba(0, 0, 0, 0)"),  # Transparent color for value 0
        ),
    )

    if save_to is not None:
        assert str(save_to).endswith("html"), "Can only save in html format"
        (hm + tx).save(save_to, format="html")
        Info(f"Saved chart at {save_to}")
    else:
        return hm + tx


CM = confusion_matrix
