"""
Coding: UTF-8
Author: Randal
Time: 2021/2/20
E-mail: RandalJin@163.com

Description: This is a simple toolkit for data extraction of text.
The most important function in the script is about word frequency statistics.
Using re, I generalized the process in words counting, regardless of  any preset
word segmentation. Besides, many interesting functions, like getting top sentences are built here.

All rights reserved.

"""

import xlwings as xw
import pandas as pd
import numpy as np
import os
import re
from alive_progress import alive_bar
from alive_progress import show_bars, show_spinners
import jieba
import datetime
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
import math

class jieba_vectorizer(CountVectorizer):

    def __init__(self, tf, userdict, stopwords, orient=False):
        """
        :param tf: 输入的样本框，{axis: 1, 0: id, 1: 标题, 2: 正文, 3: 来源, 4: freq}
        :param stopwords: 停用词表的路径
        :param user_dict_link: 关键词清单的路径
        :param orient: {True: 返回的 DTM 只包括关键词清单中的词，False: 返回 DTM 中包含全部词语}
        :return: 可以直接使用的词向量样本
        """
        self.userdict = userdict
        self.orient = orient
        self.stopwords = stopwords
        jieba.load_userdict(self.userdict)  # 载入关键词词典

        tf = tf.copy()  # 防止对函数之外的原样本框造成改动
        print('切词中，请稍候……')
        rule = re.compile(u'[^\u4e00-\u9fa5]')  # 清洗所有样本，只保留汉字
        for i in range(0, tf.shape[0]):
            try:
                tf.iloc[i, 2] = rule.sub('', tf.iloc[i, 2])
            except TypeError:
                print('样本清洗Error: doc_id = ' + str(i))
            continue

        if self.stopwords is not None:
            stopwords = txt_to_list(self.stopwords)  # 载入停用词表
        else:
            stopwords = []

        # 开始切词

        words = []
        items = range(0, len(tf))
        with alive_bar(len(items), force_tty=True, bar='circles') as bar:
            for i, row in tf.iterrows():
                item = row['正文']
                result = jieba.cut(item)
                # 同时过滤停用词
                word = ''
                for element in result:
                    if element not in stopwords:
                        if element != '\t':
                            word += element
                            word += " "
                words.append(word)
                bar()

        # CountVectorizer() 可以自动完成词频统计，通过fit_transform生成文本向量和词袋库
        # 如果需要换成 tfidfVectorizer, 把下面三行修改一下就可以了
        vect = CountVectorizer()
        X = vect.fit_transform(words)
        self.vectorizer = vect

        matrix = X
        X = X.toarray()
        # 二维ndarray可以展示在pycharm里，但是和DataFrame性质完全不同
        # ndarray 没有 index 和 column
        features = vect.get_feature_names()
        XX = pd.DataFrame(X, index=tf['id'], columns=features)

        self.DTM0 = matrix
        self.DTM = XX
        self.features = features

        # # 下面是之前走的弯路，不足一哂
        # words_bag = vect.vocabulary_
        # # 字典的转置（注意只适用于vk一一对应的情况，1v多k请参考setdefault)
        # bag_words = dict((v, k) for k, v in words_bag.items())
        #
        # # 字典元素的排列顺序不等于字典元素值的排列顺序
        # lst = []
        # for i in range(0, len(XX.columns)):
        #     lst.append(bag_words[i])
        # XX.columns = lst

        if orient:
            dict_filter = txt_to_list(self.userdict)
            for word in features:
                if word not in dict_filter:
                    XX.drop([word], axis=1, inplace=True)
            self.DTM_key = XX

    def get_feature_names(self):
        return self.features

    def strip_non_keywords(self, df):
        ff = df.copy()
        dict_filter = txt_to_list(self.userdict)
        for word in self.features:
            if word not in dict_filter:
                ff.drop([word], axis=1, inplace=True)
        return ff

def make_doc_freq(word, doc):
    """
    :param word: 指的是要对其进行词频统计的关键词
    :param doc: 指的是要遍历的文本
    :return: lst: 返回字典，记录关键词在文本当中出现的频次以及上下文
    """
    # 使用正则表达式进行匹配, 拼接成pattern

    # re.S表示会自动换行
    # finditer是findall的迭代器版本，通过遍历可以依次打印出子串所在的位置
    it = re.finditer(word, doc, re.S)
    # match.group()可以返回子串，match.span()可以返回索引
    lst = []
    for match in it:
        lst.append(match.span())
    freq = dict()
    freq['Frequency'] = len(lst)
    # 将上下文结果也整理为一个字典
    context = dict()
    for i in range(0, len(lst)):
        # 将span的范围前后各扩展不多于10个字符，得到上下文
        try:
            # 为了划出适宜的前后文范围，需要设定索引的最大值和最小值
            # 因此要比较span+10和doc极大值，span-10和doc极小值
            # 最大值在两者间取小，最小值在两者间取大
            MAX = min(lst[i][1] + 10, len(doc))
            MIN = max(0, lst[i][0] - 10)
            # 取得上下文
            context[str(i)] = doc[MIN: MAX]
        except IndexError:
            print('IndexError: ' + word)
    freq['Context'] = context
    return freq

def make_info_freq(name, pattern, doc):
    """
    :param name: 指的是对其进行词频统计的形式
    :param pattern: 指的是对其进行词频统计的正则表达式
    :param doc: 指的是要遍历的文本
    :return: lst: 返回字典，记录关键词在文本当中出现的频次以及上下文
    注：该函数返回字典中的context元素为元组：（关键词，上下文）
    """
    # 使用正则表达式进行匹配, 拼接成pattern
    # re.S表示会自动换行
    # finditer是findall的迭代器版本，通过遍历可以依次打印出子串所在的位置
    it = re.finditer(pattern[0], doc, re.S)
    # match.group()可以返回子串，match.span()可以返回索引
    cls = pattern[1]
    lst = []
    for match in it:
        lst.append(match.span())
    freq = dict()
    freq['Frequency'] = len(lst)
    freq['Name'] = name
    # 将上下文结果也整理为一个字典
    context = dict()
    for i in range(0, len(lst)):
        # 将span的范围前后各扩展不多于10个字符，得到上下文
        try:
            # 为了划出适宜的前后文范围，需要设定索引的最大值和最小值
            # 因此要比较span+10和doc极大值，span-10和doc极小值
            # 最大值在两者间取小，最小值在两者间取大
            MAX = min(lst[i][1] + 10, len(doc))
            MIN = max(0, lst[i][0] - 10)
            # 取得匹配到的关键词，并做掐头去尾处理
            word = match_cut(doc[lst[i][0]: lst[i][1]], cls)
            # 将关键词和上下文打包，存储到 context 条目中
            context[str(i)] = (word, doc[MIN: MAX])
        except IndexError:
            print('IndexError: ' + name)
    freq['Context'] = context
    return freq

def make_docs_freq(word, docs):
    """
    :param word: 指的是要对其进行词频统计的关键词
    :param docs: 是要遍历的文本的集合，必须是pandas DataFrame的形式，至少包含id列 (iloc: 0)，正文列 (iloc: 2) 和预留出的频次列 (iloc: 4)
    :return: 返回字典，其中包括“单关键词-单文本”的词频字典集合，以及计数结果汇总
    """
    freq = dict()
    # 因为总频数是通过"+="的方式计算，不是简单赋值，所以要预设为0
    freq['Total Frequency'] = 0
    docs = docs.copy()  # 防止对函数之外的原样本框造成改动

    for i in range(0, len(docs)):
        # 对于每个文档，都形成一个字典，字典包括关键词在该文档出现的频数和上下文
        # id需要在第0列，正文需要在第2列
        freq['Doc' + str(docs.iloc[i, 0])] = make_doc_freq(word, docs.iloc[i, 2])
        # 在给每个文档形成字典的同时，对于总概率进行滚动加总
        freq['Total Frequency'] += freq['Doc' + str(docs.iloc[i, 0])]['Frequency']
        docs.iloc[i, 4] = freq['Doc' + str(docs.iloc[i, 0])]['Frequency']

    # 接下来建立一个DFC(doc-freq-context)统计面板，汇总所有文档对应的词频数和上下文
    # 首先构建(id, freq)的字典映射
    xs = docs['id']
    ys = docs['freq']
    # zip(迭代器)是一个很好用的方法，建议多用
    id_freq = {x: y for x, y in zip(xs, ys)}

    # 新建一个空壳DataFrame，接下来把数据一条一条粘贴进去
    data = pd.DataFrame(columns=['id', 'freq', 'word', 'num', 'context'])
    for item in xs:
        doc = freq['Doc' + str(item)]
        num = doc['Frequency']
        context = doc['Context']
        for i in range(0, num):
            strip = {'id': item, 'freq': id_freq[item], 'word': word, 'num': i, 'context': context[str(i)]}
            # 默认orient参数等于columns
            # 如果字典的值是标量，那就必须传递一个index，这是规定
            strip = pd.DataFrame(strip, index=[None])
            # df的append方法只能通过重新赋值来进行修改
            data = data.append(strip)
    data.set_index(['id', 'freq', 'word'], drop=True, inplace=True)
    freq['DFC'] = data
    return freq

def make_infos_freq(name, pattern, docs):
    """
    :param name: 指的是对其进行词频统计的形式
    :param pattern: 指的是对其进行词频统计的（正则表达式, 裁剪方法）
    :param docs: 是要遍历的文本的集合，必须是pandas DataFrame的形式，至少包含id列(iloc: 0)和正文列(iloc: 2)
    :return: 返回字典，其中包括“单关键词-单文本”的词频字典集合，以及计数结果汇总
    """
    freq = dict()
    # 因为总频数是通过"+="的方式计算，不是简单赋值，所以要预设为0
    freq['Total Frequency'] = 0
    docs = docs.copy()  # 防止对函数之外的原样本框造成改动

    items = range(0, len(docs))
    with alive_bar(len(items), force_tty=True, bar='circles') as bar:
        for i in items:
            # 对于每个文档，都形成一个字典，字典包括关键词在该文档出现的频数和上下文
            # id需要在第0列，正文需要在第2列
            # pattern 要全须全尾地传递进去，因为make_info_freq两个参数都要用
            freq['Doc' + str(docs.iloc[i, 0])] = make_info_freq(name, pattern, docs.iloc[i, 2])
            # 在给每个文档形成字典的同时，对于总概率进行滚动加总
            freq['Total Frequency'] += freq['Doc' + str(docs.iloc[i, 0])]['Frequency']
            docs.iloc[i, 4] = freq['Doc' + str(docs.iloc[i, 0])]['Frequency']
            bar()

    # 接下来建立一个DFC(doc-freq-context)统计面板，汇总所有文档对应的词频数和上下文
    # 首先构建(id, freq)的字典映射
    xs = docs['id']
    ys = docs['freq']
    # zip(迭代器)是一个很好用的方法，建议多用
    id_freq = {x: y for x, y in zip(xs, ys)}

    # 新建一个空壳DataFrame，接下来把数据一条一条粘贴进去
    data = pd.DataFrame(columns=['id', 'freq', 'form', 'word', 'num', 'context'])
    for item in xs:
        doc = freq['Doc' + str(item)]
        num = doc['Frequency']
        # 从（关键词，上下文）中取出两个元素
        context = doc['Context']
        for i in range(0, num):
            # context 中的关键词已经 match_cut 完毕，不需要重复处理
            strip = {'id': item, 'form': name, 'freq': id_freq[item], 'word': context[str(i)][0],
                     'num': i, 'context': context[str(i)][1]}
            # 默认orient参数等于columns
            # 如果字典的值是标量，那就必须传递一个index，这是规定
            strip = pd.DataFrame(strip, index=[None])
            # df的append方法只能通过重新赋值来进行修改
            data = data.append(strip)
    data.set_index(['id', 'freq', 'form', 'word'], drop=True, inplace=True)
    freq['DFC'] = data
    print(name + '  Completed')
    return freq

def words_docs_freq(words, docs):
    """
    :param words: 表示要对其做词频统计的关键词清单
    :param docs: 是要遍历的文本的集合，必须是pandas DataFrame的形式，至少包含id列、正文列、和频率列
    :return: 返回字典，其中包括“单关键词-多文本”的词频字典集合，以及最终的DFC(doc-frequency-context)和DTM(doc-term matrix)
    """
    freqs = dict()
    # 与此同时新建一个空壳DataFrame，用于汇总DFC
    data = pd.DataFrame()
    # 新建一个空壳，用于汇总DTM(Doc-Term-Matrix)
    dtm = pd.DataFrame(None, columns=words, index=docs['id'])
    # 来吧，一个循环搞定所有
    items = range(len(words))
    with alive_bar(len(items), force_tty=True, bar='blocks') as bar:
        for word in words:
            freq = make_docs_freq(word, docs)
            freqs[word] = freq
            data = data.append(freq['DFC'])
            for item in docs['id']:
                dtm.loc[item, word] = freq['Doc' + str(item)]['Frequency']
            bar()
    # 记得要sort一下，不然排序的方式不对（应该按照doc id来排列）
    data.sort_index(inplace=True)

    freqs['DFC'] = data
    freqs['DTM'] = dtm
    return freqs

