from typing import List, Callable
from thinc.api import Model, chain
from thinc.types import Ragged, Floats2d
from spacy.tokens import Doc

from .layers import TransformerModel, TransformerListener
from .layers import trfs2arrays, split_trf_batch
from .util import registry
from .data_classes import FullTransformerBatch


@registry.architectures.register("spacy-transformers.TransformerListener.v1")
def transformer_listener_tok2vec_v1(
    pooling: Model[Ragged, Floats2d], grad_factor: float = 1.0, upstream: str = "*"
) -> Model[List[Doc], List[Floats2d]]:
    """Create a 'TransformerListener' layer, which will connect to a Transformer
    component earlier in the pipeline.

    The layer takes a list of Doc objects as input, and produces a list of
    2d arrays as output, with each array having one row per token. Most spaCy
    models expect a sublayer with this signature, making it easy to connect them
    to a transformer model via this sublayer.
    Transformer models usually operate over wordpieces, which usually don't align
    one-to-one against spaCy tokens. The layer therefore requires a reduction
    operation in order to calculate a single token vector given zero or more
    wordpiece vectors.

    pooling (Model[Ragged, Floats2d]): A reduction layer used to calculate
        the token vectors based on zero or more wordpiece vectors. If in doubt,
        mean pooling (see `thinc.layers.reduce_mean`) is usually a good choice.
    grad_factor (float): Reweight gradients from the component before passing
        them upstream. You can set this to 0 to "freeze" the transformer weights
        with respect to the component, or use it to make some components more
        significant than others. Leaving it at 1.0 is usually fine.
    upstream (str): A string to identify the 'upstream' Transformer
        to communicate with. The upstream name should either be the wildcard
        string '*', or the name of the `Transformer` component. You'll almost
        never have multiple upstream Transformer components, so the wildcard
        string will almost always be fine.
    """
    listener = TransformerListener(upstream_name=upstream)
    model: Model = chain(listener, trfs2arrays(pooling, grad_factor))
    model.set_ref("listener", listener)
    return model


@registry.architectures.register("spacy-transformers.Tok2VecTransformer.v1")
def transformer_tok2vec_v1(
    name: str,
    get_spans,
    tokenizer_config: dict,
    pooling: Model[Ragged, Floats2d],
    grad_factor: float = 1.0,
) -> Model[List[Doc], List[Floats2d]]:
    """Use a transformer as a "Tok2Vec" layer directly. This does not allow
    multiple components to share the transformer weights, and does not allow
    the transformer to set annotations into the `Doc` object, but it's a
    simpler solution if you only need the transformer within one component.

    get_spans (Callable[[List[Doc]], List[List[Span]]]): A function to extract
        spans from the batch of Doc objects. See the "TransformerModel" layer
        for details.
    tokenizer_config (dict): Settings to pass to the transformers tokenizer.
    pooling (Model[Ragged, Floats2d]): A reduction layer used to calculate
        the token vectors based on zero or more wordpiece vectors. If in doubt,
        mean pooling (see `thinc.layers.reduce_mean`) is usually a good choice.
     grad_factor (float): Reweight gradients from the component before passing
        them to the transformer. You can set this to 0 to "freeze" the transformer
        weights with respect to the component, or to make it learn more slowly.
        Leaving it at 1.0 is usually fine.
    """
    return chain(
        TransformerModel(name, get_spans, tokenizer_config),
        split_trf_batch(),
        trfs2arrays(pooling, grad_factor),
    )


@registry.architectures.register("spacy-transformers.Tok2VecTransformer.v2")
def transformer_tok2vec_v2(
    name: str,
    get_spans,
    tokenizer_config: dict,
    pooling: Model[Ragged, Floats2d],
    grad_factor: float = 1.0,
    transformer_config: dict = {},
) -> Model[List[Doc], List[Floats2d]]:
    """Use a transformer as a "Tok2Vec" layer directly. This does not allow
    multiple components to share the transformer weights, and does not allow
    the transformer to set annotations into the `Doc` object, but it's a
    simpler solution if you only need the transformer within one component.

    get_spans (Callable[[List[Doc]], List[List[Span]]]): A function to extract
        spans from the batch of Doc objects. See the "TransformerModel" layer
        for details.
    tokenizer_config (dict): Settings to pass to the transformers tokenizer.
    pooling (Model[Ragged, Floats2d]): A reduction layer used to calculate
        the token vectors based on zero or more wordpiece vectors. If in doubt,
        mean pooling (see `thinc.layers.reduce_mean`) is usually a good choice.
    grad_factor (float): Reweight gradients from the component before passing
        them to the transformer. You can set this to 0 to "freeze" the transformer
        weights with respect to the component, or to make it learn more slowly.
        Leaving it at 1.0 is usually fine.
    transformers_config (dict): Settings to pass to the transformers forward pass
        of the transformer.
    """
    return chain(
        TransformerModel(name, get_spans, tokenizer_config, transformer_config),
        split_trf_batch(),
        trfs2arrays(pooling, grad_factor),
    )


