# -*- coding: utf-8 -*-
# file: preprocess.py
# author: jackie
# Copyright (C) 2021. All Rights Reserved.

import os
import pandas as pd
import argparse
import emoji
import re
from sklearn.model_selection import train_test_split

parser = argparse.ArgumentParser()
parser.add_argument("--inpath", type=str, required=True, default='./raw_data/data1.csv')
parser.add_argument("--folder_name", type=str, required=False, default='./custom')
parser.add_argument("--task", type=str, required=False, default='aptepc')
args = parser.parse_args()


def convert(text, labels):
    # convert label to list
    try:
        labels = eval(labels)
        tags = ['O'] * len(text)
        sentiment = ['-999'] * len(text)

        for j in range(len(labels)):
            label = labels[j]
            sentiment_key = labels[j][3]

            if sentiment_key == '正':
                sentiment_value = 'Positive'
            elif sentiment_key == '负':
                sentiment_value = 'Negative'
            else:
                sentiment_value = 'Others'

            tags[label[4][0]] = 'B-ASP'
            sentiment[label[4][0]] = sentiment_value

            k = label[4][0] + 1
            while k < label[4][1]:
                tags[k] = 'I-ASP'
                sentiment[k] = sentiment_value

                k += 1
        return text, tags, sentiment
    except:
        print ("labels", labels)
        print ("text", text)

def convert_tag(text, labels):
    # convert label to list
    try:
        labels = eval(labels)
        tags = ['O'] * len(text)
        sentiment = ['-999'] * len(text)

        for j in range(len(labels)):
            label = labels[j]
            sentiment_key = labels[j][3]

            if sentiment_key == '正':
                sentiment_value = 'Positive'
            elif sentiment_key == '负':
                sentiment_value = 'Negative'
            else:
                sentiment_value = 'Others'

            tags[label[4][0]] = 'B-'+label[1]
            sentiment[label[4][0]] = sentiment_value

            k = label[4][0] + 1
            while k < label[4][1]:
                tags[k] = 'I-'+label[1]
                sentiment[k] = sentiment_value

                k += 1
        return text, tags, sentiment
    except:
        print ("labels", labels)
        print ("text", text)

def convert_sentiment(sentiment_key):
    if sentiment_key == '正':
        sentiment_value = 'Positive'
    else:
        sentiment_value = 'Negative'
    return sentiment_value


def convert_apc(text, label):
    label_update = [(i[0], i[3], i[4]) for i in eval(label)]
    label_update = list(set(label_update))

    str1_list = []
    str2_list = []
    str3_list = []

    for j in range(len(label_update)):
        str1 = text[:label_update[j][2][0]] + '$T$ ' + text[label_update[j][2][1]:]
        str1_list.append(str1)
        str2_list.append(label_update[j][0])
        str3_list.append(convert_sentiment(label_update[j][1]))

    return str1_list, str2_list, str3_list


def filter_emoji(desstr, restr=''):
    # 过滤表情
    try:
        co = re.compile(u'[\U00010000-\U0010ffff]')
    except re.error:
        co = re.compile(u'[\uD800-\uDBFF][\uDC00-\uDFFF]')
    return co.sub(restr, desstr)


def convert_to_atepc(inpath, dist_fname, flag):
    # 写之前，先检验文件是否存在，存在就删掉
    if os.path.exists(dist_fname):
        os.remove(dist_fname)
    f1 = open(dist_fname, 'w', encoding='utf8')

    data = pd.read_csv(inpath)
    data.columns = ['text', 'tag_sentiment_list']

    # preprocess for emoji
    data['text'] = data['text'].map(lambda x: filter_emoji(x, restr='xx'))

    # 只保留review的长度小于600的
    data = data[data['text'].str.len() <= 600]

    # train test split
    x_train, x_test = train_test_split(data, test_size=0.2, random_state=42)

    if flag == 'train':
        data_res = x_train.iloc[:, :].reset_index()
    else:
        data_res = x_test.iloc[:, :].reset_index()
    # print (data_res.head())
    for i in range(len(data_res)):
        text, label = data_res['text'][i], data_res['tag_sentiment_list'][i]
        text, tags, sentiment = convert(text, label)

        for word, tag, sen in zip(text, tags, sentiment):
            if word not in ['，', '。', ' ', '\xa0', '\u2006', '\u3000', '\u2002', '\u2003', '\u2005', '\x0c', '\u2028',
                            '\u2009', '\u200a']:
                f1.write(word + ' ' + tag + ' ' + sen + '\n')
            else:
                f1.write("\n")

        f1.write("\n")
    f1.close()
    print ("process atepc finished!")

