# 导入所需的库
import random
import jieba
import collections

# 定义n-gram模型的类
class NGramModel:

  # 初始化模型，设定n的值和平滑参数
  def __init__(self, n, alpha):
    self.n = n
    self.alpha = alpha
    self.ngrams = collections.defaultdict(int) # 存储n-gram的计数
    self.contexts = collections.defaultdict(int) # 存储n-1-gram的计数
    self.vocabulary = set() # 存储词汇表

  # 根据语料库训练模型，更新n-gram和n-1-gram的计数和词汇表
  def train(self, corpus):
    for sentence in corpus:
      # 在句首和句尾添加开始和结束标记
      sentence = ["<s>"] * (self.n - 1) + sentence + ["</s>"]
      # 对每个n-gram进行计数
      for i in range(len(sentence) - self.n + 1):
        ngram = tuple(sentence[i:i+self.n])
        self.ngrams[ngram] += 1
        # 对每个n-1-gram进行计数
        context = tuple(sentence[i:i+self.n-1])
        self.contexts[context] += 1
        # 更新词汇表
        self.vocabulary.update(ngram)

  # 根据n-1-gram的上下文预测下一个词的概率分布，使用加法平滑
  def predict(self, context):
    # 初始化概率分布字典
    probabilities = {}
    # 遍历词汇表中的每个词
    for word in self.vocabulary:
      # 构造n-gram
      ngram = tuple(context) + (word,)
      # 计算n-gram的概率，使用加法平滑
      probability = (self.ngrams[ngram] + self.alpha) / (self.contexts[tuple(context)] + self.alpha * len(self.vocabulary))
      # 将概率存入字典
      probabilities[word] = probability
    # 返回概率分布字典
    return probabilities

  # 根据概率分布字典随机选择一个词，使用轮盘赌算法
  def sample(self, probabilities):
    # 计算概率分布的总和
    total = sum(probabilities.values())
    # 生成一个0到总和之间的随机数
    random_number = random.uniform(0, total)
    # 初始化累积概率
    cumulative_probability = 0.0
    # 遍历概率分布字典中的每个词和概率
    for word, probability in probabilities.items():
      # 累加概率
      cumulative_probability += probability
      # 如果累积概率大于等于随机数，返回该词
      if cumulative_probability >= random_number:
        return word

  # 根据模型生成一句话，给定一个初始的上下文
  def generate(self, context):
    # 初始化生成的句子列表，包含初始的上下文
    sentence = list(context)
    # 循环生成下一个词，直到遇到结束标记或达到最大长度
    while True:
     # 预测下一个词的概率分布
      probabilities = self.predict(context)
      # 随机选择一个词
      word = self.sample(probabilities)
      # 将词添加到句子列表
      sentence.append(word)
      # 如果词是结束标记，跳出循环
      if word == "</s>":
        break
      # 更新上下文，去掉第一个词，加上最后一个词
      context = context[1:] + (word,)
    # 返回生成的句子列表，去掉开始和结束标记
    return sentence[self.n-1:-1]

# 读取语料库文件，分词并存入列表
corpus = []
with open("corpus.txt", encoding="utf-8") as f:
  for line in f:
    line = line.strip()
    if line:
      words = list(jieba.cut(line))
      corpus.append(words)
print("语料库中的句子数：", len(corpus))
print(corpus)
# 创建一个3-gram模型，平滑参数设为0.01
model = NGramModel(3, 0.01)
# 根据语料库训练模型
model.train(corpus)
print("词汇表中的词数：", len(model.vocabulary))
print("n-1-gram的计数：", model.contexts.items())
      
# 生成一句话，初始上下文设为("我", "爱")
sentence = model.generate(("我", "爱"))
# 将生成的句子列表拼接成字符串并打印
print("".join(sentence))
sentence = model.generate(("我",))
print("".join(sentence))