import matplotlib.pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix


def save_confusion_matrix(truth, pred, out_file, norm=None):
    label_names = list(set(truth) | set(pred))
    label_names.sort()
    truth_compact = [label_names.index(x) for x in truth]
    pred_compact = [label_names.index(x) for x in pred]
    cm = confusion_matrix(
        truth_compact, pred_compact, labels=list(range(len(label_names))),
        normalize=norm)
    if norm is not None:
        cm *= 100
    fig = plt.figure(figsize=(20, 20))
    ax = fig.add_subplot(111)
    disp = ConfusionMatrixDisplay(
        confusion_matrix=cm, display_labels=label_names)
    disp.plot(ax=ax, xticks_rotation='vertical',
              values_format='.1f' if norm is not None else 'd')
    plt.tight_layout()
    plt.savefig(out_file)
    plt.close(fig)

[((15, 10, 15, 38), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((17, 11, 18, 56), 'sklearn.metrics.ConfusionMatrixDisplay', 'ConfusionMatrixDisplay', (), '', False, 'from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix\n'), ((21, 4, 21, 22), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((22, 4, 22, 25), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(22, 16, 22, 24): 'out_file'}, {}), '(out_file)', True, 'import matplotlib.pyplot as plt\n'), ((23, 4, 23, 18), 'matplotlib.pyplot.close', 'plt.close', ({(23, 14, 23, 17): 'fig'}, {}), '(fig)', True, 'import matplotlib.pyplot as plt\n')]