def infos_docs_freq(infos, docs):
    """
    :param docs: 是要遍历的文本的集合，必须是pandas DataFrame的形式，至少包含id列和正文列
    :param infos: 指的是正则表达式的列表，格式为字典，key是示例，如“（1）”，value 是正则表达式，如“（[0-9]）”
    :return: 返回字典，其中包括“单关键词-多文本”的词频字典集合，以及最终的DFC(doc-frequency-context)和DTM(doc-term matrix)
    """
    freqs = dict()
    # 与此同时新建一个空壳DataFrame，用于汇总DFC
    data = pd.DataFrame()
    # 新建一个空壳，用于汇总DTM(Doc-Term-Matrix)
    dtm = pd.DataFrame(None, columns=list(infos.keys()), index=docs['id'])
    # 来吧，一个循环搞定所有
    items = range(len(infos))
    with alive_bar(len(items), force_tty=True, bar='blocks') as bar:
        for k, v in infos.items():
            freq = make_infos_freq(k, v, docs)
            freqs[k] = freq
            data = data.append(freq['DFC'])
            for item in docs['id']:
                dtm.loc[item, k] = freq['Doc' + str(item)]['Frequency']
            bar()
    # 记得要sort一下，不然排序的方式不对（应该按照doc id来排列）
    data.sort_index(inplace=True)

    freqs['DFC'] = data
    freqs['DTM'] = dtm
    return freqs

def massive_pop(infos, doc):
    """
    :param infos: List，表示被删除内容对应的正则表达式
    :param doc: 表示正文
    :return: 返回一个完成删除的文本
    """
    for info in infos:
        doc = re.sub(info, '', doc)
    return doc

def massive_sub(infos, doc):
    """
    :param infos: Dict, 表示被替换内容对应的正则表达式及替换对象
    :param doc: 表示正文
    :return: 返回一个完成替换的文本
    """
    for v, k in infos:
        doc = re.sub(v, k, doc)
    return doc

# 接下来取每个样本的前n句话(或者不多于前n句话的内容)，再做一次进行对比
# 取前十句话的原理是，对！？。等表示语义结束的符号进行计数，满十次为止
def top_n_sent(n, doc, percentile=1):
    """
    :param n: n指句子的数量，这个函数会返回一段文本中前n句话，若文本内容不多于n句，则全文输出
    :param word: 指正文内容
    :param percentile: 按照分位数来取句子时，要输入的分位，比如一共有十句话，取50%分位就是5句
    如果有11句话，向下取整也是输出5句
    :return: 返回字符串：前n句话
    """
    info = '[。？！]'
    # 在这个函数体内，函数主体语句的作用域大于循环体，因此循环内的变量相当于局部变量
    # 因此想在循环外直接返回，就会出现没有定义的错误，因此可以做一个全局声明
    # 但是不建议这样做，因为如果函数外有一个变量恰巧和局部变量重名，那函数外的变量也会被改变
    # 因此还是推荐多使用迭代器，把循环包裹成迭代器，可以解决很多问题
    # 而且已经封装好的迭代器，例如re.findall_iter，就不用另外再去写了，调用起来很方便
    # 如下，第一行代码的作用是用列表包裹迭代器，形成一个生成器的列表
    # 每个生成器都存在自己的 Attribute
    re_iter = list(re.finditer(info, doc))
    # max_iter 是 re 匹配到的最大次数
    max_iter = len(re_iter)

    # 这一句表示，正文过于简短，或者没有标点，此时直接输出全文
    if max_iter == 0:
        return doc

    # 考虑 percentile 的情况，如果总共有11句，就舍弃掉原来的 n，直接改为总句数的 percentile 对应的句子数
    # 注意是向下取整
    if percentile != 1:
        n = math.ceil(percentile * max_iter)
        # 如果匹配到至少一句，循环自然结束，输出结果
        if n > 0:
            return doc[0: re_iter[n - 1].end()]
        # 如果正文过于简短，或设定的百分比过低，一句话都凑不齐，此时直接输出第一句
        elif n == 0:
            return doc[0: re_iter[0].end()]

    # 如果匹配到的句子数大于 n，此时只取前 n 句
    if max_iter >= n:
        return doc[0: re_iter[n - 1].end()]
    # 如果匹配到的句子不足 n 句，直接输出全部内容
    elif 0 < max_iter < n:
        return doc[0: re_iter[-1].end()]

    # 为减少重名的可能，尽量在函数体内减少变量的使用

def dtm_sort_filter(dtm, keymap, name=None):
    """
    :param dtm: 前面生成的词频统计矩阵：Doc-Term-Matrix
    :param keymap: 字典，标明了  类别-关键词列表  两者关系
    :param name: 最终生成 Excel 文件的名称（需要包括后缀）
    :return: 返回一个字典，字典包含两个 pandas.DataFrame: 一个是表示各个种类是否存在的二进制表，另一个是最终的种类数
    """
    dtm = dtm.applymap(lambda x: 1 if x != 0 else 0)
    strips = {}
    for i, row in dtm.iterrows():
        strip = {}
        for k, v in keymap.items():
            strip[k] = 0
            for item in v:
                try:
                    strip[k] += row[item]
                except KeyError:
                    pass
        strips[i] = strip
    dtm_class = pd.DataFrame.from_dict(strips, orient='index')
    dtm_class = dtm_class.applymap(lambda x: 1 if x != 0 else 0)
    dtm_final = dtm_class.agg(np.sum, axis=1)
    result = {'DTM_class': dtm_class, 'DTM_final': dtm_final}
    return result

def dtm_point_giver(dtm, keymap, scoremap, name=None):
    """
    :param dtm: 前面生成的词频统计矩阵：Doc-Term-Matrix
    :param keymap: 字典，{TypeA: [word1, word2, word3, ……], TypeB: ……}
    :param scoremap: 字典，标明了  类别-分值 两者关系
    :param name: 最终生成 Excel 文件的名称（需要包括后缀）
    :return: 返回一个 pandas.DataFrame，表格有两列，一列是文本id，一列是文本的分值（所有关键词的分值取最高）
    """
    dtm = dtm.applymap(lambda x: 1 if x != 0 else 0)

    # 非 keymap 中词会被过滤掉
    strips = {}
    for i, row in dtm.iterrows():
        strip = {}
        for k, v in keymap.items():
            strip[k] = 0
            for item in v:
                try:
                    strip[k] += row[item]
                except KeyError:
                    pass
        strips[i] = strip

    dtm_class = pd.DataFrame.from_dict(strips, orient='index')
    dtm_class = dtm_class.applymap(lambda x: 1 if x != 0 else 0)

    # 找到 columns 对应的分值
    keywords = list(dtm_class.columns)
    multiplier = []
    for keyword in keywords:
        multiplier.append(scoremap[keyword])

    # DataFrame 的乘法运算，不会改变其 index 和 columns
    dtm_score = dtm_class.mul(multiplier, axis=1)
    # 取一个最大值来赋分
    dtm_score = dtm_score.agg(np.max, axis=1)
    return dtm_score

def dfc_sort_filter(dfc, keymap, name=None):
    """
    :param dfc: 前面生成的词频统计明细表：Doc-Frequency-Context
    :param keymap: 字典，标明了  关键词-所属种类  两者关系
    :param name: 最终生成 Excel 文件的名称（需要包括后缀）
    :return: 返回一个 pandas.DataFrame，表格有两列，一列是文本id，一列是文本中所包含的业务种类数
    """
    # 接下来把关键词从 dfc 的 Multi-index 中拿出来（这个index本质上就是一个ndarray)
    # 拿出来关键词就可以用字典进行映射
    # 先新建一列class-id，准备放置映射的结果
    dfc.insert(0, 'cls-id', None)

    # 开始遍历
    for i in range(0, len(dfc.index)):
        dfc.iloc[i, 0] = keymap[dfc.index[i][2]]

    # 理论上就可以直接通过 excel 的分类计数功能来看业务种类数了
    # 失败了，excel不能看种类数，只能给所有值做计数，因此还需要借助python的unique语句
    # dfc.to_excel('被监管业务统计.xlsx')

    # 可以对于每一种index做一个计数，使用loc索引到的对象是一个DataFrame
    # 先拿到一个doc id的列表
    did = []
    for item in dfc.index.unique():
        did.append(item[0])
    did = list(pd.Series(did).unique())

    # 接下来获得每一类的结果，注：多重索引的取值值得关注
    uni = {}
    for item in did:
        uni[item] = len(dfc.loc[item, :, :]['cls-id'].unique())

    # 把生成的字典转换为以键值行索引的 DataFrame
    uni = pd.DataFrame.from_dict(uni, orient='index')
    uni.fillna(0, axis=1, inplace=True)
    # uni.to_excel(name)
    return uni

def dfc_point_giver(dfc, keymap, name=None):
    """
    :param dfc: 前面生成的词频统计明细表：Doc-Frequency-Context
    :param keymap: 字典，标明了  关键词-分值 两者关系
    :param name: 最终生成 Excel 文件的名称（需要包括后缀）
    :return: 返回一个 pandas.DataFrame，表格有两列，一列是文本id，一列是文本的分值（所有关键词的分值取最高）
    """
    dfc.insert(0, 'point', None)

    # 开始遍历
    for i in range(0, len(dfc.index)):
        dfc.iloc[i, 0] = keymap[dfc.index[i][2]]

    # 可以对于每一种index做一个计数，使用loc索引到的对象是一个DataFrame
    # 先拿到一个doc id的列表
    did = []
    for item in dfc.index.unique():
        did.append(item[0])
    did = list(pd.Series(did).unique())

    # 接下来获得每一类的结果，注：多重索引的取值值得关注
    uni = {}
    for item in did:
        uni[item] = max(dfc.loc[item, :, :]['point'].unique())

    # 把生成的字典转换为以键值行索引的 DataFrame
    uni = pd.DataFrame.from_dict(uni, orient='index')
    uni.fillna(0, axis=1, inplace=True)
    # uni.to_excel(name)
    return uni

def dfc_sort_counter(dfc, name=None):
    """
    :param dfc: 前面生成的词频统计明细表：Doc-Frequency-Context
    :param name: 最终生成 Excel 文件的名称（需要包括后缀）
    :return: 返回一个 pandas.DataFrame，表格有两列，一列是文本id，一列是文本中所包含的业务种类数
    """
    # 可以对于每一种index做一个计数，使用loc索引到的对象是一个DataFrame
    dfc.insert(0, 'form', None)
    for i in range(0, dfc.shape[0]):
        dfc.iloc[i, 0] = dfc.index[i][2]

    # 先拿到一个doc id的列表
    did = []
    for item in dfc.index.unique():
        did.append(item[0])
    did = list(pd.Series(did).unique())

    # 接下来获得每一类的结果，注：多重索引的取值值得关注
    uni = {}
    for item in did:
        uni[item] = len(dfc.loc[item, :, :, :]['form'].unique())

    # 把生成的字典转换为以键值行索引的 DataFrame
    uni = pd.DataFrame.from_dict(uni, orient='index')
    uni.fillna(0, axis=1, inplace=True)
    # uni.to_excel(name)
    return uni

# 定义一个大写中文数字转阿拉伯数字的函数
CN_NUM = {
    '〇': 0, '一': 1, '二': 2, '三': 3, '四': 4, '五': 5, '六': 6, '七': 7, '八': 8, '九': 9, '零': 0,
    '壹': 1, '贰': 2, '叁': 3, '肆': 4, '伍': 5, '陆': 6, '柒': 7, '捌': 8, '玖': 9, '貮': 2, '两': 2,
}

CN_UNIT = {
    '十': 10,
    '拾': 10,
    '百': 100,
    '佰': 100,
    '千': 1000,
    '仟': 1000,
    '万': 10000,
    '万': 10000,
    '亿': 100000000,
    '亿': 100000000,
    '兆': 1000000000000,
}

def chinese_to_arabic(cn:str) -> str:
    unit = 0  # current
    ldig = []  # digest
    for cndig in reversed(cn):
        if cndig in CN_UNIT:
            unit = CN_UNIT.get(cndig)
            if unit == 10000 or unit == 100000000:
                ldig.append(unit)
                unit = 1
        else:
            dig = CN_NUM.get(cndig)
            if unit:
                dig *= unit
                unit = 0
            ldig.append(dig)
    if unit == 10:
        ldig.append(10)
    val, tmp = 0, 0
    for x in reversed(ldig):
        if x == 10000 or x == 100000000:
            val += tmp * x
            tmp = 0
        else:
            tmp += x
    val += tmp
    val = str(val)
    return val

def match_cut(matchobj, cls):
    """
    :param matchobj: 输入匹配好的字符串
    :param cls: 这是正则表达式的匹配方法，遵循掐头去尾的原则。
    0：原样奉还，1：掐一去一，2：掐一去二，3：掐二去一，4：掐一去零,  5：掐零去一
    :return: 返回一个处理好的字符串
    """
    if cls == 0:
        return matchobj
    if cls == 1:
        return matchobj[1: -1]
    if cls == 2:
        return matchobj[1: -2]
    if cls == 3:
        return matchobj[2, -1]
    if cls == 4:
        return matchobj[1:]
    if cls == 5:
        return matchobj[: -1]

