import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.linear_model import LogisticRegression


def dataShow(data, cls, categoryLabel):
    # data为

    # 生成网格采样点
    x1_min, x1_max = data[:, 0].min(), data[:, 0].max()
    x2_min, x2_max = data[:, 1].min(), data[:, 1].max()
    x1, x2 = np.mgrid[x1_min:x1_max:200j, x2_min:x2_max:200j]
    # 测试点
    grid_test = np.stack((x1.flat, x2.flat), axis=1)

    # 预测分类值
    grid_hat = cls.predict(grid_test)
    grid_hat = grid_hat.reshape(x1.shape)
    cm_light = mpl.colors.ListedColormap(['#FFFFFF', '#D3D3D3', '#708090'])
    cm_dark = mpl.colors.ListedColormap(['g', 'b', 'r'])
    marker = ['o', '*', '+']
    # 根据预测分类值，绘制表现出分类边界
    plt.pcolormesh(x1, x2, grid_hat, cmap=cm_light)
    # 绘制样本散点图
    for i in range(data.shape[0]):
        plt.scatter(data[i, 0], data[i, 1], c='black', s=80,
                    marker=np.array(marker)[categoryLabel[i]], cmap=cm_dark)


    # 坐标轴显示设置
    plt.xlim(x1_min, x1_max)
    plt.ylim(x2_min, x2_max)
    plt.grid()
    plt.show()



# Press the green button in the gutter to run the script.
if __name__ == '__main__':

    iris = datasets.load_iris()
    y = iris.target
    x = iris.data[:, [0,2]]
    cls = LogisticRegression()
    cls.fit(x, y)
    dataShow(x, cls, y)