"""Multilingual encoder with text, context and speech-state fusion."""
from pathlib import Path

import torch
from torch import nn
from transformers import AutoConfig, AutoModel, AutoTokenizer
from safetensors.torch import load_file

from .features import collate


class VoiceEncoder(nn.Module):
    def __init__(self, config, *, pretrained=None, pooling='cls', dropout=.15):
        super().__init__()
        if pooling not in ('cls', 'mean'):
            raise ValueError('Unsupported pooling')
        options = {'trust_remote_code': False,
                   'attn_implementation': 'sdpa' if config.model_type == 'modernbert' else 'eager'}
        if config.model_type == 'modernbert':
            config.reference_compile = False
        if config.model_type == 'bert':
            options['add_pooling_layer'] = False
        self.encoder = (AutoModel.from_pretrained(pretrained, config=config, local_files_only=True,
                                                 weights_only=True, **options)
                        if pretrained else AutoModel.from_config(config, **options))
        self.pooling, self.hidden_size = pooling, config.hidden_size
        width = 3*self.hidden_size+7
        self.classifier = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 256), nn.GELU(),
                                        nn.Dropout(dropout), nn.Linear(256, 2))

    def forward(self, input_ids, attention_mask, observed):
        hidden = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        if self.pooling == 'mean':
            mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
            states = (hidden*mask).sum(dim=1)/mask.sum(dim=1).clamp(min=1)
        else:
            states = hidden[:, 0]
        pairs = states.reshape(observed.shape[0], 2, self.hidden_size)
        current = pairs[:, 0]
        recent = pairs[:, 1]*observed[:, -1:]
        return self.classifier(torch.cat((current, recent, current*recent, observed), dim=-1))


def tokenizer_at(path):
    return AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False,
                                         fix_mistral_regex=False)


def load(path, metadata, device='cpu'):
    path = Path(path)
    config = AutoConfig.from_pretrained(path/'encoder', local_files_only=True, trust_remote_code=False)
    model = VoiceEncoder(config, pooling=metadata['pooling'])
    model.load_state_dict(load_file(str(path/'model.safetensors')), strict=True)
    return model.to(device).eval(), tokenizer_at(path/'tokenizer')


def predict(model, tokenizer, rows, *, device='cpu', batch_size=16, max_length=384, drop_audio=False,
            history_policy='first80'):
    model.eval()
    logits = []
    with torch.inference_mode():
        for start in range(0, len(rows), batch_size):
            inputs = collate(tokenizer, rows[start:start+batch_size], max_length, drop_audio=drop_audio,
                             history_policy=history_policy)
            values = model(**{key: value.to(device) for key, value in inputs.items()})
            logits.extend(values.float().cpu().tolist())
    return logits