def re_formatter(file_time):
    """
    这个函数是修改数据格式的一个小工具，会将输入的时间统一由YYYY-DD-MM修改为YYYY/MM/DD
    :param file_time: 输入的时间, 纯字符串格式
    :return: 返回的时间，datetime格式(%YYYY/%MM/%DD)
    """
    # 首先把空格和空行删去
    file_time = str(file_time)
    file_time = file_time.strip('\n\r\t ')
    file_time = file_time[:10]
    # 将字符串转为datetime格式
    file_time = datetime.datetime.strptime(file_time, '%Y-%m-%d')
    # 将datetime转为字符串
    file_time = datetime.datetime.strftime(file_time, '%Y/%m/%d')
    return file_time

def year_month_finder(file_time):
    """
   这个函数是修改数据格式的一个小工具，会将输入的时间统一由YYYY/DD/MM修改为YYYY年M月
    :param file_time: 输入的时间, 纯字符串格式
    :return: 返回的时间，格式为(%YYYY年%M月)
    """
    # 首先把空格和空行删去
    file_time = str(file_time)
    file_time = file_time.strip('\n\r\t ')
    # 将字符串转为datetime格式
    try:
        file_time = datetime.datetime.strptime(file_time, '%Y/%m/%d')
        file_time_year = datetime.datetime.strftime(file_time, '%Y')
        file_time_month = datetime.datetime.strftime(file_time, '%m')
        file_time = file_time_year + 'M' + file_time_month
    except ValueError:
        pass
    return file_time

def month_day_finder(file_time):
    """
   这个函数是修改数据格式的一个小工具，会将输入的时间统一由YYYY/DD/MM修改为M月D日
    :param file_time: 输入的时间, 纯字符串格式
    :return: 返回的时间，格式为(%M%月/%D%日)
    """
    # 首先把空格和空行删去
    file_time = str(file_time)
    file_time = file_time.strip('\n\r\t ')
    # 将字符串转为datetime格式
    try:
        file_time = datetime.datetime.strptime(file_time, '%Y/%m/%d')
        file_time_month = datetime.datetime.strftime(file_time, '%m')
        file_time_day = datetime.datetime.strftime(file_time, '%d')
        file_time = file_time_month + '/' + file_time_day
    except ValueError:
        pass
    return file_time

quarter_map = {'1': '1', '2': '1', '3': '1', '4': '2', '5': '2', '6': '2',
               '7': '3', '8': '3', '9': '3', '10': '4', '11': '4', '12': '4',
               '01': '1', '02': '1', '03': '1', '04': '2', '05': '2', '06': '2',
               '07': '3', '08': '3', '09': '3'}

def quarter_finder(file_time):
    """
    这个函数是修改数据格式的一个小工具，会将输入的时间统一由YYYY/DD/MM修改为YYYY年Q季度
    :param file_time: 输入的时间, 纯字符串格式
    :return: 返回的时间，格式为(%YYYY年%Q季度)
    """
    # 首先把空格和空行删去
    file_time = str(file_time)
    file_time = file_time.strip('\n\r\t ')
    # 将字符串转为datetime格式
    try:
        file_time = datetime.datetime.strptime(file_time, '%Y/%m/%d')
        file_time_year = datetime.datetime.strftime(file_time, '%Y')
        file_time_month = datetime.datetime.strftime(file_time, '%m')
        file_time_month = quarter_map[file_time_month]
        file_time = file_time_year + 'Q' + file_time_month
    except ValueError:
        pass
    return file_time

def dataframe_filter(df, keywords, axis, status):
    """
    :param df: 要过滤的样本框
    :param keywords: list, 关键词清单: 要保留/剔除的 index 或 columns
    :param axis: integer, 过滤的方向，axis=0 为竖直方向，axis=1 为水平方向
    :param status: integer, 过滤的方式，status=0 为只保留清单中的数据，status=1 为只去除清单中的数据
    :return: 过滤好的样本框
    """
    df = df.copy()
    # 如果是只保留清单中的数据
    if status == 0:
        if axis == 1:
            lst = df.columns
            items = range(len(lst))
            with alive_bar(len(items), force_tty=True, bar='blocks') as bar:
                for item in lst:
                    if item not in keywords:
                        df.drop([item], axis=1, inplace=True)
                    bar()
        elif axis == 0:
            lst = df.index
            items = range(len(lst))
            with alive_bar(len(items), force_tty=True, bar='blocks') as bar:
                for item in lst:
                    if item not in keywords:
                        df.drop(item, axis=0, inplace=True)
                    bar()
    # 如果是只去除清单中的数据
    elif status == 1:
        if axis == 1:
            lst = df.columns
            items = range(len(lst))
            with alive_bar(len(items), force_tty=True, bar='blocks') as bar:
                for item in lst:
                    if item in keywords:
                        df.drop([item], axis=1, inplace=True)
                    bar()
        elif axis == 0:
            lst = df.index
            items = range(len(lst))
            with alive_bar(len(items), force_tty=True, bar='blocks') as bar:
                for item in lst:
                    if item in keywords:
                        df.drop(item, axis=0, inplace=True)
                    bar()
    else:
        raise KeyError('Axis not properly set')
    return df

def main():
    # 主程序运行，开始！
    # 读取文件
    os.chdir('E:/ANo.3/base')
    app1 = xw.App(visible=False, add_book=False)
    wb = app1.books.open('20210116样本分筛.xlsx')
    try:
        sht = wb.sheets('0次')
        df = pd.DataFrame(sht.used_range.value)
        df.columns = list(df.loc[0])
        df.drop(0, axis=0, inplace=True)
        df.reset_index(drop=True, inplace=True)
        # 记得把id变成整数，否则之后索引字符化后会有麻烦
        df['id'] = df['id'].apply(lambda x: int(x))
        # 提前插入一列频率，方便遍历时记录结果
        df.insert(4, 'freq', None)

        # 获取关键词清单
        f = open('add_words_dict.txt', 'r', encoding='UTF-8')
        n = f.read()
        f.close()
        n = n.split('\n')

        freqs = words_docs_freq(n, df)
        wb = app1.books.add()
        sht = wb.sheets('Sheet1')
        sht.name = '词频统计表'
        sht['A1'].value = freqs['DTM']
        sht = wb.sheets.add('词频统计明细')
        sht['A1'].value = freqs['DFC']
        wb.save('20210126简单词频统计.xlsx')
    finally:
        app1.quit()

def list_to_txt(name, content, sep = '\n'):
    """
    :param name: 指要创建的 txt 文件名称， 不存在的话会自动创建
    :param content: 指要写入的内容，此处必须为列表
    :param sep: 指列表的分隔符，默认为换行符
    :return: 无返回值
    这个函数可以把一个列表写入txt, 元素之间自动以换行符形式进行粘连
    """
    # 首先把列表使用分隔符进行粘连
    content = sep.join(content)
    # 使用with open方式打开，可以保证结束操作的同时关闭句柄
    with open(name, 'w', encoding='UTF-8') as f:
        f.write(content)

def txt_to_list(path, sep = '\n'):
    """
    :param path: 指要读取的 txt 文件所在位置， 不存在的话会报错，path 支持列表格式
    :param sep: 指列表的分隔符，默认为换行符
    :return: 返回一个列表
    这个函数可以把一个 txt 转化为列表 t, 导入的元素之间需要以制定换行符形式分隔开
    如果同时传入多个txt，则自动拼接在同一个列表当中
    """
    if path is not list:
        path = [path]
    lst = []
    for link in path:
        # 使用with open方式打开，可以保证结束操作的同时关闭句柄
        with open(link, 'r', encoding='UTF-8') as f:
            content = str(f.read())
        content = content.split(sep)
        try:
            content.remove('\n')
            content.remove('\t')
            content.remove('\r')
            content.remove('')
            content.remove(' ')
            content.remove('  ')
        except ValueError:
            pass
        lst.extend(content)
    return lst

def doc_filter(df, target):
    """
    :param df: 要进行样本筛选的数据框
    :param target: 符合要求的id列表
    :return: 只保留符合要求id对应样本的数据框
    """
    for doc_id in df['id']:
        if doc_id not in target:
            df.drop(df[df['id'] == doc_id].index, inplace=True)
    return df

def vec_cos(arr1, arr2):
    """
    :param arr1:  the first np.array
    :param arr2:  the second np.array
    :return:  cos(θ) (θ: the angle between two vectors)

    This measures the similarity between two vectors: the closer the value is to 1, the more similar between two vectors
    """
    num = np.sum(arr1 * arr2)
    den = np.sqrt(sum(arr1**2) * sum(arr2**2))
    sim = num/den
    return sim

def cos_rank(matrix, keymap=None, threshold=0.9):
    """
    :param matrix: (np.array) it contains all the keywords frequency vectors
    :param keymap: (dictionary) you can pass a preset dictionary if you wanna keep the original vector id
    :param threshold: (float) you can pass a threshold of cosθ above which the id will be recorded
    :return df: (pd.DataFrame) the rank of similarity

    This function splits the vectors into pairs and calculate the cos similarity pair by pair.
    Then it will generate a DataFrame that ranks the pairs from 1 to 0
    """
    pairs = []
    for i in range(0, len(matrix), 1):
        for t in range(i+1, len(matrix), 1):
            pairs.append((i, t))
    sim = []
    if keymap is None:
        for pair in pairs:
            sim.append((pair, vec_cos(matrix[pair[0]], matrix[pair[1]])))
    else:
        for pair in pairs:
            sim.append(((keymap[pair[0]], keymap[pair[1]]), vec_cos(matrix[pair[0]], matrix[pair[1]])))

    df = pd.DataFrame(sim)
    df.columns = ['id', 'Similarity']
    df.sort_values('Similarity', axis=0, ascending=False, inplace=True)

    # You can choose desired filtering condition
    df = df[df['Similarity'] > threshold]

    return df

def quarter2date(date):
    # 将字符串 date 转化为 datetime 格式
    quarter = int(re.findall('Q[0-9]', date)[0][1])
    month = {1: 1, 2: 4, 3: 7, 4: 10}
    rule = re.compile('Q[0-9]')
    year = rule.sub('', date)
    return datetime.date(int(year), month[quarter], 1)

if __name__ == '__main__':
    main()


################################################################################
##
## BY: WANDERSON M.PIMENTA
## PROJECT MADE WITH: Qt Designer and PySide2
## V: 1.0.0
##
################################################################################

import sys
import platform
from PySide2 import QtCore, QtGui, QtWidgets
from PySide2.QtCore import (QCoreApplication, QPropertyAnimation, QDate, QDateTime, QMetaObject, QObject, QPoint, QRect, QSize, QTime, QUrl, Qt, QEvent)
from PySide2.QtGui import (QBrush, QColor, QConicalGradient, QCursor, QFont, QFontDatabase, QIcon, QKeySequence, QLinearGradient, QPalette, QPainter, QPixmap, QRadialGradient)
from PySide2.QtWidgets import *

# GUI FILE
from ui_splash_screen import Ui_SplashScreen
from ui_main import Ui_MainWindow

# GLOBALS
counter = 0
jumper = 10

## ==> YOUR APPLICATION WINDOW
class MainWindow(QMainWindow):
    def __init__(self):
        QMainWindow.__init__(self)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        ## ==> SET VALUES TO DEF progressBarValue
        def setValue(self, slider, labelPercentage, progressBarName, color):

            # GET SLIDER VALUE
            value = slider.value()

            # CONVERT VALUE TO INT
            sliderValue = int(value)

            # HTML TEXT PERCENTAGE
            htmlText = """
{VALUE}%
"""
            labelPercentage.setText(htmlText.replace("{VALUE}", str(sliderValue)))

            # CALL DEF progressBarValue
            self.progressBarValue(sliderValue, progressBarName, color)

        ## ==> APPLY VALUES TO PROGREESBAR
        self.ui.sliderCPU.valueChanged.connect(lambda: setValue(self, self.ui.sliderCPU, self.ui.labelPercentageCPU, self.ui.circularProgressCPU, "rgba(85, 170, 255, 255)"))
        self.ui.sliderGPU.valueChanged.connect(lambda: setValue(self, self.ui.sliderGPU, self.ui.labelPercentageGPU, self.ui.circularProgressGPU, "rgba(85, 255, 127, 255)"))
        self.ui.sliderRAM.valueChanged.connect(lambda: setValue(self, self.ui.sliderRAM, self.ui.labelPercentageRAM, self.ui.circularProgressRAM, "rgba(255, 0, 127, 255)"))

        ## ==> DEF START VALUES
        self.ui.sliderCPU.setValue(25)
        self.ui.sliderGPU.setValue(65)
        self.ui.sliderRAM.setValue(45)

    ## DEF PROGRESS BAR VALUE
    ########################################################################
    def progressBarValue(self, value, widget, color):

        # PROGRESSBAR STYLESHEET BASE
        styleSheet = """
        QFrame{
        	border-radius: 110px;
        	background-color: qconicalgradient(cx:0.5, cy:0.5, angle:90, stop:{STOP_1} rgba(255, 0, 127, 0), stop:{STOP_2} {COLOR});
        }
        """

        # GET PROGRESS BAR VALUE, CONVERT TO FLOAT AND INVERT VALUES
        # stop works of 1.000 to 0.000
        progress = (100 - value) / 100.0

        # GET NEW VALUES
        stop_1 = str(progress - 0.001)
        stop_2 = str(progress)

        # FIX MAX VALUE
        if value == 100:
            stop_1 = "1.000"
            stop_2 = "1.000"

        # SET VALUES TO NEW STYLESHEET
        newStylesheet = styleSheet.replace("{STOP_1}", stop_1).replace("{STOP_2}", stop_2).replace("{COLOR}", color)

        # APPLY STYLESHEET WITH NEW VALUES
        widget.setStyleSheet(newStylesheet)