# Note: when updating, also make sure to update 'replace_listener_cfg' in _util.py
@registry.architectures.register("spacy-transformers.Tok2VecTransformer.v3")
def transformer_tok2vec_v3(
    name: str,
    get_spans,
    tokenizer_config: dict,
    pooling: Model[Ragged, Floats2d],
    grad_factor: float = 1.0,
    transformer_config: dict = {},
    mixed_precision: bool = False,
    grad_scaler_config: dict = {},
) -> Model[List[Doc], List[Floats2d]]:
    """Use a transformer as a "Tok2Vec" layer directly. This does not allow
    multiple components to share the transformer weights, and does not allow
    the transformer to set annotations into the `Doc` object, but it's a
    simpler solution if you only need the transformer within one component.

    get_spans (Callable[[List[Doc]], List[List[Span]]]): A function to extract
        spans from the batch of Doc objects. See the "TransformerModel" layer
        for details.
    tokenizer_config (dict): Settings to pass to the transformers tokenizer.
    pooling (Model[Ragged, Floats2d]): A reduction layer used to calculate
        the token vectors based on zero or more wordpiece vectors. If in doubt,
        mean pooling (see `thinc.layers.reduce_mean`) is usually a good choice.
    grad_factor (float): Reweight gradients from the component before passing
        them to the transformer. You can set this to 0 to "freeze" the transformer
        weights with respect to the component, or to make it learn more slowly.
        Leaving it at 1.0 is usually fine.
    transformers_config (dict): Settings to pass to the transformers forward pass
        of the transformer.
    mixed_precision (bool): Enable mixed-precision. Mixed-precision replaces
        whitelisted ops to half-precision counterparts. This speeds up training
        and prediction on modern GPUs and reduces GPU memory use.
    grad_scaler_config (dict): Configuration for gradient scaling in mixed-precision
        training. Gradient scaling is enabled automatically when mixed-precision
        training is used.

        Setting `enabled` to `False` in the gradient scaling configuration disables
        gradient scaling. The `init_scale` (default: `2 ** 16`) determines the
        initial scale. `backoff_factor` (default: `0.5`) specifies the factor
        by which the scale should be reduced when gradients overflow.
        `growth_interval` (default: `2000`) configures the number of steps
        without gradient overflows after which the scale should be increased.
        Finally, `growth_factor` (default: `2.0`) determines the factor by which
        the scale should be increased when no overflows were found for
        `growth_interval` steps.
    """
    # Note that this is a chain of chain on purpose, to match the structure of
    # TransformerListener.v1 after it is run through replace_listener (cf PR #310)
    return chain(  # type: ignore
        chain(
            TransformerModel(
                name,
                get_spans,
                tokenizer_config,
                transformer_config,
                mixed_precision,
                grad_scaler_config,
            ),
            split_trf_batch(),
        ),
        trfs2arrays(pooling, grad_factor),
    )


@registry.architectures.register("spacy-transformers.TransformerModel.v1")
def create_TransformerModel_v1(
    name: str,
    get_spans: Callable,
    tokenizer_config: dict = {},
) -> Model[List[Doc], FullTransformerBatch]:
    model = TransformerModel(name, get_spans, tokenizer_config)
    return model


@registry.architectures.register("spacy-transformers.TransformerModel.v2")
def create_TransformerModel_v2(
    name: str,
    get_spans: Callable,
    tokenizer_config: dict = {},
    transformer_config: dict = {},
) -> Model[List[Doc], FullTransformerBatch]:
    model = TransformerModel(name, get_spans, tokenizer_config, transformer_config)
    return model


@registry.architectures.register("spacy-transformers.TransformerModel.v3")
def create_TransformerModel_v3(
    name: str,
    get_spans: Callable,
    tokenizer_config: dict = {},
    transformer_config: dict = {},
    mixed_precision: bool = False,
    grad_scaler_config: dict = {},
) -> Model[List[Doc], FullTransformerBatch]:
    """Pretrained transformer model that can be finetuned for downstream tasks.

    name (str): Name of the pretrained Huggingface model to use.
    get_spans (Callable[[List[Doc]], List[List[Span]]]): A function to extract
        spans from the batch of Doc objects. See the "TransformerModel" layer
        for details.
    tokenizer_config (dict): Settings to pass to the transformers tokenizer.
    transformers_config (dict): Settings to pass to the transformers forward pass
        of the transformer.
    mixed_precision (bool): Enable mixed-precision. Mixed-precision replaces
        whitelisted ops to half-precision counterparts. This speeds up training
        and prediction on modern GPUs and reduces GPU memory use.
    grad_scaler_config (dict): Configuration for gradient scaling in mixed-precision
        training. Gradient scaling is enabled automatically when mixed-precision
        training is used.

        Setting `enabled` to `False` in the gradient scaling configuration disables
        gradient scaling. The `init_scale` (default: `2 ** 16`) determines the
        initial scale. `backoff_factor` (default: `0.5`) specifies the factor
        by which the scale should be reduced when gradients overflow.
        `growth_interval` (default: `2000`) configures the number of steps
        without gradient overflows after which the scale should be increased.
        Finally, `growth_factor` (default: `2.0`) determines the factor by which
        the scale should be increased when no overflows were found for
        `growth_interval` steps.
    """
    model = TransformerModel(
        name,
        get_spans,
        tokenizer_config,
        transformer_config,
        mixed_precision,
        grad_scaler_config,
    )
    return model
