#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Author : 陈坤泽
# @Email  : 877362867@qq.com
# @Date   : 2023/11/05

""" 基于jieba库的一些文本处理功能 """

from collections import Counter
import re

from tqdm import tqdm
import pandas as pd

import jieba
import jieba.posseg as pseg

from pyxllib.prog.pupil import DictTool, run_once
from pyxllib.file.specialist import XlPath
from pyxllib.algo.stat import update_dataframes_to_excel


def jieba_add_words(words):
    for w in words:
        jieba.add_word(w)


def jieba_del_words(words):
    for w in words:
        jieba.del_word(w)


@run_once('str')
def jieba_cut(text):
    return tuple(jieba.cut(text))


@run_once('str')
def pseg_cut(text):
    return tuple(pseg.cut(text))


def _count_word_frequency(texts, function_word=True):
    """ 统计关键词出现频数 (主要是协助计算tf-idf)

    :param texts: 输入字符串列表
    :param function_word: 是否要统计虚词
    :return: 一个dict
        key: 分词名称
        values: [x, y]，x是出现总频数，y是这个词在多少篇文章中出现过

    >>> _count_word_frequency(['正正正正', '正反正', '反反反反'])
    {'正正': [1, 1], '反反': [2, 1]}

    原没有过滤词性的结果：{'正正': [2, 1], '正': [1, 1], '反正': [1, 1], '反反': [2, 1]}
    """

    d = dict()
    for text in tqdm(texts, '词频统计'):
        wordflags = list(pseg.cut(text))
        words = set()
        for word, flag in wordflags:
            # 虚词不做记录
            if (not function_word) and flag in ('uj', 'd', 'p', 'c', 'u', 'xc'):
                continue
            words.add(word)
            if word not in d:
                d[word] = [0, 0]
            d[word][0] += 1
        for word in words:
            d[word][1] += 1
    return d


def analyse_tf_idf(texts, outfile=None, sheet_name='tf-idf', *, function_word=True):
    """ 分析tf-idf值

    :param list[str] texts: 多份文件的文本内容
    :return: 一个DataFrame数据

    这个算法jieba可能有些自带库可以搞，但是自己写一下也不难啦
    注意我这里返回的tf-idf中，是放大了总频数倍的，这样显示的数值大一点，看起来舒服~
    """
    from math import log10

    frequency = _count_word_frequency(texts, function_word)
    DictTool.isub(frequency, [' ', '\t', '\n'])

    n = len(texts)
    sum_frequency = sum([v[0] for v in frequency.values()])

    li = []
    for k, v in frequency.items():
        idf = log10(n / v[1])
        # idf = 1
        li.append([k, v[0], v[0] / sum_frequency, v[1], idf, v[0] * idf])
    df = pd.DataFrame.from_records(li, columns=('词汇', '频数', '频率', '出现该词文章数', 'idf', 'tf-idf'))
    df.sort_values(by='tf-idf', ascending=False, inplace=True)

    if outfile:
        update_dataframes_to_excel(outfile, {sheet_name: df})

    return df


