# encoding: utf-8
# 此文件包含模型的训练。 给定数据集，训练出情感分类模型，并把模型文件存放在 model文件夹里。 
from sklearn.feature_extraction.text import TfidfVectorizer
from nltk.tokenize import word_tokenize
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import pickle
import json
import pandas as pd
from tqdm import tqdm

class SentimentModel(object):
    def __init__(self,df_review=None,pos_star=4,neg_star=2):
        if df_review is not None:
            self.pos_review = []
            self.neg_review = []
            for idx in tqdm(range(len(df_review))):
                review = df_review.iloc[idx]
                if review.stars >= pos_star:
                    self.pos_review.append(" ".join([ word for word in word_tokenize(review.text)]))
                elif review.stars <= neg_star: 
                    self.neg_review.append(" ".join([ word for word in word_tokenize(review.text)]))
            print("样本统计:\n正例: {}\n负例: {}".format(len(self.pos_review),len(self.neg_review)))
            # print("样本示例:\n正例: {}\n负例: {}".format("\n".join(self.pos_review[:2]),"\n".join(self.neg_review[:2])))
            self.vectorizer = TfidfVectorizer()
            corpus = self.pos_review + self.neg_review
            self.X = self.vectorizer.fit_transform(corpus)
            self.y = [1] * len(self.pos_review) + [0] * len(self.neg_review)
        else:
            self.clf = None
            self.vectorizer = None

    
    def train(self,test_size=0.25,*args,**kwargs):
        if kwargs["save_model_path"] is not None:
            save_model_path = kwargs["save_model_path"]
        if kwargs["save_vectorizer_path"] is not None:
            save_vectorizer_path = kwargs["save_vectorizer_path"]
        self.clf = LogisticRegression(random_state=0,C=10)
        X_train, X_test, y_train, y_test = train_test_split(self.X, self.y, test_size=test_size, random_state=0)
        self.clf.fit(X_train,y_train)
        score = self.clf.score(X_test,y_test)
        with open(save_model_path,"wb") as f:
            pickle.dump(self.clf,f)
        with open(save_vectorizer_path,"wb") as f:
            pickle.dump(self.vectorizer,f)
        print("Finish training, score is {}.".format(str(score)))
 
    def predict(self,*args,**kwargs):
        if self.clf is None:
            if kwargs["model_path"] is not None:
                model_path = kwargs["model_path"]
                with open(model_path,'rb') as f:
                    self.clf = pickle.load(f)
            else:
                raise Exception("model_path missing")
        if self.vectorizer is None:
            if kwargs["vectorizer_path"] is not None:
                vectorizer_path = kwargs["vectorizer_path"]
                with open(vectorizer_path,'rb') as f:
                    self.vectorizer = pickle.load(f)
            else:
                raise Exception("vectorizer_path missing")
        data = args[0]
        tmp = []
        for d in data:
            tmp.append(" ".join([word for word in word_tokenize(d)]))
        data = tmp
        features = self.vectorizer.transform(data)
        result = self.clf.predict(features)
        score = self.clf.predict_proba(features)
        return result,score

if __name__ == "__main__":
    review_path = "data/review.json"
    model_path = "model/model.pkl"
    vector_path = "model/vector.pkl"
    is_train = False
    if is_train:
        reviews = []
        with open(review_path,"r",encoding="utf-8") as f:
            for line in f.readlines():
                if len(line.strip()) == 0:
                    continue
                ele = json.loads(line.strip())
                reviews.append(ele)
        review_df = pd.DataFrame(reviews)
        print(review_df.head())
        model = SentimentModel(review_df)
        model.train(save_model_path=model_path,save_vectorizer_path=vector_path)
    else:
        model = SentimentModel()
    data = ["Total bill for this horrible service? Over $8Gs. These crooks actually had the nerve to charge us $69 for 3 pills. I checked online the pills can be had for 19 cents EACH! Avoid Hospital ERs at all costs.","kkkkkk"]
    result,score = model.predict(data,model_path=model_path,vectorizer_path=vector_path)
    print("\n".join(["{}\t{}\t{}".format(str(label),str(s[label]),d) for label,s,d in zip(result,score,data)]))