import nltk.tokenize as tk
import nltk.stem as ns
import pandas as pd
import numpy as np
import csv
from tqdm import tqdm
from typing import *
import pickle
import random

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import KFold
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier

from sklearn.metrics import f1_score
from sklearn.metrics import accuracy_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score


def read_labeled_file(filepath: str):
    """
    :param filepath: 训练集csv文件路径
    :return:
        news：news列表，List[str]
        labels: 标签列标，List[int]
    """
    print("Now start to read training data from disk...")
    News, Labels = [], []
    with open(filepath, encoding='utf-8') as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')
        for row in tqdm(csv_reader):
            id = row[0]
            news, label = row[2], int(row[1])
            News.append(news)
            Labels.append(label)
    return News, Labels


def read_unlabeled_data(filepath: str):
    print("Now start to read testing data from disk...")
    Ids, News = [], []
    with open(filepath, encoding='utf-8') as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')
        for row in tqdm(csv_reader):
            Ids.append(row[0])
            News.append(row[2])
    return Ids, News


def preprocess(News: List[str]):
    """
    :param News: 新闻列表
    :return: 词根化处理之后的新闻列表，词根化需要分词但是后续的tf-idf-transform要一个整句子所以我分词后又拼回去了, 返回List[str]
    """
    print("Now start to stemming tokens in news...")
    sb_stemmer = ns.SnowballStemmer("english")
    stem_news = []
    for doc in tqdm(News):
        words = tk.word_tokenize(doc)
        words = [sb_stemmer.stem(each) for each in words]
        stem_news.append(' '.join(words))
    return stem_news


if __name__=='__main__':

    train_file_path = 'D:\\tmp_data_space\\xfy_data\\gossip_train_data_2020_fall.csv'
    test_file_path = 'D:\\tmp_data_space\\xfy_data\\gossip_test_data_2020_fall.csv'
    preprocessed_train_data_path = "D:\\tmp_data_space\\xfy_data\\gossip_train_X_y_Ids_Xtest.pkl"
    final_pred_test_file_path ='D:\\tmp_data_space\\xfy_data\\gossip_test_pred_data_2020_fall.csv'

    # """
    # 预处理阶段: 将训练集和测试集向量化
    # 这一阶段耗时比较长，所以将样本向量化之后我们将其写入到文件中，然后将这部分代码注释掉
    # 之后训练的时候直接从文件里加载数据就好了
    # """
    # # 准备训练数据
    # labeled_news, labels = read_labeled_file(train_file_path)
    # Ids, unlabeled_news = read_unlabeled_data(test_file_path)
    #
    # labeled_num = len(labeled_news)
    # # # sample 100:
    # # News, Labels = News[:100], Labels[:100]
    #
    # News = labeled_news
    # News.extend(unlabeled_news)
    # News = preprocess(News)
    #
    # # 转化成tf-idf向量
    # vectorizer = TfidfVectorizer(stop_words='english', max_features=5000, max_df=0.5)
    # X = vectorizer.fit_transform(News)
    # X = X.toarray()
    # X, X_test = X[:labeled_num], X[labeled_num:]
    # y = np.array(labels)
    #
    # # 将预处理好的数据存入文件
    # pickle.dump([X, y, Ids, X_test], open(preprocessed_train_data_path, 'wb'))

    """
    模型训练阶段: 做五折交叉验证
    首先从文件中加载向量化后的代码，然后选择一个模型做五折交叉验证，这里我选的是LR，跑的比较快，这个任务很简单所以效果还行
    """
    X, y, Ids, X_test = pickle.load(open(preprocessed_train_data_path, 'rb'))

    kf = KFold(n_splits=5)
    kf.get_n_splits(X)
    for train_index, valid_index in kf.split(X):
        X_train, X_valid = X[train_index], X[valid_index]
        y_train, y_valid = y[train_index], y[valid_index]

        cls = LogisticRegression()
        # cls = RandomForestClassifier(n_estimators=100)
        cls.fit(X_train, y_train)

        y_valid_pred = cls.predict(X_valid)
        f1 = f1_score(y_valid, y_valid_pred)
        acc = accuracy_score(y_valid, y_valid_pred)
        p = precision_score(y_valid, y_valid_pred)
        r = recall_score(y_valid, y_valid_pred)
        print("f1: {}, acc: {}, p: {}, r: {}".format(f1, acc, p, r))

    """
    训练验证结束后，我们观察了实验结果，然后将所有的train样本拿来做训练集，训练出一个更强的模型，为测试集生成标签
    """
    cls = LogisticRegression()
    # cls = RandomForestClassifier(n_estimators=100)
    cls.fit(X, y)
    y_test = cls.predict(X_test)

    # TODO: 将y_test按照指定方式写入文件提交即可
    doc_label = dict()
    for i in range(len(Ids)):
        doc_label[Ids[i]] = str(int(y_test[i]))

    dst_file = open(final_pred_test_file_path, 'w', encoding='utf-8', newline='')
    csv_writer = csv.writer(dst_file)
    with open(test_file_path, encoding='utf-8') as test_csv_file:
        csv_reader = csv.reader(test_csv_file, delimiter=',')
        for row in tqdm(csv_reader):
            id, news = row[0], row[2]
            label = doc_label[id]
            csv_writer.writerow([id, label, news])

    dst_file.close()

    print("Done. ")

