from collections import OrderedDict
from typing import Optional, List
from thinc.api import (
    with_array,
    chain,
    Model,
    PyTorchWrapper,
    PyTorchLSTM,
)
from thinc.types import Floats2d

from spacy.tokens import Doc
from spacy.util import registry
import torch
from torch import nn


@registry.architectures("TorchEntityRecognizer.v1")
def build_torch_ner_model(
    tok2vec: Model[List[Doc], List[Floats2d]],
    hidden_width: int,
    dropout: Optional[float] = None,
    nO: Optional[int] = None,
    width: Optional[int] = None,
) -> Model[List[Doc], List[Floats2d]]:
    """Build a tagger model, using a provided token-to-vector component. The tagger
    model simply adds a linear layer with softmax activation to predict scores
    given the token vectors.
    tok2vec (Model[List[Doc], List[Floats2d]]): The token-to-vector subnetwork.
    nO (int or None): The number of tags to output. Inferred from the data if None.
    RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
    """
    if width == None :
        t2v_width = tok2vec.maybe_get_dim("nO")
    else :
        t2v_width = width
    torch_model = TorchEntityRecognizer(t2v_width, hidden_width, nO, dropout)
    wrapped_pt_model = PyTorchWrapper(torch_model)
    wrapped_pt_model.attrs["set_dropout_rate"] = torch_model.set_dropout_rate
    
    model = chain(tok2vec, with_array(wrapped_pt_model))
    model.set_ref("tok2vec", tok2vec)
    model.set_ref("torch_model", wrapped_pt_model)
    model.init = init
    return model


def init(
    model: Model[List[Doc], Floats2d],
    X: Optional[List[Doc]] = None,
    Y: Optional[List[str]] = None,
) -> Model[List[Doc], List[Floats2d]]:
    """Dynamically set PyTorch Output Layer shape based on labels data
    model (Model[List[Doc], Floats2d]): Thinc Model wrapping tok2vec and PyTorch model
    X (Optional[List[Doc]], optional): Sample of Doc objects.
    Y (Optional[List[Ints2d]], optional): Available model labels.
    RETURNS (Model[List[Doc], List[Floats2d]]): Initialized Model
    """

    tok2vec = model.get_ref("tok2vec")

    torch_model = model.get_ref("torch_model")
    listener = tok2vec.maybe_get_ref("listener")

    t2v_width = listener.maybe_get_dim("nO") if listener else None

    if t2v_width:
        torch_model.shims[0]._model.set_input_shape(t2v_width)
        torch_model.set_dim("nI", t2v_width)

    if Y is not None:
        nO = len(Y)
        torch_model.shims[0]._model.set_output_shape(nO)
        torch_model.set_dim("nO", nO)

    tok2vec = model.get_ref("tok2vec")
    tok2vec.initialize()
    torch_model = model.get_ref("torch_model")
    return model


def is_dropout_module(
    module: nn.Module,
    dropout_modules: List[nn.Module] = [nn.Dropout, nn.Dropout2d, nn.Dropout3d],
) -> bool:
    """Detect if a PyTorch Module is a Dropout layer
    module (nn.Module): Module to check
    dropout_modules (List[nn.Module], optional): List of Modules that count as Dropout layers.
    RETURNS (bool): True if module is a Dropout layer.
    """
    for m in dropout_modules:
        if isinstance(module, m):
            return True
    return False


class TorchEntityRecognizer(nn.Module):
    """Torch Entity Recognizer Model Head"""

    def __init__(self, nI: int, nH: int, nO: int, dropout: float):
        """Initialize TorchEntityRecognizer.
        nI (int): Input Dimension
        nH (int): Hidden Dimension Width
        nO (int): Output Dimension Width
        dropout (float): Dropout ratio (0 - 1.0)
        """
        super(TorchEntityRecognizer, self).__init__()

        # Just for initialization of PyTorch layer. Output shape set during Model.init
        nI = nI or 1
        nO = nO or 1

        self.nH = nH
        self.model = nn.Sequential(
            OrderedDict(
                {
                    "input_layer": nn.Linear(nI, nH),
                    "input_activation": nn.ReLU(),
                    "input_dropout": nn.Dropout2d(dropout),
                    "output_layer": nn.Linear(nH, nO),
                    "output_dropout": nn.Dropout2d(dropout),
                    "softmax": nn.Softmax(dim=1),
                }
            )
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Forward pass of the model.
        inputs (torch.Tensor): Batch of outputs from spaCy tok2vec layer
        RETURNS (torch.Tensor): Batch of results with a score for each tag for each token
        """
        return self.model(inputs)

    def _set_layer_shape(self, name: str, nI: int, nO: int):
        """Dynamically set the shape of a layer
        name (str): Layer name
        nI (int): New input shape
        nO (int): New output shape
        """
        with torch.no_grad():
            layer = getattr(self.model, name)
            #print(layer)
            layer.out_features = nO
            layer.weight = nn.Parameter(torch.Tensor(nO, nI))
            if layer.bias is not None:
                layer.bias = nn.Parameter(torch.Tensor(nO))
            layer.reset_parameters()

    def set_input_shape(self, nI: int):
        """Dynamically set the shape of the input layer
        nI (int): New input layer shape
        """
        self._set_layer_shape("input_layer", nI, self.nH)

    def set_output_shape(self, nO: int):
        """Dynamically set the shape of the output layer
        nO (int): New output layer shape
        """
        self._set_layer_shape("output_layer", self.nH, nO)

    def set_dropout_rate(self, dropout: float):
        """Set the dropout rate of all Dropout layers in the model.
        dropout (float): Dropout rate to set
        """
        dropout_layers = [
            module for module in self.modules() if is_dropout_module(module)
        ]
        for layer in dropout_layers:
            layer.p = dropout