# ruff:noqa:E402

import pytest

try:

    import torch.nn

except ImportError:

    torch = None

if torch is None:

    pytest.skip("torch not installed", allow_module_level=True)

pytest.importorskip("rich")

import shutil

from typing import (

    Optional,

    Sequence,

    Union,

)

import pytest

import spacy.tokenizer

import torch.nn

from confit import Config

from confit.utils.random import set_seed

from spacy.tokens import Doc, Span

from edsnlp.core.registries import registry

from edsnlp.data.converters import AttributesMappingArg, get_current_tokenizer

from edsnlp.metrics.dep_parsing import DependencyParsingMetric

from edsnlp.training.optimizer import LinearSchedule, ScheduledOptimizer

from edsnlp.training.trainer import GenericScorer, train

from edsnlp.utils.span_getters import SpanSetterArg, set_spans

@registry.factory.register("myproject.custom_dict2doc", spacy_compatible=False)

class CustomSampleGenerator:

    def __init__(

        self,

        *,

        tokenizer: Optional[spacy.tokenizer.Tokenizer] = None,

        name: str = "myproject.custom_dict2doc",

        span_setter: SpanSetterArg = {"ents": True, "*": True},

        bool_attributes: Union[str, Sequence[str]] = [],

        span_attributes: Optional[AttributesMappingArg] = None,

    ):

        self.tokenizer = tokenizer

        self.name = name

        self.span_setter = span_setter

        self.bool_attributes = bool_attributes

        self.span_attributes = span_attributes

    def __call__(self, obj):

        tok = get_current_tokenizer() if self.tokenizer is None else self.tokenizer

        doc = tok(obj["note_text"] or "")

        doc._.note_id = obj.get("note_id", obj.get("__FILENAME__"))

        doc._.note_datetime = obj.get("note_datetime")

        spans = []

        if self.span_attributes is not None:

            for dst in self.span_attributes.values():

                if not Span.has_extension(dst):

                    Span.set_extension(dst, default=None)

        for ent in obj.get("entities") or ():

            ent = dict(ent)

            span = doc.char_span(

                ent.pop("start"),

                ent.pop("end"),

                label=ent.pop("label"),

                alignment_mode="expand",

            )

            for label, value in ent.items():

                new_name = (

                    self.span_attributes.get(label, None)

                    if self.span_attributes is not None

                    else label

                )

                if self.span_attributes is None and not Span.has_extension(new_name):

                    Span.set_extension(new_name, default=None)

                if new_name:

                    span._.set(new_name, value)

            spans.append(span)

        set_spans(doc, spans, span_setter=self.span_setter)

        for attr in self.bool_attributes:

            for span in spans:

                if span._.get(attr) is None:

                    span._.set(attr, False)

        return doc

def test_ner_qualif_train_diff_bert(run_in_test_dir, tmp_path):

    set_seed(42)

    config = Config.from_disk("ner_qlf_diff_bert_config.yml")

    shutil.rmtree(tmp_path, ignore_errors=True)

    kwargs = Config.resolve(config["train"], registry=registry, root=config)

    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)

    scorer = GenericScorer(**kwargs["scorer"])

    val_data = kwargs["val_data"]

    last_scores = scorer(nlp, val_data)

    # Check empty doc

    nlp("")

    assert last_scores["ner"]["micro"]["f"] > 0.4

    assert last_scores["qual"]["micro"]["f"] > 0.4

def test_ner_qualif_train_same_bert(run_in_test_dir, tmp_path):

    set_seed(42)

    config = Config.from_disk("ner_qlf_same_bert_config.yml")

    shutil.rmtree(tmp_path, ignore_errors=True)

    kwargs = Config.resolve(config["train"], registry=registry, root=config)

    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)

    scorer = GenericScorer(**kwargs["scorer"])

    val_data = kwargs["val_data"]

    last_scores = scorer(nlp, val_data)

    # Check empty doc

    nlp("")

    assert last_scores["ner"]["micro"]["f"] > 0.4

    assert last_scores["qual"]["micro"]["f"] > 0.4

def test_qualif_train(run_in_test_dir, tmp_path):

    set_seed(42)

    config = Config.from_disk("qlf_config.yml")

    shutil.rmtree(tmp_path, ignore_errors=True)

    kwargs = Config.resolve(config["train"], registry=registry, root=config)

    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)

    scorer = GenericScorer(**kwargs["scorer"])

    val_data = kwargs["val_data"]

    last_scores = scorer(nlp, val_data)

    # Check empty doc

    nlp("")

    assert last_scores["qual"]["micro"]["f"] >= 0.4

def test_dep_parser_train(run_in_test_dir, tmp_path):

    set_seed(42)

    config = Config.from_disk("dep_parser_config.yml")

    shutil.rmtree(tmp_path, ignore_errors=True)

    kwargs = Config.resolve(config["train"], registry=registry, root=config)

    nlp = train(**kwargs, output_dir=tmp_path, cpu=True)

    scorer = GenericScorer(**kwargs["scorer"])

    val_data = list(kwargs["val_data"])

    last_scores = scorer(nlp, val_data)

    scorer_bis = GenericScorer(parser=DependencyParsingMetric(filter_expr="False"))

    # Just to test what happens if the scores indicate 2 roots

    val_data_bis = [Doc.from_docs([val_data[0], val_data[0]])]

    nlp.pipes.parser.decoding_mode = "mst"

    last_scores_bis = scorer_bis(nlp, val_data_bis)

    assert last_scores_bis["parser"]["uas"] == 0.0

    # Check empty doc

    nlp("")

    assert last_scores["dep"]["las"] >= 0.4

def test_optimizer():

    net = torch.nn.Linear(10, 10)

    optim = ScheduledOptimizer(

        torch.optim.AdamW,

        module=net,

        total_steps=10,

        groups={

            ".*": {

                "lr": 9e-4,

                "schedules": LinearSchedule(

                    warmup_rate=0.1,

                    start_value=0,

                ),

            }

        },

    )

    for param in net.parameters():

        assert "exp_avg" not in optim.optim.state[param]

    optim.initialize()

    for param in net.parameters():

        assert "exp_avg" in optim.optim.state[param]

    lr_values = [optim.optim.param_groups[0]["lr"]]

    for i in range(10):

        optim.step()

        lr_values.append(optim.optim.param_groups[0]["lr"])

    # close enough

    assert lr_values == pytest.approx(

        [

            0.0,

            0.0009,

            0.0008,

            0.0007,

            0.0006,

            0.0005,

            0.0004,

            0.0003,

            0.0002,

            0.0001,

            0.0,

        ]

    )