import os
import sys

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

from .configuration_fela import FelaConfig
from .cpu_delta import CPUGatedDeltaNet as _cd
from .cpu_landmark import CPULandmark as _cl
from .cpu_swa import CPUSlidingWindow as _cs
from .model_cpu_gpt2 import CPUGPT, CPUGPTConfig
from .cpu_patch import enable_cpu_delta

_KEEP = (_cd, _cl, _cs)


class FelaForCausalLM(PreTrainedModel):
    config_class = FelaConfig
    base_model_prefix = "model"
    _tied_weights_keys = []
    _no_split_modules = []

    def __init__(self, config):
        super().__init__(config)
        cfg = CPUGPTConfig(
            vocab_size=config.vocab_size,
            seq_len=config.seq_len,
            n_layer=config.n_layer,
            n_embd=config.n_embd,
            n_head=config.n_head,
            ffn_hidden=config.ffn_hidden,
            layer_pattern=config.layer_pattern,
            gla_delta=config.gla_delta,
            fno_modes=config.fno_modes,
            gla_chunk=config.gla_chunk,
            landmark_layer_every=config.landmark_layer_every,
            landmark_chunk=config.landmark_chunk,
            landmark_max=config.landmark_max,
            attn_layer_every=config.attn_layer_every,
            dropout=0.0,
        )
        self.model = CPUGPT(cfg)
        self.model.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self._prepared = False
        self.post_init()

    def get_input_embeddings(self):
        return self.model.wte

    def set_input_embeddings(self, value):
        self.model.wte = value

    def get_output_embeddings(self):
        return self.model.lm_head

    def set_output_embeddings(self, value):
        self.model.lm_head = value

    def _ensure_prepared(self):
        if not self._prepared:
            enable_cpu_delta(self.model)
            self.model.prepare_inference()
            self._prepared = True

    def forward(
        self, input_ids=None, attention_mask=None, labels=None, use_cache=None, **kwargs
    ):
        self._ensure_prepared()
        logits = self.model(input_ids)
        loss = None
        if labels is not None:
            sl = logits[..., :-1, :].contiguous()
            lb = labels[..., 1:].contiguous()
            loss = nn.functional.cross_entropy(sl.view(-1, sl.size(-1)), lb.view(-1))
        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

    def can_generate(self):
        return True