from typing import List, Tuple
import torch
from SciAssist import ReferenceStringParsing

device = "gpu" if torch.cuda.is_available() else "cpu"
rsp_pipeline = ReferenceStringParsing(os_name="nt", device=device)


def rsp_for_str(input, dehyphen=False) -> List[Tuple[str, str]]:
    results = rsp_pipeline.predict(input, type="str", dehyphen=dehyphen)
    output = []
    for res in results:
        for token, tag in zip(res["tokens"], res["tags"]):
            output.append((token, tag))
        output.append(("\n\n", None))
    return output


def rsp_for_file(input, dehyphen=False) -> List[Tuple[str, str]]:
    if input == None:
        return None
    filename = input.name
    # Identify the format of input and parse reference strings
    if filename[-4:] == ".txt":
        results = rsp_pipeline.predict(filename, type="txt", dehyphen=dehyphen, save_results=False)
    elif filename[-4:] == ".pdf":
        results = rsp_pipeline.predict(filename, dehyphen=dehyphen, save_results=False)
    else:
        return [("File Format Error !", None)]
    # Prepare for the input gradio.HighlightedText accepts.
    output = []
    for res in results:
        for token, tag in zip(res["tokens"], res["tags"]):
            output.append((token, tag))
        output.append(("\n\n", None))
    return output