"""Prepare shared speech observations without importing a model runtime."""
import numpy as np

from studio.core.utils.reply_modes.decision_features import AUDIO_FIELDS, VERSION, audio_vector, texts
from .history import history_view

NETWORK_TYPE = 'voicemem-jev-fusion'


def prepare_row(row, *, drop_audio=False, history_policy='first80'):
    current, recent = texts(row['text'], history_view(row.get('history'), history_policy))
    observed = np.concatenate((audio_vector(None if drop_audio else row.get('audio_state')),
                               np.asarray([bool(recent)], np.float32)))
    return current, recent, observed.astype(np.float32)


def collate(tokenizer, rows, max_length=384, *, drop_audio=False, tensors='pt', history_policy='first80'):
    prepared = [prepare_row(row, drop_audio=drop_audio, history_policy=history_policy) for row in rows]
    pairs = [text for current, recent, _ in prepared for text in (current, recent)]
    encoded = tokenizer(pairs, padding=True, truncation=True, max_length=max_length,
                        return_tensors=tensors)
    values = np.stack([row[2] for row in prepared])
    if tensors == 'pt':
        import torch
        values = torch.from_numpy(values)
    return {'input_ids': encoded['input_ids'], 'attention_mask': encoded['attention_mask'],
            'observed': values}


def contract(history_policy='first80'):
    history_view([], history_policy)
    result = {'version': VERSION, 'audio_fields': [name for name, _ in AUDIO_FIELDS],
            'scales': [scale for _, scale in AUDIO_FIELDS],
            'layout': ['eot', 'silence', 'span', 'eot_present', 'silence_present',
                       'span_present', 'history_present'],
            'text_preprocessing': 'voicemem-decision-text',
            'missing_value': 0, 'present_zero_has_mask': True}
    if history_policy != 'first80':
        result['history_policy'] = history_policy
    return result