from collections import defaultdict
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from loguru import logger


class NaiveBayesScratch():
 """樸素貝葉斯算法Scratch實現(xiàn)"""
 def __init__(self):
  # 存儲先驗概率 P(Y=ck)
  self._prior_prob = defaultdict(float)
  # 存儲似然概率 P(X|Y=ck)
  self._likelihood = defaultdict(defaultdict)
  # 存儲每個類別的樣本在訓練集中出現(xiàn)次數
  self._ck_counter = defaultdict(float)
  # 存儲每一個特征可能取值的個數
  self._Sj = defaultdict(float)

 def fit(self, X, y):
  """
  模型訓練，參數估計使用貝葉斯估計
  X:
   訓練集，每一行表示一個樣本，每一列表示一個特征或屬性
  y:
   訓練集標簽
  """
  n_sample, n_feature = X.shape
  # 計算每個類別可能的取值以及每個類別樣本個數
  ck, num_ck = np.unique(y, return_counts=True)
  self._ck_counter = dict(zip(ck, num_ck))
  for label, num_label in self._ck_counter.items():
   # 計算先驗概率，做了拉普拉斯平滑處理，即計算P（y）
   self._prior_prob[label] = (num_label + 1) / (n_sample + ck.shape[0])

  # 記錄每個類別樣本對應的索引
  ck_idx = []
  for label in ck:
   label_idx = np.squeeze(np.argwhere(y == label))
   ck_idx.append(label_idx)

  # 遍歷每個類別
  for label, idx in zip(ck, ck_idx):
   xdata = X[idx]
   # 記錄該類別所有特征對應的概率
   label_likelihood = defaultdict(defaultdict)
   # 遍歷每個特征
   for i in range(n_feature):
    # 記錄該特征每個取值對應的概率
    feature_val_prob = defaultdict(float)
    # 獲取該列特征可能的取值和每個取值出現(xiàn)的次數
    feature_val, feature_cnt = np.unique(xdata[:, i], return_counts=True)
    self._Sj[i] = feature_val.shape[0]
    feature_counter = dict(zip(feature_val, feature_cnt))
    for fea_val, cnt in feature_counter.items():
     # 計算該列特征每個取值的概率，做了拉普拉斯平滑，即為了計算P（x|y）
     feature_val_prob[fea_val] = (cnt + 1) / (self._ck_counter[label] + self._Sj[i])
    label_likelihood[i] = feature_val_prob
   self._likelihood[label] = label_likelihood

 def predict(self, x):
  """
  輸入樣本，輸出其類別，本質上是計算后驗概率
  **注意計算后驗概率的時候對概率取對數**，概率連乘可能導致浮點數下溢，取對數將連乘轉化為求和
  """
  # 保存分類到每個類別的后驗概率，即計算P（y|x）
  post_prob = defaultdict(float)
  # 遍歷每個類別計算后驗概率
  for label, label_likelihood in self._likelihood.items():
   prob = np.log(self._prior_prob[label])
   # 遍歷樣本每一維特征
   for i, fea_val in enumerate(x):
    feature_val_prob = label_likelihood[i]
    # 如果該特征值出現(xiàn)在訓練集中則直接獲取概率
    if fea_val in feature_val_prob:
     prob += np.log(feature_val_prob[fea_val])
    else:
     # 如果該特征沒有出現(xiàn)在訓練集中則采用拉普拉斯平滑計算概率
     laplace_prob = 1 / (self._ck_counter[label] + self._Sj[i])
     prob += np.log(laplace_prob)
   post_prob[label] = prob
  prob_list = list(post_prob.items())
  prob_list.sort(key=lambda v: v[1], reverse=True)
  # 返回后驗概率最大的類別作為預測類別
  return prob_list[0][0]


def main():
 X, y = load_iris(return_X_y=True)
 xtrain, xtest, ytrain, ytest = train_test_split(X, y, train_size=0.8, shuffle=True)

 model = NaiveBayesScratch()
 model.fit(xtrain, ytrain)

 n_test = xtest.shape[0]
 n_right = 0
 for i in range(n_test):
  y_pred = model.predict(xtest[i])
  if y_pred == ytest[i]:
   n_right += 1
  else:
   logger.info("該樣本真實標簽為：{}，但是Scratch模型預測標簽為：{}".format(ytest[i], y_pred))
 logger.info("Scratch模型在測試集上的準確率為：{}%".format(n_right * 100 / n_test))

if __name__ == "__main__":
 main()