"""MAJEPPA / LLM-JEPA inference wrapper for the HF Space demo.

Loads the Lightning checkpoint, exposes generate(score_midi_path, performer, recording)
that returns the path of a freshly written performance MIDI.
"""
from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Optional

import torch

from src.models.llm_jepa_midi import LLMJEPAMidiModule
from src.models.components.aria_lora import PERFORMER_TYPES, RECORDING_TYPES

log = logging.getLogger(__name__)

DEFAULT_HPARAMS = dict(
    model_name="loubb/aria-medium",
    lora_r=512,
    lora_alpha=32,
    lora_dropout=0.1,
    k_pred=1,
    rope_factor=2.0,
    max_position_embeddings=4096,
    lambda_jepa=8.0,
    gamma_gen=1.0,
    temperature=0.07,
    beta_condition=5.0,
)
BASE_VOCAB = 17727  # Aria base vocab; anything >= is a special token we strip on decode


class MajeppaInference:
    def __init__(self, ckpt_path: str | Path, device: str | torch.device | None = None):
        self.device = torch.device(device) if device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.module = self._load(Path(ckpt_path))
        self.module.eval()
        self.module._reset_aria_inference_buffers()
        log.info("MAJEPPA loaded on %s", self.device)

    def _load(self, ckpt_path: Path) -> LLMJEPAMidiModule:
        ckpt = torch.load(str(ckpt_path), map_location=self.device, weights_only=False)
        if "hyper_parameters" in ckpt:
            module = LLMJEPAMidiModule.load_from_checkpoint(str(ckpt_path), map_location=self.device)
        elif "state_dict" in ckpt:
            module = LLMJEPAMidiModule(**DEFAULT_HPARAMS)
            module.load_state_dict(ckpt["state_dict"])
        else:
            module = LLMJEPAMidiModule(**DEFAULT_HPARAMS)
            module.load_state_dict(ckpt)
        return module.to(self.device)

    @torch.no_grad()
    def generate(
        self,
        score_midi_path: str | Path,
        performer_type: str,
        recording_type: str,
        max_new_tokens: int = 512,
        temperature: float = 0.8,
        top_k: int = 50,
        seed: int | None = None,
    ) -> Optional[str]:
        """Generate a performance MIDI from a score MIDI under the given condition.

        If ``seed`` is given, sets torch RNGs so the sampling is reproducible.
        Returns the path of the written .mid file, or None on failure.
        """
        m = self.module
        score_midi_path = Path(score_midi_path)

        if seed is not None:
            torch.manual_seed(int(seed))
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(int(seed))

        enc = m.tokenizer.encode_from_file(str(score_midi_path), return_tensors="pt")
        score_ids = enc.input_ids.squeeze(0).tolist()

        perf_id = m.special_token_ids.get(
            f"<COND_{performer_type}>", m.special_token_ids["<COND_performer_unk>"]
        )
        rec_id = m.special_token_ids.get(
            f"<COND_{recording_type}>", m.special_token_ids["<COND_recording_unk>"]
        )

        first_note = [score_ids[2]] if len(score_ids) > 2 else []
        prompt = score_ids + [m.eos_id, perf_id, rec_id] + m.pred_token_ids + first_note
        prompt_tensor = torch.tensor([prompt], dtype=torch.long, device=self.device)

        gen_ids = m._generate_manual(
            prompt_tensor,
            max_new_tokens=max_new_tokens,
            prompt_len=len(prompt),
            temperature=temperature,
            top_k=top_k,
        )
        if not gen_ids:
            return None

        clean = [t for t in gen_ids if 0 <= t < BASE_VOCAB and t != m.eos_id]
        if not clean:
            return None

        aria_tok = m.tokenizer._tokenizer
        toks = aria_tok.decode(clean)
        toks = [t for t in toks if t not in (("prefix", "instrument", "piano"), "<S>", "<E>")]
        if not toks:
            return None

        full_seq = [("prefix", "instrument", "piano"), "<S>"] + toks + ["<E>"]
        midi_dict = aria_tok.detokenize(full_seq)
        if not midi_dict.note_msgs:
            return None

        out_dir = Path("outputs"); out_dir.mkdir(exist_ok=True)
        stem = re.sub(r"[^a-zA-Z0-9_-]", "_", score_midi_path.stem)[:60]
        seed_tag = f"__seed{seed}" if seed is not None else ""
        out_path = out_dir / f"{stem}__{performer_type}__{recording_type}{seed_tag}.mid"
        midi_dict.to_midi().save(str(out_path))
        return str(out_path)


PERFORMER_LABELS = {
    "virtuoso":            "Virtuoso",
    "piano_teacher":       "Piano Teacher",
    "child_professional":  "Child (professional)",
    "adult_intermediate":  "Adult Intermediate",
    "adult_beginner":      "Adult Beginner",
    "child_beginner":      "Child Beginner",
}
RECORDING_LABELS = {
    "concert_performance": "Concert Performance",
    "performance":         "Performance",
    "demo_class":          "Demo / Teaching",
    "practice":            "Practice",
    "sight_read":          "Sight-read",
}
# Re-export the sets defined in aria_lora.py (kept here for clarity)
PERFORMER_OPTIONS = list(PERFORMER_LABELS.keys())
RECORDING_OPTIONS = list(RECORDING_LABELS.keys())
assert set(PERFORMER_OPTIONS) <= set(PERFORMER_TYPES)
assert set(RECORDING_OPTIONS) <= set(RECORDING_TYPES)