def convert_to_atepc_tag(inpath, dist_fname, flag):
    # 写之前，先检验文件是否存在，存在就删掉
    if os.path.exists(dist_fname):
        os.remove(dist_fname)
    f1 = open(dist_fname, 'w', encoding='utf8')

    data = pd.read_csv(inpath)
    data.columns = ['text', 'tag_sentiment_list']

    # preprocess for emoji
    data['text'] = data['text'].map(lambda x: filter_emoji(x, restr='xx'))

    # drop id list not able to process
    # print (data.iloc[8832,:])
    # data = data.drop([8832])

    # 只保留review的长度小于600的
    data = data[data['text'].str.len() <= 600]

    # train test split
    x_train, x_test = train_test_split(data, test_size=0.2, random_state=42)

    if flag == 'train':
        data_res = x_train.iloc[:, :].reset_index()
    else:
        data_res = x_test.iloc[:, :].reset_index()
    # print (data_res.head())
    for i in range(len(data_res)):
        text, label = data_res['text'][i], data_res['tag_sentiment_list'][i]
        text, tags, sentiment = convert(text, label)

        for word, tag, sen in zip(text, tags, sentiment):
            if word not in ['，', '。', ' ', '\xa0', '\u2006', '\u3000', '\u2002', '\u2003', '\u2005', '\x0c', '\u2028',
                            '\u2009', '\u200a']:
                f1.write(word + ' ' + tag + ' ' + sen + '\n')
            else:
                f1.write("\n")

        f1.write("\n")
    f1.close()
    print ("process atepc finished!")

def convert_to_apc(inpath, dist_fname, flag):
    # 写之前，先检验文件是否存在，存在就删掉
    if os.path.exists(dist_fname):
        os.remove(dist_fname)
    f1 = open(dist_fname, 'w', encoding='utf8')

    data = pd.read_csv(inpath)
    # train test split
    x_train, x_test = train_test_split(data, test_size=0.2, random_state=42)

    if flag == 'train':
        data_res = x_train.iloc[:, :].reset_index()
    else:
        data_res = x_test.iloc[:, :].reset_index()
    # print (data_res.head())
    for i in range(len(data_res)):
        text, label = data_res['text'][i], data_res['tag_sentiment_list'][i]
        str1_list, str2_list, str3_list = convert_apc(text, label)

        for x1, x2, x3 in zip(str1_list, str2_list, str3_list):
            f1.write(x1 + '\n')
            f1.write(x2 + '\n')
            f1.write(x3 + '\n')

    f1.close()

    print ("process apc finished!")


def main(inpath, folder_name, task):
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)

    if task == 'aptepc':
        # get folder name
        print ("start process for an aptepc task")
        folder_name_prefix = folder_name.split('/')[-1]
        dist_train_fname = os.path.join(folder_name_prefix, folder_name_prefix + '.train.txt.atepc')
        dist_test_fname = os.path.join(folder_name_prefix, folder_name_prefix + '.test.txt.atepc')
        # process train
        convert_to_atepc(inpath, dist_train_fname, 'train')
        print ("<<< finish training data preprocess")
        # process test
        convert_to_atepc(inpath, dist_test_fname, 'test')
        print ("<<< finish test data preprocess")
    elif task == 'apc':
        # get folder name
        folder_name_prefix = folder_name.split('/')[-1]
        dist_train_fname = os.path.join(folder_name_prefix, folder_name_prefix + '.train.txt')
        dist_test_fname = os.path.join(folder_name_prefix, folder_name_prefix + '.test.txt')
        # process train
        convert_to_apc(inpath, dist_train_fname, 'train')
        print ("<<< finish training data preprocess")
        # process test
        convert_to_apc(inpath, dist_test_fname, 'test')
        print ("<<< finish test data preprocess")
    elif task == 'aptepc-tag':
        # get folder name
        print ("start process for an aptepc tag task")
        folder_name_prefix = folder_name.split('/')[-1]
        dist_train_fname = os.path.join(folder_name_prefix, folder_name_prefix + '.train.txt.atepc')
        dist_test_fname = os.path.join(folder_name_prefix, folder_name_prefix + '.test.txt.atepc')
        # process train
        convert_to_atepc_tag(inpath, dist_train_fname, 'train')
        print ("<<< finish training data preprocess")
        # process test
        convert_to_atepc_tag(inpath, dist_test_fname, 'test')
        print ("<<< finish test data preprocess")


