import numpy as np
import matplotlib.pyplot as plt

def plot_classification_report(cr, title='Classification report ', with_avg_total=False, cmap=plt.cm.Blues):

    lines = cr.split('\n')

    classes = []
    plotMat = []
    for line in lines[2 : (len(lines) - 3)]:
        #print(line)
        t = line.split()
        # print(t)
        classes.append(t[0])
        v = [float(x) for x in t[1: len(t) - 1]]
        print(v)
        plotMat.append(v)

    if with_avg_total:
        aveTotal = lines[len(lines) - 1].split()
        classes.append('avg/total')
        vAveTotal = [float(x) for x in t[1:len(aveTotal) - 1]]
        plotMat.append(vAveTotal)


    plt.imshow(plotMat, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    x_tick_marks = np.arange(3)
    y_tick_marks = np.arange(len(classes))
    plt.xticks(x_tick_marks, ['precision', 'recall', 'f1-score'], rotation=45)
    plt.yticks(y_tick_marks, classes)
    plt.tight_layout()
    plt.ylabel('Classes')
    plt.xlabel('Measures')
    plt.show()


# Test example of classification report
sampleClassificationReport = """             precision    recall  f1-score   support

          1       0.62      1.00      0.76        66
          2       0.93      0.93      0.93        40
          3       0.59      0.97      0.73        67
          4       0.47      0.92      0.62       272
          5       1.00      0.16      0.28       413

avg / total       0.77      0.57      0.49       858"""

plot_classification_report(sampleClassificationReport)

import pandas as pd

from sklearn.datasets import load_wine
from sklearn.metrics import classification_report
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

data = load_wine()
X_train, X_test, y_train, y_test = train_test_split(data.data, data.target, shuffle=True)

model = MultinomialNB().fit(X_train, y_train)
y_pred = model.predict(X_test)
classificationReport = classification_report(y_test, y_pred, target_names=data.target_names)
plot_classification_report(classificationReport)

model = LogisticRegression().fit(X, y)
y_pred = model.predict(X)
classificationReport = classification_report(y, y_pred, target_names=data.target_names)
plot_classification_report(classificationReport)