## ==> SPLASHSCREEN WINDOW
class SplashScreen(QMainWindow):
    def __init__(self):
        QMainWindow.__init__(self)
        self.ui = Ui_SplashScreen()
        self.ui.setupUi(self)

        ## ==> SET INITIAL PROGRESS BAR TO (0) ZERO
        self.progressBarValue(0)

        ## ==> REMOVE STANDARD TITLE BAR
        self.setWindowFlags(QtCore.Qt.FramelessWindowHint) # Remove title bar
        self.setAttribute(QtCore.Qt.WA_TranslucentBackground) # Set background to transparent

        ## ==> APPLY DROP SHADOW EFFECT
        self.shadow = QGraphicsDropShadowEffect(self)
        self.shadow.setBlurRadius(20)
        self.shadow.setXOffset(0)
        self.shadow.setYOffset(0)
        self.shadow.setColor(QColor(0, 0, 0, 120))
        self.ui.circularBg.setGraphicsEffect(self.shadow)

        ## QTIMER ==> START
        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.progress)
        # TIMER IN MILLISECONDS
        self.timer.start(15)

        ## SHOW ==> MAIN WINDOW
        ########################################################################
        self.show()
        ## ==> END ##

    ## DEF TO LOANDING
    ########################################################################
    def progress (self):
        global counter
        global jumper
        value = counter

        # HTML TEXT PERCENTAGE
        htmlText = """
{VALUE}%
"""

        # REPLACE VALUE
        newHtml = htmlText.replace("{VALUE}", str(jumper))

        if(value > jumper):
            # APPLY NEW PERCENTAGE TEXT
            self.ui.labelPercentage.setText(newHtml)
            jumper += 10

        # SET VALUE TO PROGRESS BAR
        # fix max value error if > than 100
        if value >= 100: value = 1.000
        self.progressBarValue(value)

        # CLOSE SPLASH SCREE AND OPEN APP
        if counter > 100:
            # STOP TIMER
            self.timer.stop()

            # SHOW MAIN WINDOW
            self.main = MainWindow()
            self.main.show()

            # CLOSE SPLASH SCREEN
            self.close()

        # INCREASE COUNTER
        counter += 0.5

    ## DEF PROGRESS BAR VALUE
    ########################################################################
    def progressBarValue(self, value):

        # PROGRESSBAR STYLESHEET BASE
        styleSheet = """
        QFrame{
        	border-radius: 150px;
        	background-color: qconicalgradient(cx:0.5, cy:0.5, angle:90, stop:{STOP_1} rgba(255, 0, 127, 0), stop:{STOP_2} rgba(85, 170, 255, 255));
        }
        """

        # GET PROGRESS BAR VALUE, CONVERT TO FLOAT AND INVERT VALUES
        # stop works of 1.000 to 0.000
        progress = (100 - value) / 100.0

        # GET NEW VALUES
        stop_1 = str(progress - 0.001)
        stop_2 = str(progress)

        # SET VALUES TO NEW STYLESHEET
        newStylesheet = styleSheet.replace("{STOP_1}", stop_1).replace("{STOP_2}", stop_2)

        # APPLY STYLESHEET WITH NEW VALUES
        self.ui.circularProgress.setStyleSheet(newStylesheet)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = SplashScreen()
    sys.exit(app.exec_())


# this is the Code for 3rd part of the Assignment

import random
import matplotlib.pyplot as plt
import numpy as np

def RC_Multiply(a,b):
	sum = int("0",10)
	for i in range(3):
		sum += a[i]*b[i]
	return sum;

Data = []
plt_dt = []
a = 0
for i in range(100):
	r_fra1 = random.randint(-1000,1000)/1000;
	r_fra2 = random.randint(-1000,1000)/1000;
	dat_plt = [r_fra1, r_fra2]
	plt_dt.append(dat_plt)
	a = [r_fra1, r_fra2,1]
	if(i < 50):
		clas = "A"
	else:
		clas = "B"
	Data.append([a,clas])

w1 = [1,1,0]  #plot this line
w2 = [-1,-1,0]
w3 = [0,0.5,0]
w4 = [1, -1,5] # plot this line
w5 = [1,1,0.3] # plot this line

acc_w1 = 0;
acc_w2 = 0;
acc_w3 = 0;
acc_w4 = 0;
acc_w5 = 0;
for i in range(100):
	if(RC_Multiply(w1,Data[i][0]) > 0 and Data[i][1]=="A"):
		acc_w1 += 1;
	elif(RC_Multiply(w1,Data[i][0]) <= 0 and Data[i][1]=="B"):
		acc_w1 += 1;
	if(RC_Multiply(w1,Data[i][0]) > 0 and Data[i][1]=="A"):
		acc_w2 += 1;
	elif(RC_Multiply(w2,Data[i][0]) <= 0 and Data[i][1]=="B"):
		acc_w2 += 1;
	if(RC_Multiply(w3,Data[i][0]) > 0 and Data[i][1]=="A"):
		acc_w3 += 1;
	elif(RC_Multiply(w3,Data[i][0]) <= 0 and Data[i][1]=="B"):
		acc_w3 += 1;
	if(RC_Multiply(w4,Data[i][0]) > 0 and Data[i][1]=="A"):
		acc_w4 += 1;
	elif(RC_Multiply(w4,Data[i][0]) <= 0 and Data[i][1]=="B"):
		acc_w4 += 1;
	if(RC_Multiply(w5,Data[i][0]) > 0 and Data[i][1]=="A"):
		acc_w5 += 1;
	elif(RC_Multiply(w5,Data[i][0]) <= 0 and Data[i][1]=="B"):
		acc_w5 += 1;

print("PART (i)")
print("Accuracy of a : "+str(acc_w1)+"%")
print("Accuracy of b : "+str(acc_w2)+"%")
print("Accuracy of c : "+str(acc_w3)+"%")
print("Accuracy of d : "+str(acc_w4)+"%")
print("Accuracy of e : "+str(acc_w5)+"%")

print("PART (ii)")
fig, ax = plt.subplots()
data_plot = np.array(plt_dt);
x,y = data_plot.T
ax.scatter(x,y)
x = np.array(range(-1,2))
y = (w1[0]/-w1[1])*x + (w1[2]/-w1[1])
la = ax.plot(x,y,label='line for a')
ax.set_title('With Boundry a')

fig,bx = plt.subplots()
data_plot = np.array(plt_dt);
x,y = data_plot.T
bx.scatter(x,y)
x = np.array(range(-1,2))
y = (w4[0]/-w4[1])*x + (w4[2]/-w4[1])
bx.plot(x,y,label='line for d')
bx.set_title('With Boundry b')

fig,cx = plt.subplots()
data_plot = np.array(plt_dt);
x,y = data_plot.T
cx.scatter(x,y)
x = np.array(range(-1,2))
y = (w5[0]/-w5[1])*x + (w5[2]/-w5[1])
le = cx.plot(x,y,label='line for e')
cx.set_title('With Boundry c')
plt.show()


from sys import (stdout, stderr)
from argparse import (ArgumentParser, ArgumentError)

from .color_display import ColorDisplay

########################################################################
# Command-line arguments

timeout_dft = -1

parser = ArgumentParser(
        description="Python script to show off terminal colors.",
        epilog="Run this script from the terminal whose colors " +
               "you want to showcase. " +
               "Only xterm-like terminals are supported " +
               "(see the README).")

mode_group = parser.add_mutually_exclusive_group()

p_choices = [16, 88, 256]

arg_p = mode_group.add_argument(
        '-p', '--pretty',
        action='store_true', default=False,
        help="show colors 0 through N-1 in a pretty format.  " +
             ("N must belong to %r.  " % p_choices) +
             "If N > 16, it should be the actual number of colors " +
             "supported by the terminal, or the output will almost " +
             "certainly not be pretty.")

mode_group.add_argument(
        '-f', '--flat',
        action='store_true', default=False,
        help="show a simple table with colors 0 through N-1.  ")

parser.add_argument(
        'n', nargs='?', metavar='N',
        type=int, default=16,
        help="number of colors to show.  " +
             "Unless you explicitly supply -p/--pretty or -f/--flat, " +
             "--pretty is used if possible and --flat is used " +
             "otherwise.  " +
             "N defaults to 16, showing the ANSI colors 0-15.  " +
             "If N is 0, the script will attempt to determine the " +
             "maximum number of colors automatically " +
             "(which may be slow).")

parser.add_argument(
        '--no-fgbg',
        action='store_false', dest='fgbg', default=True,
        help="suppress display of foreground/background colors.")

parser.add_argument(
        '--no-query',
        action='store_false', dest='do_query', default=True,
        help="don't try to query any RGB values from the terminal " +
             "and just use placeholders.")

parser.add_argument(
        '-t', '--timeout', metavar='T',
        type=int, default=timeout_dft,
        help="how long to wait for the terminal to "
             "respond to a query, in milliseconds  " +
             "[default: {0}].  ".format(timeout_dft) +
             "If your output has '?' characters " +
             "instead of RGB values " +
             "or junk printed after the script runs, " +
             "increasing this value may or may not " +
             "help, depending on the terminal.  " +
             "A negative T will behave like infinity.")

parser.add_argument(
        '-l', '--level', metavar='L',
        type=int, default=3,
        help="choose how much color to use in the output.  " +
             "(0 = no color; 3 = most color [default])")

parser.add_argument(
        '--screen-forward', '--tmux-forward',
        action='store_true', default=False,
        help="attempt to pass terminal queries through a tmux " +
             "session (ignored if this is not a tmux session)")

########################################################################

def main():
    args = parser.parse_args()

    assert not (args.pretty and args.flat)

    if args.pretty:
        if args.n not in p_choices:
            raise ArgumentError(
                    arg_p,
                    "N must belong to %r" % p_choices)

    with ColorDisplay(0, args.timeout, args.level, args.do_query,
                      args.screen_forward) as C:

        if args.n == 0:
            args.n = C.get_num_colors(args.timeout)

        if not (args.pretty or args.flat):
            if args.n in p_choices:
                args.pretty = True
            else:
                args.flat = True

        if args.level >= 1:
            stdout.write(C.reset)

        if args.fgbg:
            C.show_fgbg()

        if args.pretty:
            assert args.n in p_choices

            stdout.write('\n    ANSI colors:\n\n')
            C.show_ansi()

            if args.n > 16:
                stdout.write('\n    RGB cube:\n\n')
                C.show_color_cube(args.n)

                stdout.write('    Grayscale ramp:\n\n')
                C.show_grayscale_ramp(args.n)
        else:
            C.show_colors(args.n)

        if C.num_errors > 0:
            stderr.write("Warning: not all queries succeeded\n" +
                         "Warning:     (output contains " +
                         "placeholders and may be inaccurate)\n")

if __name__ == "__main__":
    main()


import numpy as np
import matplotlib.pyplot as plt
import os
import pylab

from mpl_toolkits.mplot3d import Axes3D

# experiments variables
datasets = ["kotanchek-test-0"]
schemes = ["proximity-xy", "surrounding-xy", "remoteness-xy", "nonlinearity"]
dist_metrics = ["L2.0"]
nums_neighbors = ["k5"]

# local paths
base_dir = os.path.expanduser("~/Dropbox/my_files/research/ISR/")
data_folder = base_dir + "datasets/original/"
weights_folder = base_dir + "datasets/2-08.06.2017/original/weights/"
plots_folder = base_dir + "plots/5-01.09.2017/original/"

if not os.path.exists(plots_folder):
    os.makedirs(plots_folder)