main(args.inpath, args.folder_name, args.task)
[((13, 9, 13, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((120, 7, 120, 33), 'os.path.exists', 'os.path.exists', ({(120, 22, 120, 32): 'dist_fname'}, {}), '(dist_fname)', False, 'import os\n'), ((124, 11, 124, 30), 'pandas.read_csv', 'pd.read_csv', ({(124, 23, 124, 29): 'inpath'}, {}), '(inpath)', True, 'import pandas as pd\n'), ((134, 22, 134, 76), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((158, 7, 158, 33), 'os.path.exists', 'os.path.exists', ({(158, 22, 158, 32): 'dist_fname'}, {}), '(dist_fname)', False, 'import os\n'), ((162, 11, 162, 30), 'pandas.read_csv', 'pd.read_csv', ({(162, 23, 162, 29): 'inpath'}, {}), '(inpath)', True, 'import pandas as pd\n'), ((176, 22, 176, 76), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((200, 7, 200, 33), 'os.path.exists', 'os.path.exists', ({(200, 22, 200, 32): 'dist_fname'}, {}), '(dist_fname)', False, 'import os\n'), ((204, 11, 204, 30), 'pandas.read_csv', 'pd.read_csv', ({(204, 23, 204, 29): 'inpath'}, {}), '(inpath)', True, 'import pandas as pd\n'), ((206, 22, 206, 76), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((112, 13, 112, 51), 're.compile', 're.compile', ({(112, 24, 112, 50): 'u"""[𐀀-\U0010ffff]"""'}, {}), "(u'[𐀀-\\U0010ffff]')", False, 'import re\n'), ((121, 8, 121, 29), 'os.remove', 'os.remove', ({(121, 18, 121, 28): 'dist_fname'}, {}), '(dist_fname)', False, 'import os\n'), ((159, 8, 159, 29), 'os.remove', 'os.remove', ({(159, 18, 159, 28): 'dist_fname'}, {}), '(dist_fname)', False, 'import os\n'), ((201, 8, 201, 29), 'os.remove', 'os.remove', ({(201, 18, 201, 28): 'dist_fname'}, {}), '(dist_fname)', False, 'import os\n'), ((228, 11, 228, 38), 'os.path.exists', 'os.path.exists', ({(228, 26, 228, 37): 'folder_name'}, {}), '(folder_name)', False, 'import os\n'), ((229, 8, 229, 32), 'os.makedirs', 'os.makedirs', ({(229, 20, 229, 31): 'folder_name'}, {}), '(folder_name)', False, 'import os\n'), ((235, 27, 235, 100), 'os.path.join', 'os.path.join', ({(235, 40, 235, 58): 'folder_name_prefix', (235, 60, 235, 99): "folder_name_prefix + '.train.txt.atepc'"}, {}), "(folder_name_prefix, folder_name_prefix + '.train.txt.atepc')", False, 'import os\n'), ((236, 26, 236, 98), 'os.path.join', 'os.path.join', ({(236, 39, 236, 57): 'folder_name_prefix', (236, 59, 236, 97): "folder_name_prefix + '.test.txt.atepc'"}, {}), "(folder_name_prefix, folder_name_prefix + '.test.txt.atepc')", False, 'import os\n'), ((114, 13, 114, 58), 're.compile', 're.compile', ({(114, 24, 114, 57): "u'[\\ud800-\\udbff][\\udc00-\\udfff]'"}, {}), "(u'[\\ud800-\\udbff][\\udc00-\\udfff]')", False, 'import re\n'), ((246, 27, 246, 94), 'os.path.join', 'os.path.join', ({(246, 40, 246, 58): 'folder_name_prefix', (246, 60, 246, 93): "folder_name_prefix + '.train.txt'"}, {}), "(folder_name_prefix, folder_name_prefix + '.train.txt')", False, 'import os\n'), ((247, 26, 247, 92), 'os.path.join', 'os.path.join', ({(247, 39, 247, 57): 'folder_name_prefix', (247, 59, 247, 91): "folder_name_prefix + '.test.txt'"}, {}), "(folder_name_prefix, folder_name_prefix + '.test.txt')", False, 'import os\n'), ((258, 27, 258, 100), 'os.path.join', 'os.path.join', ({(258, 40, 258, 58): 'folder_name_prefix', (258, 60, 258, 99): "folder_name_prefix + '.train.txt.atepc'"}, {}), "(folder_name_prefix, folder_name_prefix + '.train.txt.atepc')", False, 'import os\n'), ((259, 26, 259, 98), 'os.path.join', 'os.path.join', ({(259, 39, 259, 57): 'folder_name_prefix', (259, 59, 259, 97): "folder_name_prefix + '.test.txt.atepc'"}, {}), "(folder_name_prefix, folder_name_prefix + '.test.txt.atepc')", False, 'import os\n')]