from pathlib import Path


def getTrainData(datapath):
    """
    Load crf's train.data to get the train data in the format below.

    Args:
        train.data with "\n" divides the articles:
            只 O\n
            是 O\n
            前 B-time\n
            天 I-time\n
            ...
            \n
            只 O\n

    Returns:
        train_texts (dataset_size, content_size):
            [
                ['只', '是', '前', '天', '好', '很', ...],
                ['只', '是', '前', '天', '好', '很', ...],
                ...
            ]

        train_tags (dataset_size, content_size):
            [
                ['O', 'O', 'B-time', 'I-time', 'O', 'O', ...],
                ['O', 'O', 'B-time', 'I-time', 'O', 'O', ...],
                ...
            ]
    """
    train_texts = []
    train_tags = []
    with Path(datapath).open(encoding="UTF8") as f:
        train_data = f.readlines()
        texts = []
        tags = []

        # content of each line: "中 B-name\n"
        # line[0] = 中
        # line[2:-1] = B-name
        for line in train_data:
            if line == "\n":
                train_texts.append(texts)
                train_tags.append(tags)
                texts = []
                tags = []
                continue

            text, tag = line[0], line[2:-1]
            texts.append(text)
            tags.append(tag)

    return train_texts, train_tags