# encoding=utf-8

import pandas as pd

import numpy as np

from sklearn.model_selection import train_test_split

import time

def binaryzation(array):

    #二值化

    # 大于50的为1，小于50的为0，二值化后，参数总量从255降到了2

    array[array[:, :] <= 50] = 0

    array[array[:, :] > 50] = 1

    return array

def read_data(path):

    raw_data = pd.read_csv(path, header=0)

    data = raw_data.values

    imgs = data[:, 1:]

    labels = data[:, 0]

    imgs = binaryzation(imgs)

    # 选2/3的数据作为数据集，1/3的数据作为测试集

    tra_x, test_x, tra_y, test_y = train_test_split(imgs, labels, test_size=0.33, random_state=2019526)

    return tra_x, test_x, tra_y, test_y

def naive_Bayes(x, y, class_num=10, feature_len=784, point_len=2, lamba=1):

    """

    :param x: 学习的数据集

    :param y: 学习的标签

    :param class_num: y的类别的数目

    :param feature_len: 每个训练数据的特征向量有多少维

    :param point_len: 每个训练集数据的特征向量有多少种取值情况

    :return: 先验概率和条件概率

    """

    # 这里我们使用贝叶斯估计法来做

    prior_probability = np.zeros(class_num)

    conditional_probability = np.zeros((class_num, feature_len, point_len))

    # 先得到每个对应的数目,计算先验概率

    for i in range(10):

        prior_probability[i] += len(y[y[:] == i])

    # 计算条件概率

    for i in range(class_num):

        # 取出x中label对应数据为i的点

        data_ith = x[y == i]

        for j in range(feature_len):

            for k in range(point_len):

                # 在之前的数据中取出第j列为1，或者为0的点

                conditional_probability[i, j, k] += len(data_ith[data_ith[:, j] == k]) + lamba

        conditional_probability[i, :, :] /= (prior_probability[i] + point_len * lamba)

    prior_probability += lamba

    prior_probability /= (len(y) + class_num * lamba)

    return prior_probability, conditional_probability

def predict(pp, cp, data, class_num=10,feature_len=784):

    """

    预测

    :param pp: prior_probability

    :param cp: conditional_probability

    :param data: 输入数据

    :return: 预测结果

    """

    pre_label = np.zeros(len(data)+1)

    for i in range(len(data)):

        max_possibility = 0

        max_label = 0

        for j in range(class_num):

            tmp=pp[j]

            for k in range(feature_len):

                tmp*=cp[j,k,data[i,k]]

            if tmp>max_possibility:

                max_possibility=tmp

                max_label=j

        pre_label[i]=max_label

    return pre_label

if __name__ == '__main__':

    time_1=time.time()

    tra_x, test_x, tra_y, test_y = read_data('data/Mnist/mnist_train.csv')

    time_2=time.time()

    prior_probability, conditional_probability = naive_Bayes(tra_x, tra_y, 10, 784, 2, 1)

    time_3=time.time()

    pre_label = predict(prior_probability, conditional_probability, test_x, 10,784)

    cous=0

    for i in range(len(test_y)):

        if pre_label[i]==test_y[i]:

            cous+=1

    print(cous/len(pre_label))

    time_4=time.time()

    print("time of reading data: ",int(time_2-time_1))

    print("time of getting model: ",int(time_3-time_2))

    print("time of predictting: ",int(time_4-time_3))