for dataset in datasets:
    print("Dataset: " + dataset)

    fold = data_folder + dataset + ".csv"

    for scheme in schemes:
        print("  Scheme: " + scheme)

        for num_neighbors in nums_neighbors:
            for dist_metric in dist_metrics:
                # for the nonlinearity scheme only the Euclidean distance was used
                if scheme == "nonlinearity" and dist_metric != "L2.0":
                    continue

                weights = weights_folder + "-".join([scheme, num_neighbors, dist_metric, dataset]) + ".csv"

                print("    Distance metric: " + dist_metric)
                print("    Data file: " + fold)
                print("    Weight file: " + weights)

                # read and process the datasets
                data = np.loadtxt(fold, delimiter=",")
                X = data[:, 0]
                Y = data[:, 1]
                Z = data[:, 2]

                # read and process the weight files
                W = np.loadtxt(weights, delimiter=",")
                labels = ['%.2f' % elem for elem in W.tolist()]
                min_weight = min(W) * 0.95
                max_weight = max(W) * 0.95

                areas = [(w - min_weight) / (max_weight - min_weight) * 400 for w in W]

                plt.figure(figsize=(16, 8))

                fig = pylab.figure(figsize=(16, 10))
                ax = Axes3D(fig)

                # version with black rings: ax.scatter(X, Y, Z, s=areas, facecolors='none', edgecolors='black')
                ax.scatter(X, Y, Z, s=areas, c="#1F77B4", facecolors="none", alpha=0.4)

                plt.xlabel("X", fontsize=12)
                plt.ylabel("Y", fontsize=12)

                plt.title("Weights - " + scheme + ", " + dist_metric + ", " + dataset)
                plt.savefig(plots_folder + "-".join([dataset, scheme, num_neighbors, dist_metric]) + ".pdf",
                            bbox_inches="tight")

                plt.close()


from sqlalchemy import asc
from sqlalchemy import func
from transaction import manager

from parks.logic.math_logic import latitude_to_miles
from parks.logic.math_logic import latitude_longitude_distance
from parks.logic.math_logic import longitude_to_miles
from parks.models import DBSession
from parks.models import Park
from parks.models import Stamp
from parks.models import StampCollection
from parks.models import StampHistory
from parks.models import StampLocation
from parks.models import StampToLocation

def get_stamps_by_park_id(park_id):
    max_time_subquery = DBSession.query(
        StampCollection.stamp_id.label('stamp_id'),
        func.max(StampCollection.date_collected).label('most_recent_time'),
    ).group_by(
        StampCollection.stamp_id
    ).subquery()

    return DBSession.query(
        Stamp,
        StampLocation,
        max_time_subquery.c.most_recent_time,
    ).join(
        StampToLocation,
    ).join(
        StampLocation,
    ).outerjoin(
        max_time_subquery,
        Stamp.id == max_time_subquery.c.stamp_id
    ).filter(
        StampLocation.park_id == park_id
    ).group_by(
        Stamp.id
    ).group_by(
        # This group_by needs to be included for the query to run in Postgres
        StampLocation.id
    ).group_by(
        # This group_by needs to be included for the query to run in Postgres
        max_time_subquery.c.most_recent_time
    ).all()

def get_stamps_by_location_id(location_id):
    max_time_subquery = DBSession.query(
        StampCollection.stamp_id.label('stamp_id'),
        func.max(StampCollection.date_collected).label('most_recent_time'),
    ).group_by(
        StampCollection.stamp_id
    ).subquery()

    return DBSession.query(
        Stamp,
        StampLocation,
        max_time_subquery.c.most_recent_time,
    ).join(
        StampToLocation,
    ).join(
        StampLocation,
    ).outerjoin(
        max_time_subquery,
        Stamp.id == max_time_subquery.c.stamp_id,
    ).filter(
        StampLocation.id == location_id,
    ).group_by(
        Stamp.id
    ).group_by(
        # This group_by needs to be included for the query to run in Postgres
        StampLocation.id
    ).group_by(
        # This group_by needs to be included for the query to run in Postgres
        max_time_subquery.c.most_recent_time
    ).all()

def get_nearby_stamps(latitude, longitude, distance_miles):
    lower_latitude = latitude - (distance_miles / latitude_to_miles())
    upper_latitude = latitude + (distance_miles / latitude_to_miles())
    lower_longitude = longitude - (distance_miles / longitude_to_miles(longitude))
    upper_longitude = longitude + (distance_miles / longitude_to_miles(longitude))

    max_time_subquery = DBSession.query(
        StampCollection.stamp_id.label('stamp_id'),
        func.max(StampCollection.date_collected).label('most_recent_time'),
    ).join(
        Stamp
    ).group_by(
        StampCollection.stamp_id
    ).subquery()

    stamps = DBSession.query(
        StampLocation,
        Stamp,
        Park,
        max_time_subquery.c.most_recent_time.label('last_seen'),
    ).join(
        StampToLocation
    ).join(
        Stamp
    ).join(
        Park
    ).outerjoin(
        max_time_subquery,
        max_time_subquery.c.stamp_id == Stamp.id
    ).filter(
        StampLocation.latitude.between(lower_latitude, upper_latitude)
    ).filter(
        StampLocation.longitude.between(lower_longitude, upper_longitude)
    ).order_by(
        asc(
            (latitude - StampLocation.latitude) * (latitude - StampLocation.latitude) +
            (longitude - StampLocation.longitude) * (longitude - StampLocation.longitude)
        )
    ).all()

    # We're doing square distance calculation in SQL so that we can use
    # indexes, so we might include some stamps that aren't within the radius
    # that we need to post-filter
    stamps = [
        s for s in stamps if latitude_longitude_distance(
            latitude,
            longitude,
            s.StampLocation.latitude,
            s.StampLocation.longitude
        ) < distance_miles
    ]

    return stamps

def stamp_exists(stamp_text):
    stamp_text_count = DBSession.query(
        Stamp.text
    ).filter(
        Stamp.text == stamp_text
    ).count()

    if stamp_text_count > 0:
        return True
    else:
        return False

def create_new_stamp(stamp_text, stamp_type, added_by_user_id):
    # The model allows for nulls here because my initial database population
    # script doesn't have a user, so just enforce this rule here
    if added_by_user_id is None:
        raise ValueError('added_by_user should not be None')

    stamp = Stamp(
        text=stamp_text,
        type=stamp_type,
        status='active',
        added_by_user_id=added_by_user_id,
    )
    DBSession.add(stamp)

    stamp_id = DBSession.query(
        Stamp.id
    ).filter(
        Stamp.text == stamp_text
    ).one()

    return stamp_id[0]

def edit_stamp(stamp_id, stamp_text, stamp_type, stamp_status, edited_by_user_id):
    # The model allows for nulls here because my initial database population
    # script doesn't have a user, so just enforce this rule here
    if edited_by_user_id is None:
        raise ValueError('edited_by_user should not be None')

    with manager:
        stamp = DBSession.query(
            Stamp
        ).filter(
            Stamp.id == stamp_id
        ).one()

        stamp.text = stamp_text
        stamp.type = stamp_type
        stamp.status = stamp_status

        stamp_history = StampHistory(
            stamp_id=stamp_id,
            text=stamp_text,
            type=stamp_type,
            status='active',
            edited_by_user_id=edited_by_user_id,
        )
        DBSession.add(stamp_history)

def get_stamps_beginning_with_string(text, limit=None):
    """Returns stamps with text that start with the substring."""
    if limit is None:
        limit = 10

    stamps = DBSession.query(
        Stamp
    ).filter(
        Stamp.text.like(text + '%')
    ).limit(
        limit
    ).all()

    return stamps


# -*- coding:utf-8 -*-
from vcf.parser import _Info as VcfInfo, field_counts as vcf_field_counts
from multiprocessing import Queue, Process, Pool, cpu_count
from maxentpy.maxent import load_matrix5, load_matrix3
from maxentpy.maxent_fast import score5, score3
from functools import partial, reduce
from argparse import ArgumentParser
from collections import namedtuple
from itertools import count
import pandas as pd
import pyfaidx
import queue
import vcf
import os

ROOT = os.path.abspath(os.path.dirname(__file__))
ANNOTATION = os.path.join(ROOT, 'annotation')
VariantRecord = namedtuple('VariantRecord', ['chromosome', 'pos', 'ref', 'alt'])
VariantAnnotation = namedtuple('VariantAnnotation', ['location', 'exon_nearby', 'exon_start', 'exon_end', 'strand', 'gene'])

class MaxEntScan:
    def __init__(self, annotation, reference):
        self.annotation = pd.read_csv(os.path.join(ANNOTATION, annotation + '.exon.gff.tsv'), sep='\t', header=None, dtype={0: str})
        self.matrix5 = load_matrix5()
        self.matrix3 = load_matrix3()
        self.reference = pyfaidx.Fasta(reference)

    def mes_score(self, record):
        if not (record.ref in ['A', 'T', 'C', 'G', 'a', 't', 'c', 'g'] and record.alt in ['A', 'T', 'C', 'G', 'a', 't', 'c', 'g']):
            return record.alt + '|.|.|.|.'
        record_anno = MaxEntScan.variant_annotation(self.annotation, record, 0)
        seq_ss, seq_ss_mut, ss = '', '', ''
        if record_anno.location == 'exon':
            if record_anno.strand == '+':
                if record.pos + 2 >= record_anno.exon_end:
                    seq_ss, seq_ss_mut, ss = MaxEntScan.forward_exon_ss5(self.reference, record, record_anno.exon_end)
                elif record.pos - 2 <= record_anno.exon_start:
                    seq_ss, seq_ss_mut, ss = MaxEntScan.forward_exon_ss3(self.reference, record, record_anno.exon_start)
            else:
                if record.pos - 2 <= record_anno.exon_start:
                    seq_ss, seq_ss_mut, ss = MaxEntScan.reverse_exon_ss5(self.reference, record, record_anno.exon_start)
                elif record.pos + 2 >= record_anno.exon_end:
                    seq_ss, seq_ss_mut, ss = MaxEntScan.reverse_exon_ss3(self.reference, record, record_anno.exon_end)
            return '|'.join([record.alt, record_anno.gene, *MaxEntScan.score_ss_seq(self.matrix5, self.matrix3, seq_ss, seq_ss_mut, ss)])
        else:
            record_anno = MaxEntScan.variant_annotation(self.annotation, record, -6)
            if record_anno.exon_nearby:
                if record_anno.strand == '+':
                    seq_ss, seq_ss_mut, ss = MaxEntScan.forward_intron_ss5(self.reference, record, record_anno.exon_end)
                    return '|'.join([record.alt, record_anno.gene, *MaxEntScan.score_ss_seq(self.matrix5, self.matrix3, seq_ss, seq_ss_mut, ss)])
            record_anno = MaxEntScan.variant_annotation(self.annotation, record, 20)
            if record_anno.exon_nearby:
                if record_anno.strand == '+':
                    seq_ss, seq_ss_mut, ss = MaxEntScan.forward_intron_ss3(self.reference, record, record_anno.exon_start)
                    return '|'.join([record.alt, record_anno.gene, *MaxEntScan.score_ss_seq(self.matrix5, self.matrix3, seq_ss, seq_ss_mut, ss)])
            record_anno = MaxEntScan.variant_annotation(self.annotation, record, 6)
            if record_anno.exon_nearby:
                if record_anno.strand == '-':
                    seq_ss, seq_ss_mut, ss = MaxEntScan.reverse_intron_ss5(self.reference, record, record_anno.exon_start)
                    return '|'.join([record.alt, record_anno.gene, *MaxEntScan.score_ss_seq(self.matrix5, self.matrix3, seq_ss, seq_ss_mut, ss)])
            record_anno = MaxEntScan.variant_annotation(self.annotation, record, -20)
            if record_anno.exon_nearby:
                if record_anno.strand == '-':
                    seq_ss, seq_ss_mut, ss = MaxEntScan.reverse_intron_ss3(self.reference, record, record_anno.exon_end)
                    return '|'.join([record.alt, record_anno.gene, *MaxEntScan.score_ss_seq(self.matrix5, self.matrix3, seq_ss, seq_ss_mut, ss)])
            return '|'.join([record.alt, '.', *MaxEntScan.score_ss_seq(self.matrix5, self.matrix3, seq_ss, seq_ss_mut, ss)])

    @staticmethod
    def variant_annotation(annotation, record, offset):
        exon = annotation[(annotation[0] == record.chromosome) & (annotation[1] <= record.pos + offset) & (record.pos + offset <= annotation[2])]
        if exon.empty:
            location = 'intron'
            exon_nearby = False
            exon_start = None
            exon_end = None
            strand = None
            gene = None
        else:
            location = 'exon'
            exon_nearby = True
            exon_start = exon[1].values[0]
            exon_end = exon[2].values[0]
            strand = exon[3].values[0]
            gene = exon[4].values[0]
        return VariantAnnotation(location, exon_nearby, exon_start, exon_end, strand, gene)

    @staticmethod
    def forward_exon_ss5(reference, record, exon_end):
        seq_ss5 = str(reference.get_seq(record.chromosome, exon_end - 2, exon_end + 6))
        seq_ss5 = seq_ss5[0:3].lower() + seq_ss5[3:5] + seq_ss5[5:].lower()
        seq_ss5_mut = seq_ss5[0:record.pos - exon_end + 2] + record.alt + seq_ss5[record.pos - exon_end + 3:]
        return seq_ss5, seq_ss5_mut, 'ss5'

    @staticmethod
    def reverse_exon_ss5(reference, record, exon_start):
        seq_ss5 = str(reference.get_seq(record.chromosome, exon_start - 6, exon_start + 2).reverse.complement)
        seq_ss5 = seq_ss5[0:3].lower() + seq_ss5[3:5] + seq_ss5[5:].lower()
        seq_ss5_mut = seq_ss5[0:exon_start - record.pos + 2] + MaxEntScan.base_complement(record.alt) + seq_ss5[exon_start - record.pos + 3:]
        return seq_ss5, seq_ss5_mut, 'ss5'

    @staticmethod
    def forward_exon_ss3(reference, record, exon_start):
        seq_ss3 = str(reference.get_seq(record.chromosome, exon_start - 20, exon_start + 2))
        seq_ss3 = seq_ss3[0:18].lower() + seq_ss3[18:20] + seq_ss3[20:].lower()
        seq_ss3_mut = seq_ss3[0:20] + seq_ss3[20:record.pos - exon_start + 20] + record.alt + seq_ss3[record.pos - exon_start + 21:]
        return seq_ss3, seq_ss3_mut, 'ss3'

    @staticmethod
    def reverse_exon_ss3(reference, record, exon_end):
        seq_ss3 = str(reference.get_seq(record.chromosome, exon_end - 2, exon_end + 20).reverse.complement)
        seq_ss3 = seq_ss3[0:18].lower() + seq_ss3[18:20] + seq_ss3[20:].lower()
        seq_ss3_mut = seq_ss3[0:20] + seq_ss3[20:exon_end - record.pos + 20] + MaxEntScan.base_complement(record.alt) + seq_ss3[exon_end - record.pos + 21:]
        return seq_ss3, seq_ss3_mut, 'ss3'

    @staticmethod
    def forward_intron_ss5(reference, record, exon_end):
        seq_ss5 = str(reference.get_seq(record.chromosome, exon_end - 2, exon_end + 6))
        seq_ss5 = seq_ss5[0:3].lower() + seq_ss5[3:5] + seq_ss5[5:].lower()
        seq_ss5_mut = seq_ss5[0:3] + seq_ss5[3:record.pos - exon_end + 2] + record.alt + seq_ss5[record.pos - exon_end + 3:]
        return seq_ss5, seq_ss5_mut, 'ss5'

    @staticmethod
    def reverse_intron_ss5(reference, record, exon_start):
        seq_ss5 = str(reference.get_seq(record.chromosome, exon_start - 6, exon_start + 2).reverse.complement)
        seq_ss5 = seq_ss5[0:3].lower() + seq_ss5[3:5] + seq_ss5[5:].lower()
        seq_ss5_mut = seq_ss5[0:3] + seq_ss5[3:exon_start - record.pos + 2] + MaxEntScan.base_complement(record.alt) + seq_ss5[exon_start - record.pos + 3:]
        return seq_ss5, seq_ss5_mut, 'ss5'

    @staticmethod
    def forward_intron_ss3(reference, record, exon_start):
        seq_ss3 = str(reference.get_seq(record.chromosome, exon_start - 20, exon_start + 2))
        seq_ss3 = seq_ss3[0:18].lower() + seq_ss3[18:20] + seq_ss3[20:].lower()
        seq_ss3_mut = seq_ss3[0:record.pos - exon_start + 20] + record.alt + seq_ss3[record.pos - exon_start + 21:]
        return seq_ss3, seq_ss3_mut, 'ss3'

    @staticmethod
    def reverse_intron_ss3(reference, record, exon_end):
        seq_ss3 = str(reference.get_seq(record.chromosome, exon_end - 2, exon_end + 20).reverse.complement)
        seq_ss3 = seq_ss3[0:18].lower() + seq_ss3[18:20] + seq_ss3[20:].lower()
        seq_ss3_mut = seq_ss3[0:exon_end - record.pos + 20] + MaxEntScan.base_complement(record.alt) + seq_ss3[exon_end - record.pos + 21:]
        return seq_ss3, seq_ss3_mut, 'ss3'

    @staticmethod
    def score_ss_seq(matrix5, matrix3, seq_ss, seq_ss_mut, ss):
        if 'N' in seq_ss or 'N' in seq_ss_mut or 'n' in seq_ss or 'n' in seq_ss_mut:
            return '.', '.', '.'
        if ss == 'ss5':
            score_wt = score5(seq_ss, matrix=matrix5)
            score_mut = score5(seq_ss_mut, matrix=matrix5)
            return str(round(score_wt, 3)), str(round(score_mut, 3)), str(round((score_mut - score_wt) / score_wt, 3))
        elif ss == 'ss3':
            score_wt = score3(seq_ss, matrix=matrix3)
            score_mut = score3(seq_ss_mut, matrix=matrix3)
            return str(round(score_wt, 3)), str(round(score_mut, 3)), str(round((score_mut - score_wt) / score_wt, 3))
        else:
            return '.', '.', '.'

    @staticmethod
    def base_complement(dna_base):
        complement_rule = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C', 'a': 'T', 't': 'A', 'c': 'G', 'g': 'C'}
        return complement_rule[dna_base]

