# -*- coding: utf-8 -*-
import numpy as np
import time

class NaiveBayes:
    """
    算法 : 朴素贝叶斯
    主要步骤为: 
    1. 根据训练集数据和标签, 确定维度、特征的数量
    2. 根据训练集计算出 先验概率 和 条件概率
    3. 对测试集进行测试, 取后验概率最大化的标签, 符合期望风险最小化
    Default: lambda=1 采用贝叶斯平滑处理, 防止某标签或数据从未出现而使概率为0
    """
    
    def __init__(self, trainData, trainLabel, _lambda=1):
        """
        param: trainData [numpy ndarray] 训练集数据
        param: trainLabel [numpy 1darray] 训练集标签
        param: _lambda [int] 平滑项 采用 Default=1 拉普拉斯平滑
        """
        t0 = time.time()
        self.trainData = trainData
        self.trainLabel = trainLabel
        self._lambda = _lambda
        self.labelNum = len(np.unique(trainLabel)) # 可取的标签总数量
        self.featureNum = len(trainData[0]) # 特征的维度
        self.featureCount = [len(np.unique([trainData[i][k] for i in 
                            range(len(trainData))]))for k in range(self.featureNum)]# 特征的每个维度可取值的数量
        t1 = time.time()
        print("完成初始化: 耗时 {:.2f}s".format(t1-t0))
    
    def get_all_prob(self):
        """
        根据提供已有的训练数据集和标签集,
        构造拟合出 先验概率分布 和 条件概率分布
        
        @ret: None
        """
        trainData, trainLabel, t0 = self.trainData, self.trainLabel, time.time()
        labelNum, featureNum, maxFeatureCount, _lambda = self.labelNum, self.featureNum, np.max(self.featureCount), self._lambda
        Py = np.zeros((labelNum, 1)) # 初始化先验概率分布向量, 全部为0
        for i in range(labelNum): # 采用平滑化, 每一类的分子计数加一个lambda平滑, 分母加lambda*总标签数 
            Py[i] = (np.sum(trainLabel == i) + _lambda)/(len(trainLabel) + labelNum * _lambda)
        Py = np.log(Py) # 取对数化, 防止太多的低小数相乘而导致浮点数下溢floating point underflow
        Px_y = np.zeros((labelNum, featureNum, maxFeatureCount)) 
        # 初始化条件概率分布: 第一维labelNum--y取哪个标签
        #                   第二维featureNum--第几维度的数据
        #                   第三维maxFeatureCount--当前维度所能取几个值
        #
        # 先验概率分布计算
        for i in range(len(trainLabel)):
            x = trainData[i]
            label = trainLabel[i]
            for j in range(featureNum):
                Px_y[label][j][x[j]] += 1
        # 条件概率分布计算       
        for label in range(labelNum):
            for j in range(featureNum):
                value_count = list()
                for k in range(maxFeatureCount): # 计数:在标签取label, 第j维度取值为k的数据有多少个
                    value_count.append(Px_y[label][j][k])
                for k in range(maxFeatureCount): # 除法:加入平滑化后, 拟合条件概率分布
                    Px_y[label][j][k] = np.log((Px_y[label][j][k] + _lambda) / (np.sum(value_count) + _lambda * maxFeatureCount))
        t1 = time.time()
        print("完成先验/条件概率分布计算: 耗时 {:.2f}s".format(t1-t0))
        self.Py, self.Px_y = Py, Px_y 
    
    def predict(self,new_x):
        """
        朴素贝叶斯: 
        根据已计算拟合过了的先验概率和条件概率分布
        预测出新数据点new_x的类别
        采取后验概率最大化的准则, 理论上符合期望风险最小化的目标
        
        param: [lst/np.1darray] new_x 新的待预测的数据实例点
        @ret:  [int] new_x 的预测类别
        """
        Py, Px_y, featureNum, labelNum = self.Py, self.Px_y, self.featureNum, self.labelNum
        prob = [0.0] * labelNum # 初始化每个类别的预测概率
        for label in range(labelNum):
            curr_prob = 0.0
            for j in range(featureNum): # 因为概率分布已为Log对数化
                curr_prob += Px_y[label][j][new_x[j]] # 所以连乘变成了连加
            prob[label] = curr_prob + Py[label]
        return prob.index(max(prob)) # 返回后验概率最大化的类别为预测值
    
    def test(self, testData, testLabel):
        """
        对外界输入的测试数据集进行预测,
        并对比测试标签集, 得出测试准确率以及平均预测耗时
        
        param: [np.ndarry/lst[lst]] testData 测试数据集
        param: [np.1darray/lst] testLabel 测试标签集
        @ret: [float] accu 测试准确率
        """
        t0 = time.time()
        correct = 0
        for i in range(len(testLabel)):
            if self.predict(testData[i]) == testLabel[i]:
                correct += 1
        t1 = time.time()
        print("平均每次预测计算: 耗时{:.5f}s".format((t1-t0)/len(testData)))
        accu = correct / len(testData)
        print("测试成绩: {}/{} 准确率{:.2f}%".format(correct, len(testData), 100*accu))
        return accu