import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_asr import FelaAsrConfig
from .fela_ctc2 import FELACTC2
from .model_cpu_gpt2 import CPUGPTConfig


class FelaAsrModel(PreTrainedModel):
    config_class = FelaAsrConfig
    base_model_prefix = "model"
    main_input_name = "wav"

    def __init__(self, config):
        super().__init__(config)
        cfg = CPUGPTConfig(
            vocab_size=config.vocab,
            seq_len=config.seq_len,
            n_layer=config.n_layer,
            n_embd=config.n_embd,
            n_head=config.n_head,
            fno_modes=config.fno_modes,
            gla_chunk=config.gla_chunk,
            ffn_hidden=config.ffn_hidden,
            layer_pattern=config.layer_pattern,
            dropout=0.0,
        )
        if hasattr(cfg, "gla_delta"):
            cfg.gla_delta = bool(config.gla_delta)
        with torch.device("cpu"):
            self.model = FELACTC2(cfg, vocab=config.vocab, n_mels=config.n_mels)
        self.post_init()

    def forward(self, wav=None, input_values=None, **kwargs):
        if wav is None:
            wav = input_values
        out = self.model(wav)
        return CausalLMOutput(logits=out)