def split_df(df, split_num):
    df.reset_index(drop=True, inplace=True)
    df_list = list()
    step = round(df.shape[0]/split_num)
    for i in range(split_num):
        if i == 0:
            df_list.append(df.loc[0: step-1])
        elif i == split_num-1:
            df_list.append(df.loc[step*i:])
        else:
            df_list.append(df.loc[step*i:step*(i+1)-1])
    return df_list

def score_pseudo_vcf(annotation, reference, format_in, df):
    mes = MaxEntScan(annotation, reference)
    if format_in == 'vcf-4cols':
        df['MaxEntScan'] = df.apply(lambda x: mes.mes_score(VariantRecord(x['#CHROM'], x['POS'], x['REF'], x['ALT'])), axis=1)
    else:
        df['MaxEntScan'] = df.apply(lambda x: mes.mes_score(VariantRecord(x['#Chr'], x['Stop'], x['Ref'], x['Call'])), axis=1)
    return df

def annotation_pseudo_vcf(parsed_args, process_num):
    if parsed_args.format_in == 'vcf-4cols':
        df = pd.read_csv(parsed_args.file_in, sep='\t', dtype={'#CHROM': str})
    else:
        df = pd.read_csv(parsed_args.file_in, sep='\t', dtype={'#Chr': str})
    df_list = split_df(df, process_num)
    score_partial = partial(score_pseudo_vcf, parsed_args.annotation, parsed_args.reference, parsed_args.format_in)
    with Pool(process_num) as pool:
        df_list_score = pool.map(score_partial, df_list)
    df_score = reduce(lambda x, y: x.append(y), df_list_score)
    df_score.to_csv(parsed_args.file_out, sep='\t', index=False)

def score_vcf(records, results, annotation, reference):
    mes = MaxEntScan(annotation, reference)
    while True:
        try:
            record = records.get(False)
        except queue.Empty:
            continue
        if record != 'END':
            record_id, record_infos, record_score_list = record[0], record[1], list()
            for record_info in record_infos:
                record_score_list.append(mes.mes_score(record_info))
            results.put((record_id, ','.join(record_score_list)))
        else:
            records.put('END')
            break

def annotation_vcf(parsed_args, process_num):
    records, results = Queue(100 * process_num), Queue()
    input_finished = False
    output_finished = False
    wait_records = dict()
    processes = list()
    records_id = count()
    for i in range(process_num):
        p = Process(target=score_vcf, args=(records, results, parsed_args.annotation, parsed_args.reference))
        processes.append(p)
        p.start()
    vcf_reader = vcf.Reader(filename=parsed_args.file_in)
    vcf_reader.infos['MaxEntScan'] = VcfInfo('MaxEntScan', vcf_field_counts['A'], 'String',
                                             'MaxEntScan Score for VCF record alleles, related_score = (mut_score-wild_score)/wild_score Format: ALLELE|SYMBOL|wild_score|mut_score|related_score',
                                             version=None, source=None)
    vcf_writer = vcf.Writer(open(parsed_args.file_out, 'w'), vcf_reader)
    while True:
        while not records.full() and not input_finished:
            try:
                record = next(vcf_reader)
                record_id = next(records_id)
                wait_records[record_id] = record
                record_infos = list()
                chromosome = str(record.CHROM)
                pos = record.POS
                ref = record.REF
                for alt in record.ALT:
                    record_infos.append(VariantRecord(chromosome, pos, ref, str(alt)))
                records.put((record_id, record_infos))
            except StopIteration:
                input_finished = True
                records.put('END')
                break
        processes_status = list()
        for p in processes:
            processes_status.append(p.is_alive())
        if True not in processes_status:
            results.put('END')
        while True:
            try:
                result = results.get(False)
            except queue.Empty:
                break
            if result != 'END':
                record_id, record_score = result[0], result[1]
                record_write = wait_records.pop(record_id)
                record_write.add_info('MaxEntScan', record_score)
                vcf_writer.write_record(record_write)
            else:
                output_finished = True
                break
        if output_finished:
            break
    vcf_writer.close()

def main():
    parser = ArgumentParser()
    parser.add_argument('-r', help='Reference genome fasta file', required=True, dest='reference')
    parser.add_argument('-a', help='Genome annotation: GRCh37 or GRCh38', default='GRCh37', dest='annotation')
    parser.add_argument('-i', help='Input file', required=True, dest='file_in')
    parser.add_argument('-o', help='Output file', required=True, dest='file_out')
    parser.add_argument('--format_in', help='Input file format: vcf,vcf-4cols,bgianno', default='vcf-4cols', dest='format_in')
    parser.add_argument('-p', help='Process number', type=int, default=1, dest='processes')
    parsed_args = parser.parse_args()
    if parsed_args.annotation not in ['GRCh37', 'GRCh38']:
        raise Exception('Genome not recognized! Must be GRCh37 or GRCh38')
    if parsed_args.format_in not in ['vcf', 'vcf-4cols', 'bgianno']:
        raise Exception('Input file format not recognized! Must be vcf, vcf-4cols or bgianno')
    process_num = min(cpu_count(), parsed_args.processes)
    if parsed_args.format_in != 'vcf':
        annotation_pseudo_vcf(parsed_args, process_num)
    else:
        annotation_vcf(parsed_args, process_num)

if __name__ == '__main__':
    main()


import numpy as np
import pandas as pd

from utils import info, info_x, split_info_x

class C45Tree:
    leaf: bool
    partition_field: str
    leaf_class: np.ndarray
    subtrees: dict = {}
    most_popular_class: np.int64

    def __init__(self, data: pd.DataFrame) -> None:
        subset_fields = data.drop(columns='GRADE').columns

        subset_fields_gain_ratios = pd.Series(
            [gain_ratio(data, subset_field) * int(len(data[subset_field].unique()) != 1)
             for subset_field in subset_fields]).fillna(0)

        self.most_popular_class = data.GRADE.mode()[0]

        self.subtrees = {}

        self.leaf_class = data.GRADE.unique()

        if not any(subset_fields_gain_ratios) or len(data.GRADE.unique()) == 1:
            self.leaf = True
            return

        self.partition_field = subset_fields[np.argmax(
            subset_fields_gain_ratios)]
        self.leaf = False
        for partition_field_value in data[self.partition_field].unique():
            self.subtrees[partition_field_value] = C45Tree(
                data[data[self.partition_field] == partition_field_value].drop(columns=self.partition_field))

    def traverse(self, data_item: pd.DataFrame) -> np.int64:
        if not self.leaf:
            partition_value = data_item[self.partition_field]
            if partition_value in self.subtrees:
                return self.subtrees[partition_value].traverse(data_item)
            return self.most_popular_class
        final_class = self.leaf_class
        if len(final_class) == 1:
            return final_class[0]
        return self.most_popular_class

def gain_ratio(data: pd.DataFrame, partition_field: pd.Index):
    return (info(data) - info_x(data, partition_field)) / split_info_x(data, partition_field)


from collections import Counter

def analyze_feedback(feedback_list):
    if not feedback_list:
        return {
            'Most Frequent Positive Words': [],
            'Most Frequent Negative Words': [],
            'Overall Sentiment Score': 0,
            'Sentiment Report': 'Positive to Negative Ratio = N/A'
        }

    positive_words = set(['good', 'excellent', 'recommended', 'quick', 'highly'])
    negative_words = set(['terrible', 'bad', 'slow', 'poor'])

    positive_counts = Counter()
    negative_counts = Counter()

    for feedback in feedback_list:
        words = [word.lower().strip('.,?!') for word in feedback.split()]
        positive_counts.update([word for word in words if word in positive_words])
        negative_counts.update([word for word in words if word in negative_words])

    most_frequent_positive_words = [word for word, count in positive_counts.most_common()]
    most_frequent_negative_words = [word for word, count in negative_counts.most_common()]

    overall_sentiment_score = sum(positive_counts[word] for word in positive_words) - sum(negative_counts[word] for word in negative_words)

    total_words = sum(len(feedback.split()) for feedback in feedback_list)

    if total_words == 0:
        return {
            'Most Frequent Positive Words': most_frequent_positive_words,
            'Most Frequent Negative Words': most_frequent_negative_words,
            'Overall Sentiment Score': 0,
            'Sentiment Report': 'Positive to Negative Ratio = N/A'
        }

    positive_count = sum(positive_counts[word] for word in positive_words)
    negative_count = sum(negative_counts[word] for word in negative_words)
    positive_to_negative_ratio = positive_count / max(1, negative_count)

    return {
        'Most Frequent Positive Words': most_frequent_positive_words,
        'Most Frequent Negative Words': most_frequent_negative_words,
        'Overall Sentiment Score': overall_sentiment_score,
        'Sentiment Report': f'Positive to Negative Ratio = {positive_to_negative_ratio:.2f}:1'
    }

