import sys
import os

sys.path.insert(1, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.parse_data import parse_concept, parse_ast, find_char_indexes
import pandas as pd
from tqdm import tqdm
import spacy_streamlit


def visualize_record(record, task="assertion"):
    """
    Visualize a single record.
    task can be "concept" or "assertion"
    """
    assert task in ["concept", "assertion"], "task must be 'concept' or 'assertion'"
    
    filename = record["filename"]
    with open("data/train/txt" + os.sep + filename + ".txt") as f:
        text = f.read()
    lines = text.split("\n")

    # parse file concepts and assertions
    if task == "concept":
        df = pd.DataFrame(parse_concept("data/train/concept/" + filename + ".con"))
        df = df.apply(find_char_indexes, axis=1, args=(text,))
        df = df.rename(columns={"start_char_index": "start", "end_char_index": "end", "concept_type": "label"})
        possible_labels = ["problem", "test", "treatment"]
    elif task == "assertion":
        df = pd.DataFrame(parse_ast("data/train/ast/" + filename + ".ast"))
        df = df.apply(find_char_indexes, axis=1, args=(text,))
        df = df.rename(columns={"start_char_index": "start", "end_char_index": "end", "assertion_type": "label"})
        possible_labels = ["present", "possible", "absent", "conditional", "hypothetical", "associated_with_someone_else"]

    doc = [
        {
            "text": line,
            "ents": df[df["start_line"] == i + 1][["start", "end", "label"]].to_dict(orient="records"),
        }
        for i, line in enumerate(lines)
    ]

    return spacy_streamlit.visualize_ner(
        doc,
        labels=possible_labels,
        show_table=False,
        title="",
        manual=True,
        displacy_options={
            "colors": {
                "problem": "#f08080",
                "treatment": "#9bddff",
                "test": "#ffdab9",
                "present": "#f08080",
                "possible": "#00ffff",
                "absent": "#ff00ff",
                "conditional": "#ffa500",
                "hypothetical": "#ffdab9",
                "associated_with_someone_else": "#00ff7f"
            },
        },
    )


if __name__ == "__main__":
    visualize_record({"filename": "018636330_DH"})