"""LFG-2 model.

Gemma 4 31B (audio encoder + LM, both frozen at train time) with the audio
`embedding_projection` replaced by the trained `DeepAudioProjector`. The stock
Gemma 4 weights live in the model safetensors; the projector's extra parameters
live in `projector_final.pt`. This subclass loads the base model normally, then
installs the deep projector and loads its weights — so a single
`from_pretrained(..., trust_remote_code=True)` yields the full LFG-2 model.
"""
from __future__ import annotations

import os

import torch
from transformers import Gemma4ForConditionalGeneration

from .configuration_lfg2 import LFG2Config
from .deep_projector import install_deep_projector


class LFG2ForConditionalGeneration(Gemma4ForConditionalGeneration):
    config_class = LFG2Config

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        # Load the base Gemma 4 weights (stock nn.Linear projection) as usual.
        model = super().from_pretrained(
            pretrained_model_name_or_path, *model_args, **kwargs
        )
        # Swap in the trained DeepAudioProjector and load its weights.
        model._install_lfg2_projector(
            pretrained_model_name_or_path, revision=kwargs.get("revision")
        )
        return model

    def _install_lfg2_projector(self, name_or_path, revision=None):
        cfg = self.config
        filename = getattr(cfg, "projector_file", "projector_final.pt")
        path = self._resolve_projector_file(name_or_path, filename, revision)
        ckpt = torch.load(path, map_location="cpu")
        ck_cfg = ckpt.get("config", {})
        hidden = ck_cfg.get("hidden", getattr(cfg, "projector_hidden", 4096))
        mlp_layers = ck_cfg.get(
            "mlp_layers", getattr(cfg, "projector_mlp_layers", 2)
        )
        deep = install_deep_projector(
            self, hidden=hidden, n_hidden_layers=mlp_layers,
            param_dtype=torch.float32,
        )
        deep.load_state_dict(ckpt["state_dict"], strict=True)

    @staticmethod
    def _resolve_projector_file(name_or_path, filename, revision=None):
        local = os.path.join(str(name_or_path), filename)
        if os.path.isfile(local):
            return local
        from huggingface_hub import hf_hub_download
        return hf_hub_download(str(name_or_path), filename, revision=revision)