from typing import Optional

import mlx.core as mx
import mlx.nn as nn

from ..base import InputEmbeddingsFeatures, LanguageModelOutput
from .config import ModelConfig
from .language import LanguageModel


class Model(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.model_type = config.model_type
        self.language_model = LanguageModel(config)

    def get_input_embeddings(
        self,
        input_ids: Optional[mx.array] = None,
        pixel_values: Optional[mx.array] = None,
        **kwargs,
    ) -> InputEmbeddingsFeatures:
        if pixel_values is not None:
            raise ValueError("LLaDA2 MoE is a text-only model.")
        if input_ids is None:
            raise ValueError("input_ids are required for LLaDA2 MoE.")
        return InputEmbeddingsFeatures(
            inputs_embeds=self.language_model.model.word_embeddings(input_ids)
        )

    def __call__(
        self,
        input_ids: mx.array,
        pixel_values: mx.array = None,
        mask: mx.array = None,
        cache=None,
        **kwargs,
    ) -> LanguageModelOutput:
        input_embeddings_features = self.get_input_embeddings(input_ids, pixel_values)
        return self.language_model(
            input_ids,
            mask=mask,
            cache=cache,
            inputs_embeds=input_embeddings_features.inputs_embeds,
            **kwargs,
        )

    def sanitize(self, weights):
        def transform_key(key):
            if key.startswith("language_model."):
                return key
            if key.startswith("model.") or key.startswith("lm_head."):
                return f"language_model.{key}"
            return key

        weights = {transform_key(k): v for k, v in weights.items()}
        return self.language_model.sanitize(weights)

    @property
    def layers(self):
        return self.language_model.layers

    def make_cache(self):
        return self.language_model.make_cache()