"""
Llama モデル実装

Meta社の最新Llamaモデル
GQA/RoPE/SwiGLUなど最新アーキテクチャを採用
HuggingFace認証が必要
"""
from typing import List, Tuple

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from .base import BaseLanguageModel, ModelConfig


# Llama 3.2 1B設定
LLAMA_3_2_1B_CONFIG = ModelConfig(
    name="Llama 3.2 1B",
    model_id="meta-llama/Llama-3.2-1B",
    embedding_dim=2048,
    vocab_size=128256,
)

# Llama 3.2 3B設定
LLAMA_3_2_3B_CONFIG = ModelConfig(
    name="Llama 3.2 3B",
    model_id="meta-llama/Llama-3.2-3B",
    embedding_dim=3072,
    vocab_size=128256,
)


class LlamaModel(BaseLanguageModel):
    """
    Llamaモデルの実装

    Meta社の最新Llama 3.2シリーズ。
    GQA/RoPE/SwiGLU採用。HuggingFace認証が必要。
    """

    # 出力ノイズの倍率（学習バイアス破壊用）
    LOGITS_NOISE_SCALE = 10.0

    def load(self) -> None:
        """モデルとトークナイザーをロード"""
        if self._is_loaded:
            return

        try:
            self._model = AutoModelForCausalLM.from_pretrained(
                self._config.model_id,
                torch_dtype="auto",
            )
            self._tokenizer = AutoTokenizer.from_pretrained(self._config.model_id)
            self._model.eval()
            self._is_loaded = True
        except Exception as e:
            raise RuntimeError(
                f"Failed to load model {self._config.model_id}: {e}. "
                "Note: Llama models require HuggingFace authentication. "
                "Run 'huggingface-cli login' first."
            )

    def forward_with_noise(
        self, noise: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        ノイズを入力として順伝播を実行し、出力にもノイズを加算

        Args:
            noise: 入力ノイズテンソル [batch, seq_len, embedding_dim]

        Returns:
            Tuple[logits, corrupted_logits]
        """
        if not self._is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        with torch.no_grad():
            outputs = self._model(inputs_embeds=noise)
            logits = outputs.logits

            # 出力logitsにノイズを加算して学習バイアスを破壊
            logits_noise = (
                torch.randn_like(logits) * logits.std() * self.LOGITS_NOISE_SCALE
            )
            corrupted_logits = logits + logits_noise

        return logits, corrupted_logits

    def decode_indices(self, indices: List[int]) -> List[str]:
        """トークンインデックスをデコード"""
        if not self._is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        return [self._tokenizer.decode([i]) for i in indices]