#!/usr/bin/env python
# coding: utf-8

# In[8]:


import pandas as pd
from yellowbrick.classifier.confusion_matrix import *
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import LabelEncoder


# In[9]:


df = pd.read_csv('examples/data/occupancy/occupancy.csv')

features = ["temperature", "relative humidity", "light", "C02", "humidity"]
target = "occupancy"

X = df[features]
y = df[target]
classes = ["unoccupied", "occupied"]
le = LabelEncoder()
le.fit(classes)
X_train, X_test, y_train, y_test = tts(X, y, test_size=0.2)

viz = ConfusionMatrix(LogisticRegression(), classes=classes, label_encoder =le)
viz.fit(X_train, y_train)
viz.score(X_test, y_test)
viz.poof()


# In[11]:


digits = load_digits()
X = digits.data
y = digits.target

X_train, X_test, y_train, y_test = train_test_split(X,y, test_size =0.2, random_state=11)
X_train = X_train
X_test = X_test
y_train = y_train
y_test = y_test

model = LogisticRegression()
classes = ['one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine']
mapping = {'zero': '0', 'one': '1', 'two': '2', 'three': '3', 'four': '4', 'five': '5',
           'six': '6', 'seven': '7', 'eight': '8', 'nine': '9'}
cm = ConfusionMatrix(model, classes=classes, label_encoder = mapping)
cm.fit(X_train, y_train)
cm.score(X_test, y_test)
cm.poof()