#pip install rich 
#Formats any text/data printed to the console. 
from rich import print

import spacy
from spacy.tokens import DocBin
from sklearn.model_selection import train_test_split
from typing import Set, List, Tuple
import sys
import os
current = os.path.dirname(os.path.realpath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)
import aeg


nlp = spacy.load("en_core_web_md")
train_data = aeg.format_data()
valid_data = aeg.format_data()
#Categories to choose from: 
cats = ['issue', 'rule', 'policy', 'inference', 'conclusion']

print(train_data)


#Formats data/prepares for training: 
def make_docs(data: List[Tuple[str, str]], target_file: str, cats: Set[str]):
    docs = DocBin()
    for doc, label in nlp.pipe(data, as_tuples = True):
        for cat in cats: 
            doc.cats[cat] = 1 if cat == label else 0
  
        docs.add(doc)
    docs.to_disk(target_file)
    return (docs)

make_docs(train_data, "data/train.spacy", cats)
make_docs(valid_data, "data/valid.spacy", cats)