import random

import re

from random import choice, randint

from random import random as randomf

from string import ascii_letters, ascii_lowercase

import pandas as pd

import pytest

from edsnlp.connectors.omop import OmopConnector

random.seed(10)

def random_word():

    n = randint(1, 20)

    return "".join(

        [choice(ascii_letters)] + [choice(ascii_lowercase) for _ in range(n)]

    )

def random_text():

    n = randint(30, 60)

    return " ".join([random_word() for _ in range(n)])

def random_note_nlp(text):

    ents = []

    for match in re.finditer(r"\w+", text):

        if randomf() > 0.8:

            ent = dict(

                start_char=match.start(),

                end_char=match.end(),

                lexical_variant=match.group(),

                note_nlp_source_value=random_word().lower(),

                negation=randomf() > 0.5,

            )

            ents.append(ent)

    return ents

@pytest.fixture

def note():

    df = pd.DataFrame(dict(note_text=[random_text() for _ in range(10)]))

    df["note_id"] = range(len(df))

    df["note_datetime"] = "2021-10-19"

    return df

@pytest.fixture

def note_nlp(note):

    df = note.copy()

    df["ents"] = df.note_text.apply(random_note_nlp)

    df = df.explode("ents")

    df = pd.concat([df, df.ents.apply(pd.Series)], axis=1)

    df = df.drop(columns=["note_text", "ents", "note_datetime"])

    df["note_nlp_id"] = range(len(df))

    return df

@pytest.fixture

def omop(blank_nlp) -> OmopConnector:

    blank_nlp.add_pipe("negation")

    return OmopConnector(blank_nlp)

@pytest.fixture

def docs(omop: OmopConnector, note, note_nlp):

    return omop.omop2docs(note, note_nlp, extensions=["negation"])

def test_omop2docs(docs, note, note_nlp):

    lexical_variants = note_nlp.groupby("note_id")["lexical_variant"].agg(list)

    for doc, text, lvs in zip(docs, note.note_text, lexical_variants):

        assert doc.text == text

        assert len(doc.ents) == len(lvs)

        for ent, lv in zip(doc.ents, lvs):

            assert ent.text == lv

def test_docs2omop(omop: OmopConnector, docs):

    note, note_nlp = omop.docs2omop(docs, extensions=["negation"])

    lexical_variants = note_nlp.groupby("note_id")["lexical_variant"].agg(list)

    for doc, text, lvs in zip(docs, note.note_text, lexical_variants):

        assert doc.text == text

        assert len(doc.ents) == len(lvs)

        for ent, lv in zip(doc.ents, lvs):

            assert ent.text == lv

def test_roundtrip(omop: OmopConnector, docs, note, note_nlp):

    note2, note_nlp2 = omop.docs2omop(docs, extensions=["negation"])

    assert (note2 == note[note2.columns]).all().all()

    assert (note_nlp2 == note_nlp[note_nlp2.columns]).all().all()