from typing import List

import spacy
from fastapi import FastAPI
from pydantic import BaseModel

import edsnlp

app = FastAPI(title="EDS-NLP", version=edsnlp.__version__)

nlp = spacy.blank("eds")

nlp.add_pipe("eds.sentences")

config = dict(
    regex=dict(
        covid=[
            "covid",
            r"covid[-\s]?19",
            r"sars[-\s]?cov[-\s]?2",
            r"corona[-\s]?virus",
        ],
    ),
    attr="LOWER",
)
nlp.add_pipe("eds.matcher", config=config)

nlp.add_pipe("eds.negation")
nlp.add_pipe("eds.family")
nlp.add_pipe("eds.hypothesis")
nlp.add_pipe("eds.reported_speech")


class Entity(BaseModel):  # (2)

    # OMOP-style attributes
    start: int
    end: int
    label: str
    lexical_variant: str
    normalized_variant: str

    # Qualifiers
    negated: bool
    hypothesis: bool
    family: bool
    reported_speech: bool


class Document(BaseModel):  # (1)
    text: str
    ents: List[Entity]


@app.post("/process", response_model=List[Document])  # (1)
async def process(
    notes: List[str],  # (2)
):

    documents = []

    for doc in nlp.pipe(notes):
        entities = []

        for ent in doc.ents:
            entity = Entity(
                start=ent.start_char,
                end=ent.end_char,
                label=ent.label_,
                lexical_variant=ent.text,
                normalized_variant=ent._.normalized_variant,
                negated=ent._.negation,
                hypothesis=ent._.hypothesis,
                family=ent._.family,
                reported_speech=ent._.reported_speech,
            )
            entities.append(entity)

        documents.append(
            Document(
                text=doc.text,
                ents=entities,
            )
        )

    return documents