class TextClassifier:
    def __init__(self, texts=None):
        """ 文本分类器

        :param list[str] texts: 文本内容
        """

        self.texts = []
        self.tfidf = {}
        self.vecs = []  # 每份文本对应的向量化表达
        self.default_tfidf = 1  # 如果没有计算tf-idf，可以全部默认用权重1

        if texts:
            for text in texts:
                self.texts.append(text)

    def get_text_tf(self, text, *,
                    function_word_weight=0.2,
                    normalize=True,
                    ingore_words=(' ', '\t', '\n'),
                    add_flag=False):
        """ 这里可以定制提取text关键词的算法

        :param function_word_weight: 这里可以自定义功能性词汇权重，一般是设一个小数降低权重

        一般是定制一些过滤规则，比如过滤掉一些词性，或者过滤掉一些词
        """
        ct = Counter()

        # 1 初步的分词，以及是否要过滤虚词
        wordflags = list(pseg_cut(text))
        for word, flag in wordflags:
            if flag in ('uj', 'd', 'p', 'c', 'u', 'xc', 'x'):
                if add_flag:
                    ct[word + ',' + flag] += function_word_weight
                else:
                    ct[word] += function_word_weight
            else:
                if add_flag:
                    ct[word + ',' + flag] += 1
                else:
                    ct[word] += 1

        # 2 归一化一些词
        if normalize:
            ct2 = Counter()
            for k, v in ct.items():
                # 如果需要对一些词汇做归一化，也可以这里设置
                k = re.sub(r'\d', '0', k)  # 把数字都换成0
                ct2[k] += v
            ct = ct2

        # 3 过滤掉一些词
        if ingore_words:
            for k in ingore_words:
                if k in ct:
                    del ct[k]

        return ct

    def compute_tfidf(self, outfile=None, sheet_name='tf-idf', normalize=False, function_word_weight=0.2, add_flag=False):
        """ 重算tfidf表 """
        from math import log10

        # 1 统计频数和出现该词的文章数
        d = dict()
        for text in tqdm(self.texts, '词频统计'):
            ct = self.get_text_tf(text, normalize=normalize, function_word_weight=function_word_weight, add_flag=add_flag)
            for k, v in ct.items():
                if k not in d:
                    d[k] = [0, 0]
                d[k] = [d[k][0] + v, d[k][1] + 1]

        # 2 计算tfidf
        n = len(self.texts)
        sum_tf = sum([v[0] for v in d.values()])
        ls = []
        for k, v in d.items():
            idf = log10(n / v[1])
            # idf = 1
            ls.append([k, v[0], v[0] / sum_tf, v[1], idf, v[0] * idf])

        df = pd.DataFrame.from_records(ls, columns=('词汇', '频数', '频率', '出现该词文章数', 'idf', 'tf-idf'))
        df.sort_values(by='tf-idf', ascending=False, inplace=True)

        # 3 保存到文件
        if outfile:
            update_dataframes_to_excel(outfile, {sheet_name: df})

        self.tfidf = {row['词汇']: row['tf-idf'] for idx, row in df.iterrows()}
        self.default_tfidf = df.loc[len(df) - 1]['tf-idf']  # 最后条的权重作为其他未见词的默认权重

        return df

    def normalization(self, d):
        """ 向量归一化

        输入一个类字典结构表示的向量，对向量做归一化处理
        """
        length = sum([v * v for v in d.values()]) ** 0.5  # 向量长度
        return {k: v / length for k, v in d.items()}

    def get_text_vec(self, text):
        """ 获取文本的向量化表达

        :param str text: 文本内容
        """
        ct = self.get_text_tf(text)
        vec = {k: v * self.tfidf.get(k, self.default_tfidf) for k, v in ct.items()}
        vec = self.normalization(vec)
        return vec

    def compute_vecs(self):
        """ 重置向量化表达 """
        vecs = []
        for text in tqdm(self.texts, desc='query向量化'):
            vecs.append(self.get_text_vec(text))
        self.vecs = vecs
        return vecs

    def cosine_similar(self, x, y):
        """ 两个向量的余弦相似度，值越大越相似

        这里是简化的，只算两个向量的点积，请确保输入的都是单位长度的向量
        注意这里x和y都是稀疏矩阵的存储形式，传入的是dict结构
        """
        keys = x.keys() & y.keys()  # 求出x和y共有的键值
        return sum([x[k] * y[k] for k in keys])

    def find_similar_vec(self, x, maxn=10):
        """ 找与x最相近的向量，返回下标和相似度

        :pamra x: 待查找的对象
        :param maxn: 返回最相近的前maxn个对象
        """
        if isinstance(x, str):
            x = self.get_text_vec(x)

        # todo 使用并行计算？或者其实也可以向量化，但向量化是稀疏矩阵，挺占空间的
        sims = [(i, self.cosine_similar(x, v)) for i, v in enumerate(self.vecs)]
        sims.sort(key=lambda x: x[1], reverse=True)
        return sims[:maxn]

    def refine_vecs(self):
        """ 优化向量数据 """
        # 1 计算每个向量的长度
        vecs = []
        for vec in tqdm(self.vecs, '优化向量'):
            vec = [(k, v) for k, v in vec.items()]
            vec.sort(key=lambda x: x[1], reverse=True)
            vec2 = {}
            for k, v in vec:
                if v < 0.0001:
                    break
                vec2[k] = round(v, 4)
            vecs.append(vec2)

        self.vecs = vecs
        return self.vecs