# Caso 1
feedback_list_1 = [
    "The service was excellent! Highly recommended.",
    "I had a terrible experience with the product quality.",
    "Good customer support and quick resolution.",
    "Very bad customer support and slow attention.",
    "Poor packaging and delivery service."
]

result_1 = analyze_feedback(feedback_list_1)
print("Caso 1:")
print(result_1)
print("\n---\n")

# Caso 2
feedback_list_2 = [
    "This is a good product.",
    "Excellent service!",
    "Satisfied with my purchase."
]

result_2 = analyze_feedback(feedback_list_2)
print("Caso 2:")
print(result_2)
print("\n---\n")

# Caso 3
feedback_list_3 = [
    "The service was bad.",
    "Poor customer experience.",
    "Terrible quality product."
]

result_3 = analyze_feedback(feedback_list_3)
print("Caso 3:")
print(result_3)
print("\n---\n")


import numpy as np
from sklearn import metrics
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV

from sklearn import svm

data = pd.read_csv("cardio_train_clean_featureselection.csv")

# Features = xfeat, Target Variable = Y
xfeat = data[['Age', 'Gender', 'Height', 'Weight', "Systolic BP",
              "Chlosterol ", "Glucose", "Smoke", "Alcohol", "Active"]]

y = data['Cardio']

x_train, x_test, y_train, y_test = train_test_split(xfeat, y, test_size=0.2, random_state=0)

# SVM

svm_model = svm.SVC()
svm_model.fit(x_train, y_train)
y_pred_svm = svm_model.predict(x_test)

confusion_matrix_svm = pd.crosstab(y_test, y_pred_svm, rownames=['Actual'], colnames=['Predicted'])
sns.heatmap(confusion_matrix_svm, annot=True)

print('SVM Metrics:')
print('Accuracy: ', metrics.accuracy_score(y_test, y_pred_svm))
print('Precision: ', metrics.precision_score(y_test, y_pred_svm))
print('Recall: ', metrics.recall_score(y_test, y_pred_svm))
print()
plt.show()

# Tuning
'''
parameters = {'kernel': ('linear', 'rbf'), 'C': [1, 10]}

search = GridSearchCV(svm_model, parameters, scoring = 'recall')

results = search.fit(x_test, y_test)

print('Mean Recall: ', results.best_score_)
print('Config: ', results.best_params_)

'''
svm_model = svm.SVC(kernel='linear', C=10)
svm_model.fit(x_train, y_train)
y_pred_svm = svm_model.predict(x_test)

confusion_matrix_svm = pd.crosstab(y_test, y_pred_svm, rownames=['Actual'], colnames=['Predicted'])
sns.heatmap(confusion_matrix_svm, annot=True)

print('Tuned SVM Metrics:')
print('Accuracy: ', metrics.accuracy_score(y_test, y_pred_svm))
print('Precision: ', metrics.precision_score(y_test, y_pred_svm))
print('Recall: ', metrics.recall_score(y_test, y_pred_svm))
print()
plt.show()

Data Exploration, Training, Evaluation/SVM_Tuned.py


#!/usr/bin/python

# sim_index.py
import numpy as np
import wx
import wx.lib.scrolledpanel as scrolled

class Sim_Index(scrolled.ScrolledPanel):

    def __init__(self, parent, simulation, color):
        scrolled.ScrolledPanel.__init__(self, parent, id = -1,  style = wx.TAB_TRAVERSAL|wx.SUNKEN_BORDER)

        sim_color = simulation.sim_color if simulation.sim_color else color

        self.sdrdf = simulation.sdrdf#[2.5,1.5,1,.8,2.5,1.5,1,.8,2.5,1.5,1,.8,2.5,1.5,1,.8]
        self.sdrdflen = len(self.sdrdf)

        self.srdf = simulation.srdf

        #fcc_sim = [111,122,220,400,111,122,220,400,111,122,220,400,111,122,220,400]
        self.peak_index_labels = simulation.peak_index_labels

        wx.StaticText(self, -1, 'Legend Label: ', (20, 5))

        self.sim_label_tc = wx.TextCtrl(self, -1, '',  (130, 0), (130, -1))
        sim_label_text = simulation.sim_label
        self.sim_label_tc.SetValue(sim_label_text)

        wx.StaticText(self, -1, 'Simulation Intensity: ', (20, 35))

        self.sim_intens_tc = wx.TextCtrl(self, -1, '',  (200, 30), (60, -1))
        sim_intens_text = str(simulation.sim_intens)
        self.sim_intens_tc.SetValue(sim_intens_text)

        wx.StaticText(self, -1, 'Simulation Color: ', (20, 65))

        self.colorpicker = wx.ColourPickerCtrl(self, -1, colour=sim_color, pos=(200, 60),
                 size=(60, 30), name='Simulation Color')
        #self.sim_color_tc = wx.TextCtrl(self, -1, '',  (150, 60), (90, -1))
        #self.colorDlgBtn = wx.Button(self,  -1, "...",  (240, 60), (20, -1))
        #self.colorDlgBtn.Bind(wx.EVT_BUTTON, self.onColorDlg)
        #self.sim_color_tc.SetValue('#ffffff')

        vbox = wx.BoxSizer(wx.VERTICAL)

        offset = 4
        sp = 30
        i = offset
        for peak, intens in zip(self.sdrdf,self.srdf):
            static = wx.StaticText(self, -1, str(i-offset+1) + '-', (20, sp*i+5))
            static1 = wx.StaticText(self, -1, 'd:', (60, sp*i+5))
            dspace = 1/peak
            wx.StaticText(self, -1, str('%.2f' % dspace) + u' \u00c5', (75, sp*i+5))

            static2 = wx.StaticText(self, -1,'i:', (140, sp*i+5))
            wx.StaticText(self, -1, str('%.2f' % intens), (153, sp*i+5))
            i += 1

        vbox.Add((250,sp*i+5), wx.EXPAND, 0)

        if len(self.peak_index_labels)==0:
            i = offset
            self.peak_index_tc=[]
            for peak_index_label in self.sdrdf:
                self.peak_index_tc += [wx.TextCtrl(self, -1, '',  (200, sp*i), (60, -1))]
                #self.peak_index_tc[i-1].SetValue(str(au_radius))
                i += 1
        else:
            i = offset
            self.peak_index_tc=[]
            for peak_index_label in self.peak_index_labels[:self.sdrdflen]:
                self.peak_index_tc += [wx.TextCtrl(self, -1, '',  (200, sp*i), (60, -1))]
                self.peak_index_tc[i-offset].SetValue(str(peak_index_label))
                i += 1

        self.SetSizer(vbox)
        self.SetupScrolling()

#    def onColorDlg(self, event):
#        """
#        This is mostly from the wxPython Demo!
#        """
#        dlg = wx.ColourDialog(self)
#
#        # Ensure the full colour dialog is displayed,
#        # not the abbreviated version.
#        dlg.GetColourData().SetChooseFull(True)
#
#        if dlg.ShowModal() == wx.ID_OK:
#            data = dlg.GetColourData()
#            self.sim_color_tc.SetValue('#%02x%02x%02x' % data.GetColour().Get())
#            print ('#%02x%02x%02x' % data.GetColour().Get())
#
#        dlg.Destroy()

class Index(wx.Dialog):
    def __init__(self, parent, id, title):

        window_height = 500
        wx.Dialog.__init__(self, parent, id, title, wx.DefaultPosition, wx.Size(350, window_height))

        wx.StaticText(self, -1, 'Enter or edit the peak indices and title\nfor the imported simulations', (20,20))

        nb = wx.Notebook(self, -1, pos = (25, 70), size = (300, 350), style=wx.NB_TOP)

        self.parent = parent
        self.simulations = parent.simulations

        if len(self.simulations) >= 3:
            sim_len_i = 3
        else:
            sim_len_i = len(self.simulations)
        print(sim_len_i)#,self.srdfb[sim_len_i[0]],self.sdrdfb[sim_len_i[0]])

        color = ['#42D151','#2AA298','#E7E73C']

        for i, simulation in enumerate(self.simulations[-sim_len_i:]):
            sim = Sim_Index(nb,simulation,color[i])
            simulation.sim_tcs = sim
            nb.AddPage(sim, simulation.sim_name)

        set_btn = wx.Button(self, 1, 'Edit Labels', (70, window_height-55))
        set_btn.SetFocus()
        clear_btn = wx.Button(self, 2, 'Cancel', (185, window_height-55))

        self.Bind(wx.EVT_BUTTON, self.OnIndex, id=1)
        self.Bind(wx.EVT_BUTTON, self.OnClose, id=2)
        self.Bind(wx.EVT_CLOSE, self.OnClose)

    def OnIndex(self, event):
        if len(self.simulations) >= 3:
            sim_len_i = 3
        else:
            sim_len_i = len(self.simulations)
        print(sim_len_i)#,self.srdfb[sim_len_i[0]],self.sdrdfb[sim_len_i[0]])
        for simulation in self.simulations[-sim_len_i:]:
            peak_index_labels = []

            for peak_index_tc in simulation.sim_tcs.peak_index_tc:
                peak_index_labels += [peak_index_tc.GetValue()]
            print(simulation.sim_tcs.colorpicker.GetColour(), simulation.sim_tcs.colorpicker.GetColour()[0:3])
            simulation.edit_index_labels(peak_index_labels, simulation.sim_tcs.sim_label_tc.GetValue(), simulation.sim_tcs.sim_intens_tc.GetValue(), simulation.sim_tcs.colorpicker.GetColour())
        self.Destroy()

        self.parent.axes.cla()
        self.parent.plot(2,'b')

        self.parent.axes.figure.canvas.draw()

    def OnClose(self, event):
        self.Destroy()

class Simulation:
    """
    Simulation label and index
    """
    def __init__(self, sim_name, sim_open, peak_index_labels=[]):
        if len(sim_open.shape) == 1: # it's a records array
            sdrdf = sim_open['inv_d']
            srdf = sim_open['intensity']
            self.sim_intens = 1
        else:
            theta_2 = sim_open[:,0]
            inv_d = (2*np.sin(((theta_2/180)*np.pi)/2.))/.5

            intensity = sim_open[:,1]
            intensity /= intensity.max()

            print(sim_open.shape, len(theta_2), len(intensity))

            srdf = intensity
            sdrdf = inv_d
            self.sim_intens = 1

            sim_index = np.nonzero(srdf!=0)

            sdrdf = sdrdf[sim_index]
            srdf = srdf[sim_index]

        self.sdrdf = sdrdf
        self.srdf = srdf

        self.sim_name = sim_name
        self.sim_label = sim_name

        self.peak_index_labels = list(peak_index_labels)
        self.sim_tcs = []
        self.sim_color = 0

        #print(1/array(sim_scatter))

        #print(imgcal)

        #filen=(simname + '.txt')
        #data = loadtxt(filen,skiprows=0)
        ##print(data)

        #srdf = data[:,1]
        #sdrdf = data[:,0]

        #start =20
        #dsdrdf = linspace(0,sdrdf[-1],boxs)

        print(self.sdrdf, self.srdf)

    def edit_index_labels(self, index_labels, sim_label, sim_intens, sim_color):
        self.peak_index_labels = index_labels
        self.sim_label = sim_label
        self.sim_intens = float(sim_intens)
        self.sim_color = '#%02x%02x%02x' % sim_color[0:3]

#app = wx.App()
#dlg = Index(None, -1, 'Index Peaks')
#dlg.Show(True)
#app.MainLoop()


""" Test Album class """

import os
import pathlib
import tempfile

import pytest
from applescript import AppleScript

import photoscript
from tests.conftest import get_os_version, photoslib, suspend_capture
from tests.photoscript_config_data import (
    ALBUM_1_NAME,
    ALBUM_1_PATH_STR,
    ALBUM_1_PATH_STR_COLON,
    ALBUM_1_PHOTO_EXPORT_FILENAMES,
    ALBUM_1_PHOTO_UUIDS,
    ALBUM_1_POST_REMOVE_UUIDS,
    ALBUM_1_REMOVE_UUIDS,
    ALBUM_1_UUID,
    ALBUM_1_UUID_OSXPHOTOS,
    ALBUM_NAMES_ALL,
    ALBUM_NAMES_TOP,
    EMPTY_ALBUM_NAME,
    FOLDER_NAME,
    FOLDER_NAMES_ALL,
    FOLDER_NAMES_TOP,
    FOLDER_UUID,
    IMPORT_PATHS,
    IMPORT_PHOTOS,
    NUM_PHOTOS,
    PHOTO_ADD_UUID,
    PHOTO_FAVORITES_SET_UUID,
    PHOTO_FAVORITES_UNSET_UUID,
    PHOTOS_FAVORITES,
    PHOTOS_FAVORITES_SET,
    PHOTOS_FILENAMES,
    PHOTOS_PLANTS,
    PHOTOS_UUID,
    PHOTOS_UUID_FILENAMES,
    SELECTION_UUIDS,
)
from tests.utils import stemset

