"""DeBERTa-v3 with a small head. One class, three jobs, chosen by `task`:
- "scores": six numbers per essay (mean pooling, linear head)
"""
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, AutoTokenizer

from coach.data import BR, TRAITS

BACKBONE = "microsoft/deberta-v3-base"


def tokenizer(name=BACKBONE):
    tok = AutoTokenizer.from_pretrained(name)
    tok.add_special_tokens({"additional_special_tokens": [BR]})
    return tok


class ScoreModel(nn.Module):
    def __init__(self, name=BACKBONE, n_tokens=None, n_out=len(TRAITS)):
        super().__init__()
        cfg = AutoConfig.from_pretrained(name)
        cfg.hidden_dropout_prob = 0.0  # the ELL winners all switched dropout off for regression
        cfg.attention_probs_dropout_prob = 0.0
        self.backbone = AutoModel.from_pretrained(name, config=cfg)
        if n_tokens:
            self.backbone.resize_token_embeddings(n_tokens)
        self.head = nn.Linear(cfg.hidden_size, n_out)

    def forward(self, input_ids, attention_mask):
        h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        m = attention_mask.unsqueeze(-1).to(h.dtype)
        pooled = (h * m).sum(1) / m.sum(1).clamp(min=1)
        return self.head(pooled)


def param_groups(model, lr, head_lr, decay=0.9, weight_decay=0.01):
    """Layer-wise learning-rate decay: the head and top layer at `lr`, each layer below at lr * decay."""
    groups = [{"params": list(model.head.parameters()), "lr": head_lr, "weight_decay": 0.0}]
    layers = list(model.backbone.encoder.layer)
    n = len(layers)
    for i, layer in enumerate(layers):
        groups.append({"params": list(layer.parameters()), "lr": lr * decay ** (n - 1 - i), "weight_decay": weight_decay})
    emb = [p for n_, p in model.backbone.named_parameters() if not n_.startswith("encoder.layer")]
    groups.append({"params": emb, "lr": lr * decay ** n, "weight_decay": weight_decay})
    return groups


class PartsModel(nn.Module):
    """Token classification over the whole essay: which argument part each token belongs to (BIO over 7 types), and
    how effective that part is (3 classes, read off the tokens of the part and averaged at decode time)."""

    def __init__(self, name=BACKBONE, n_tokens=None, n_bio=15, n_eff=3):
        super().__init__()
        cfg = AutoConfig.from_pretrained(name)
        self.backbone = AutoModel.from_pretrained(name, config=cfg)
        if n_tokens:
            self.backbone.resize_token_embeddings(n_tokens)
        self.bio = nn.Linear(cfg.hidden_size, n_bio)
        self.eff = nn.Linear(cfg.hidden_size, n_eff)
        self.head = nn.ModuleList([self.bio, self.eff])  # so param_groups finds "head"

    def forward(self, input_ids, attention_mask):
        h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        return self.bio(h), self.eff(h)


def parts_tokenizer(name=BACKBONE):
    from tokenizers import AddedToken
    tok = AutoTokenizer.from_pretrained(name)
    tok.add_tokens([AddedToken("\n", normalized=False)])  # keep paragraph breaks; the tokenizer drops them otherwise
    return tok