import logging
import os
import time
from datetime import datetime

import pandas as pd
import pytest
import spacy
from helpers import make_nlp
from pytest import fixture

import edsnlp

os.environ["EDSNLP_MAX_CPU_WORKERS"] = "2"
os.environ["TZ"] = "Europe/Paris"

try:
    time.tzset()
except AttributeError:
    pass
logging.basicConfig(level=logging.INFO)
try:
    import torch.nn
except ImportError:
    torch = None

pytest.importorskip("rich")


def pytest_collection_modifyitems(items):
    """Run test_docs* at the end"""
    first_tests = []
    last_tests = []
    for item in items:
        if item.name.startswith("test_code_blocks"):
            last_tests.append(item)
        else:
            first_tests.append(item)
    items[:] = first_tests + last_tests


@fixture(scope="session", params=["eds", "fr"])
def lang(request):
    return request.param


@fixture(scope="session")
def nlp(lang):
    return make_nlp(lang)


@fixture(scope="session")
def nlp_eds():
    return make_nlp("eds")


@fixture
def blank_nlp(lang):
    if lang == "eds":
        model = spacy.blank("eds")
    else:
        model = edsnlp.blank("fr")
    model.add_pipe("eds.sentences")
    return model


def make_ml_pipeline():
    nlp = edsnlp.blank("eds")
    nlp.add_pipe("eds.sentences", name="sentences")
    nlp.add_pipe(
        "eds.transformer",
        name="transformer",
        config=dict(
            model="prajjwal1/bert-tiny",
            window=128,
            stride=96,
        ),
    )
    nlp.add_pipe(
        "eds.ner_crf",
        name="ner",
        config=dict(
            embedding=nlp.get_pipe("transformer"),
            mode="independent",
            target_span_getter=["ents", "ner-preds"],
            span_setter="ents",
        ),
    )
    ner = nlp.get_pipe("ner")
    ner.update_labels(["PERSON", "GIFT"])
    return nlp


@fixture()
def ml_nlp():
    if torch is None:
        pytest.skip("torch not installed", allow_module_level=False)
    return make_ml_pipeline()


@fixture(scope="session")
def frozen_ml_nlp():
    if torch is None:
        pytest.skip("torch not installed", allow_module_level=False)
    return make_ml_pipeline()


@fixture()
def text():
    return (
        "Le patient est admis pour des douleurs dans le bras droit, "
        "mais n'a pas de problème de locomotion. "
        "Historique d'AVC dans la famille. pourrait être un cas de rhume.\n"
        "NBNbWbWbNbWbNBNbNbWbWbNBNbWbNbNbWbNBNbWbNbNBWbWbNbNbNBWbNbWbNbWBNb"
        "NbWbNbNBNbWbWbNbWBNbNbWbNBNbWbWbNb\n"
        "Pourrait être un cas de rhume.\n"
        "Motif :\n"
        "Douleurs dans le bras droit.\n"
        "ANTÉCÉDENTS\n"
        "Le patient est déjà venu pendant les vacances\n"
        "d'été.\n"
        "Pas d'anomalie détectée."
    )


@fixture
def doc(nlp, text):
    return nlp(text)


@fixture
def blank_doc(blank_nlp, text):
    return blank_nlp(text)


@fixture
def df_notes():
    N_LINES = 100
    notes = pd.DataFrame(
        data={
            "note_id": list(range(N_LINES)),
            "note_text": N_LINES * [text],
            "note_datetime": N_LINES * [datetime.today()],
        }
    )

    return notes


def make_df_note(text, module):
    from pyspark.sql import types as T
    from pyspark.sql.session import SparkSession

    data = [(i, i // 5, text, datetime(2021, 1, 1)) for i in range(20)]

    if module == "pandas":
        return pd.DataFrame(
            data=data, columns=["note_id", "person_id", "note_text", "note_datetime"]
        )

    note_schema = T.StructType(
        [
            T.StructField("note_id", T.IntegerType()),
            T.StructField("person_id", T.IntegerType()),
            T.StructField("note_text", T.StringType()),
            T.StructField(
                "note_datetime",
                T.TimestampType(),
            ),
        ]
    )

    spark = SparkSession.builder.getOrCreate()
    notes = spark.createDataFrame(data=data, schema=note_schema)
    if module == "pyspark":
        return notes

    if module == "koalas":
        try:
            import databricks.koalas  # noqa F401
        except ImportError:
            pytest.skip("Koalas not installed")
        return notes.to_koalas()


@fixture
def df_notes_pandas(text):
    return make_df_note(text, "pandas")


@fixture
def df_notes_pyspark(text):
    return make_df_note(text, "pyspark")


@fixture
def df_notes_koalas(text):
    return make_df_note(text, "koalas")


@fixture
def run_in_test_dir(request, monkeypatch):
    monkeypatch.chdir(request.fspath.dirname)


@pytest.fixture(autouse=True)
def stop_spark():
    yield
    try:
        from pyspark.sql import SparkSession
    except ImportError:
        return
    try:
        getActiveSession = SparkSession.getActiveSession
    except AttributeError:

        def getActiveSession():  # pragma: no cover
            from pyspark import SparkContext

            sc = SparkContext._active_spark_context
            if sc is None:
                return None
            else:
                assert sc._jvm is not None
                if sc._jvm.SparkSession.getActiveSession().isDefined():
                    SparkSession(sc, sc._jvm.SparkSession.getActiveSession().get())
                    try:
                        return SparkSession._activeSession
                    except AttributeError:
                        try:
                            return SparkSession._instantiatedSession
                        except AttributeError:
                            return None
                else:
                    return None

    session = getActiveSession()
    if session is not None:
        session.stop()