EMPTY_ALBUM_LEN = 0

########## Interactive tests run first ##########

########## Non-interactive tests ##########

def test_album_init():

    album = photoscript.Album(ALBUM_1_UUID)
    assert isinstance(album, photoscript.Album)

def test_album_init_osxphotos_uuid():
    """test Album() with osxphotos style UUID"""

    album = photoscript.Album(ALBUM_1_UUID_OSXPHOTOS)
    assert isinstance(album, photoscript.Album)
    assert album.id == ALBUM_1_UUID
    assert album.uuid == ALBUM_1_UUID_OSXPHOTOS

def test_album_init_bad_uuid():

    with pytest.raises(ValueError):
        assert photoscript.Album("BAD_UUID")

def test_album_id():

    album = photoscript.Album(ALBUM_1_UUID)
    assert album.id == ALBUM_1_UUID

def test_album_uuid():

    album = photoscript.Album(ALBUM_1_UUID)
    assert album.uuid == ALBUM_1_UUID_OSXPHOTOS

def test_len(photoslib: photoscript.PhotosLibrary):
    """test Album.__len__"""
    album = photoslib.album(ALBUM_1_NAME)
    assert len(album) == len(ALBUM_1_PHOTO_UUIDS)

def test_album_name_title(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_1_NAME)
    assert album.name == ALBUM_1_NAME
    assert album.title == ALBUM_1_NAME

def test_album_name_setter(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_1_NAME)
    new_name = "My New Album Name"
    album.name = new_name
    assert album.name == new_name
    assert album.title == new_name

    # reset name
    album.name = ALBUM_1_NAME

def test_album_title_setter(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_1_NAME)
    new_name = "My New Album Name"
    album.title = new_name
    assert album.name == new_name
    assert album.title == new_name

    # reset title
    album.title = ALBUM_1_NAME

def test_album_parent_id(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_1_NAME)
    assert album.parent.id == FOLDER_UUID

def test_album_parent(photoslib: photoscript.PhotosLibrary):

    album = photoslib.album(ALBUM_1_NAME)
    parent = album.parent
    assert isinstance(parent, photoscript.Folder)
    assert parent.name == FOLDER_NAME

def test_album_parent_top_level(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_NAMES_TOP[0])
    assert album.parent is None

def test_album_path_str(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_1_NAME)
    assert album.path_str() == ALBUM_1_PATH_STR
    assert album.path_str(":") == ALBUM_1_PATH_STR_COLON
    with pytest.raises(ValueError):
        assert album.path_str("FOO")

def test_album_photos(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(ALBUM_1_NAME)
    photo_ids = [p.id for p in album.photos()]
    assert sorted(photo_ids) == sorted(ALBUM_1_PHOTO_UUIDS)

def test_album_photos_empty(photoslib: photoscript.PhotosLibrary):
    album = photoslib.album(EMPTY_ALBUM_NAME)
    assert album.photos() == []

def test_album_photos_add(photoslib: photoscript.PhotosLibrary):

    global EMPTY_ALBUM_LEN
    album = photoslib.album(EMPTY_ALBUM_NAME)
    assert len(album) == 0
    added = album.add([photoscript.Photo(PHOTO_ADD_UUID)])
    assert len(album) == 1
    assert len(added) == 1
    assert added[0].id == PHOTO_ADD_UUID
    assert album.photos()[0].id == PHOTO_ADD_UUID

    EMPTY_ALBUM_LEN = len(album)

def test_album_photos_add_none(photoslib: photoscript.PhotosLibrary):
    """Test add with no photos"""

    album = photoslib.album(ALBUM_1_NAME)
    album_length = len(album)
    album.add([])
    assert len(album) == album_length

def test_album_import_photos(photoslib: photoscript.PhotosLibrary):
    """Test import photo"""

    global EMPTY_ALBUM_LEN
    cwd = os.getcwd()
    photo_paths = [str(pathlib.Path(cwd) / path) for path in IMPORT_PATHS]
    album = photoslib.album(EMPTY_ALBUM_NAME)
    imported = album.import_photos(photo_paths)
    assert len(imported) == len(IMPORT_PATHS)
    assert len(album) == EMPTY_ALBUM_LEN + len(IMPORT_PATHS)

    EMPTY_ALBUM_LEN = len(album)

def test_album_import_photos_skip_dup_check(photoslib: photoscript.PhotosLibrary):
    """Attempt to import a duplicate photo with skip_duplicate_check = True
    This will cause a duplicate photo to be imported"""

    # this test must be run after the previous test
    global EMPTY_ALBUM_LEN

    album = photoslib.album(EMPTY_ALBUM_NAME)
    cwd = os.getcwd()
    photo_paths = [str(pathlib.Path(cwd) / path) for path in IMPORT_PATHS]
    imported = album.import_photos(photo_paths, skip_duplicate_check=True)
    assert len(imported) == len(IMPORT_PATHS)
    assert len(album) == EMPTY_ALBUM_LEN + len(IMPORT_PATHS)

    EMPTY_ALBUM_LEN = len(album)

def test_album_export_photos_basic(photoslib: photoscript.PhotosLibrary):

    tmpdir = tempfile.TemporaryDirectory(prefix="photoscript_test_")
    album = photoslib.album(ALBUM_1_NAME)
    exported = album.export(tmpdir.name)
    exported = [pathlib.Path(filename).name for filename in exported]
    assert sorted(exported) == ALBUM_1_PHOTO_EXPORT_FILENAMES
    files = os.listdir(tmpdir.name)
    assert sorted(files) == sorted(ALBUM_1_PHOTO_EXPORT_FILENAMES)

def test_album_export_photos_duplicate_overwrite(photoslib: photoscript.PhotosLibrary):
    """Test calling export twice resulting in
    duplicate filenames but use overwrite = true"""

    tmpdir = tempfile.TemporaryDirectory(prefix="photoscript_test_")
    album = photoslib.album(ALBUM_1_NAME)
    exported1 = album.export(tmpdir.name, overwrite=True, original=True)
    exported2 = album.export(tmpdir.name, overwrite=True, original=True)
    exported = list({pathlib.Path(filename).name for filename in exported1 + exported2})

    # some versions of macOS use .JPG and some use .jpeg
    assert stemset(exported) == stemset(ALBUM_1_PHOTO_EXPORT_FILENAMES)
    files = os.listdir(tmpdir.name)
    assert stemset(files) == stemset(ALBUM_1_PHOTO_EXPORT_FILENAMES)

def test_album_remove_by_id(photoslib: photoscript.PhotosLibrary):
    """Test album.remove_by_id"""

    album = photoslib.album(ALBUM_1_NAME)
    parent_id = album.parent.id
    new_album = album.remove_by_id(ALBUM_1_REMOVE_UUIDS)
    assert isinstance(new_album, photoscript.Album)
    uuids = [photo.id for photo in new_album.photos()]
    assert sorted(uuids) == sorted(ALBUM_1_POST_REMOVE_UUIDS)
    assert new_album.title == ALBUM_1_NAME
    assert new_album.parent.id == parent_id
    assert album.id == new_album.id
    assert album.uuid == new_album.uuid

def test_album_remove(photoslib: photoscript.PhotosLibrary):
    """Test album.remove"""

    album = photoslib.album(ALBUM_1_NAME)
    photos = [photoscript.Photo(uuid) for uuid in ALBUM_1_REMOVE_UUIDS]
    new_album = album.remove(photos)
    assert isinstance(new_album, photoscript.Album)
    uuids = [photo.id for photo in new_album.photos()]
    assert sorted(uuids) == sorted(ALBUM_1_POST_REMOVE_UUIDS)
    assert new_album.title == ALBUM_1_NAME
    assert new_album.id == album.id


#!/usr/bin/python3
'''

Adds a new city

'''
from relationship_city import City
from relationship_state import Base, State
import sys
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

if __name__ == "__main__":
    argv = sys.argv
    engine = create_engine(f"mysql+mysqldb://{argv[1]}:{argv[2]}" +
                           f"@localhost:3306/{argv[3]}", pool_pre_ping=True)
    Base.metadata.create_all(engine)
    Session = sessionmaker(bind=engine)
    session = Session()
    new_state = State(name="California")
    new_city = City(name="San Francisco")
    new_state.cities.append(new_city)
    session.add(new_state)
    session.commit()
    session.close()


import matplotlib.pyplot as plt
import numpy as np

arr1 = np.loadtxt("../experiment3_cleaned/Pretrained_aug/augVGGpreacc.csv",delimiter = ',')
arr2 = np.loadtxt("../experiment3_cleaned/Pretrained_aug/augRESNETpreacc.csv",delimiter = ',')
arr3 = np.loadtxt("../experiment3_cleaned/Pretrained_aug/augATNNpreacc.csv",delimiter = ',')
arr11 = np.loadtxt("../experiment2_cleaned_v2/Pretrained_aug/augVGGpreacc.csv",delimiter = ',')
arr22 = np.loadtxt("../experiment2_cleaned_v2/Pretrained_aug/augRESNETpreacc.csv",delimiter = ',')
arr33 = np.loadtxt("../experiment2_cleaned_v2/Pretrained_aug/augATNNpreacc.csv",delimiter = ',')

fig, axs = plt.subplots(1,2, figsize = (15,5))

axs[0].plot(arr1[:,0], label = 'VGG16 - original')
axs[0].plot(arr2[:,0], label = 'Resnet18 - original')
axs[0].plot(arr3[:,0], label = 'ViT - original')
axs[0].plot(arr11[:,0], label = 'VGG16 - fisheye')
axs[0].plot(arr22[:,0], label = 'Resnet18 - fisheye')
axs[0].plot(arr33[:,0], label = 'ViT - fisheye')
#axs[0].plot(arr44[:,0], color = 'orange', linestyle = '--', label = 'ViT')
axs[0].set_xlabel("Epochs")
axs[0].set_ylabel("Accuracy (%)")
axs[0].grid()
axs[0].legend()

axs[1].plot(arr1[:,1], label = 'VGG16 - original')
axs[1].plot(arr2[:,1], label = 'Resnet18 - original')
axs[1].plot(arr3[:,1], label = 'ViT - original')
axs[1].plot(arr11[:,1], label = 'VGG16 - fisheye')
axs[1].plot(arr22[:,1], label = 'Resnet18 - fisheye')
axs[1].plot(arr33[:,1], label = 'ViT - fisheye')
axs[1].set_xlabel("Epochs")
axs[1].set_ylabel("Accuracy (%)")
axs[1].grid()
axs[1].legend()
plt.savefig("experiment3_res.png", bbox_inches = "tight")


import unittest

from selenium import webdriver
from selenium.webdriver.chrome.options import Options

from pages.login import LoginPage

class TestDjangoAdmin(unittest.TestCase):

    def setUp(self):
        options = Options()
        options.binary_location = '/Applications/Google Chrome.app/Contents/MacOS/Google Chrome'
        options.add_argument('--headless')
        options.add_argument('--disable-gpu')
        self.driver = webdriver.Chrome(chrome_options=options)
        self.driver = webdriver.Chrome()
        self.driver.implicitly_wait(5)
        self.driver.set_page_load_timeout(30)
        self.driver.set_window_size(1920, 1080)

    def test_add_user(self):
        self.driver.get('http://127.0.0.1:8000/admin')
        login_page = LoginPage(self.driver)
        admin_page = login_page.login('admin', 'p@ssword')
        add_user_page = admin_page.go_to_add_user_page()
        user_page = add_user_page.add_user('testuser', 'dummy_p@ssword')
        self.assertIn('testuser" を追加しました。続けて編集できます。
', user_page.driver.page_source)

    def tearDown(self):
        self.driver.close()


import functools
import jwt
from runSevers.config import jwt_setting, settings, default
from common.utils.message import to_dict_msg
from apps.users.models import Users

def login_require_async(view_func):
    @functools.wraps(view_func)
    async def require(self, *arg, **kwargs):
        # 验证用户是否登录
        # 获取token值
        token = self.request.headers.get('token')
        # 解析token
        paylod = jwt.decode(token, jwt_setting['key'], [jwt_setting['algorithm']], options={'verify_exp': True}, leeway=3600 * 24 * 7)
        if paylod:
        # 通过email查询数据
            uuid = paylod.get('uuid')
            # 有数据，将数据返回
            try:
                if default["mysql"]:
                    user = await self.manager.get(Users, uuid=uuid)
                    self._uuid = user.uuid
                    self._username = user.usernmae
                else:
                    self._uuid = "001"
                    self._username = "admin"
                try:
                    await view_func(self, *arg, **kwargs)
                except:
                    self.finish(to_dict_msg(500, msg='视图函数异常'))
            except:
                self.finish(to_dict_msg(401, msg='请登录后重试'))
        else:
            self.finish(to_dict_msg(401, msg='请登录后重试'))
    return require