# bert-api/bert_api.py
# Joint Intent + Slot Filling Model — Mehearsal NLP joint (RMIT Parts 5+6 delivery)
#
# Origen: D:\XDev.Projects\MEHEARSAL-NLP-CONTROLLER-RMIT\Part_5_6\bert-api\bert-api-01
#   (versión defensiva/híbrida — la "D" del análisis del doc 40 §3.2, decisión §3.2.1)
#
# Adaptaciones MSL para HF Spaces Docker (cf. doc 34 §R3 + doc 40 §5.2.1):
#   1. import os al inicio
#   2. MODEL_NAME configurable vía env var (default: part-5-6-model)
#   3. Porting de carga LOCAL ("nlp_v5/best_joint_model/") → carga HF Hub
#      via snapshot_download(). El resto de la lógica D queda intacta.
#   4. app.run() con puerto desde env (default 7860) y debug=False
#
# Lógica intacta respecto al fuente D:
#   - JointIntentSlotModel(nn.Module): AutoModel + intent_head + slot_head
#   - INTENT_METADATA con 21 intents (incl. MUTE/SOLO/LOOP genéricos y NOT_A_COMMAND)
#   - extract_entities() completo (params_json computado por intent, no solo numbers[])
#   - Intent upgrade hack (~40 líneas) heredado de v5
#   - Slot-first extraction + regex fallback guiado por INTENT_METADATA.target_type
import os
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import snapshot_download
import torch
import torch.nn as nn
from pathlib import Path
import re
import json

# ── Hugging Face Model ────────────────────────────────────────────────────────
# MODEL_NAME es configurable por env var. HF_TOKEN se lee automáticamente por
# huggingface_hub para autenticar el acceso al modelo privado.
HF_MODEL_NAME = os.environ.get('BERT_MODEL_NAME', 'MuseSceneLab/part-5-6-model')

print(f"Downloading/loading joint model from Hugging Face: {HF_MODEL_NAME}")

MODEL_DIR = Path(
    snapshot_download(
        repo_id=HF_MODEL_NAME,
        repo_type="model"
    )
)

LABEL_MAP_PATH = MODEL_DIR / "label_maps.json"
WEIGHTS_PATH   = MODEL_DIR / "joint_model_weights.pt"

app = Flask(__name__)
CORS(app)

# ── Joint Model Architecture ──────────────────────────────────────────────────
class JointIntentSlotModel(nn.Module):
    def __init__(self, model_name, num_intents, num_slots, dropout=0.1):
        super().__init__()
        self.bert       = AutoModel.from_pretrained(model_name)
        hidden_size     = self.bert.config.hidden_size
        self.dropout    = nn.Dropout(dropout)
        self.intent_head = nn.Linear(hidden_size, num_intents)
        self.slot_head   = nn.Linear(hidden_size, num_slots)

    def forward(self, input_ids, attention_mask):
        outputs         = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = self.dropout(outputs.last_hidden_state)
        cls_output      = sequence_output[:, 0, :]
        intent_logits   = self.intent_head(cls_output)
        slot_logits     = self.slot_head(sequence_output)
        return intent_logits, slot_logits


# ── Load label maps ───────────────────────────────────────────────────────────
print(f"Loading label maps from {LABEL_MAP_PATH}...")
with open(LABEL_MAP_PATH, "r") as f:
    label_maps = json.load(f)

id2intent   = {int(k): v for k, v in label_maps["id2intent"].items()}
id2slot     = {int(k): v for k, v in label_maps["id2slot"].items()}
NUM_INTENTS = label_maps["num_intents"]
NUM_SLOTS   = label_maps["num_slots"]

print(f"  {NUM_INTENTS} intent classes")
print(f"  {NUM_SLOTS} slot classes")

# ── Load tokeniser and model ──────────────────────────────────────────────────
print(f"Loading joint model from {MODEL_DIR}...")
device    = torch.device("cpu")
tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR))

model = JointIntentSlotModel(str(MODEL_DIR), NUM_INTENTS, NUM_SLOTS)
model.load_state_dict(torch.load(str(WEIGHTS_PATH), map_location=device))
model.eval()
print("Joint model loaded successfully!")

# MAX_LEN se define aquí (antes era abajo, junto al endpoint /classify) porque
# el warmup de abajo lo necesita. El endpoint /classify lo sigue usando igual.
MAX_LEN = 64

# ── Warmup inference (doc 40 §11.8) ───────────────────────────────────────────
# Ejecuta una inferencia dummy a top-level (antes de que gunicorn forkee los
# workers) para forzar la compilación de los kernels CPU de PyTorch. Sin esto,
# la primera inferencia real post-fork supera los 300s del --timeout y mata
# el worker (incidente §11.1-11.7).
# Con --preload de gunicorn, este código corre en el master y los workers
# heredan el estado compilado via Copy-On-Write. Sumar ~10-30s al arranque
# del Space a cambio de inferencias warm rápidas (~600-800ms) desde el primer
# request real.
import time
print("Warming up model with dummy inference...")
_t0 = time.time()
_dummy_tokens   = "test warmup".lower().split()
_dummy_encoding = tokenizer(
    _dummy_tokens,
    is_split_into_words=True,
    max_length=MAX_LEN,
    padding="max_length",
    truncation=True,
    return_tensors="pt",
)
with torch.no_grad():
    _ = model(
        _dummy_encoding["input_ids"].to(device),
        _dummy_encoding["attention_mask"].to(device),
    )
print(f"Model warmed up in {time.time() - _t0:.1f}s. Ready for inference.")

# ── Entity word lists ─────────────────────────────────────────────────────────
INSTRUMENTS = [
    "drums", "guitar", "bass", "piano", "synth", "vocals", "strings",
    "voice", "singer", "keys", "keyboard", "horns", "brass",
    "batería", "guitarra", "bajo", "sintetizador", "cuerdas",
    "voz", "voces", "cantante", "teclado",
]

SECTIONS = [
    "intro", "verse", "chorus", "bridge", "outro", "hook",
    "breakdown", "refrain", "interlude",
    "introducción", "introduccion", "verso", "estribillo", "puente", "coro",
]

SPANISH_WORDS = [
    "batería", "guitarra", "bajo", "sintetizador", "cuerdas",
    "voz", "voces", "cantante", "teclado", "sube", "baja",
    "silencia", "introducción", "verso", "estribillo", "puente",
    "coro", "bucle", "compás", "tempo",
]

INTENT_METADATA = {
    "MUTE_INSTRUMENT":        {"requires_target": True,  "target_type": "instrument"},
    "UNMUTE_INSTRUMENT":      {"requires_target": True,  "target_type": "instrument"},
    "SOLO_INSTRUMENT":        {"requires_target": True,  "target_type": "instrument"},
    "MUTE_ALL":               {"requires_target": False, "target_type": None},
    "MUTE":                   {"requires_target": True,  "target_type": "instrument"},
    "SOLO":                   {"requires_target": True,  "target_type": "instrument"},
    "VOLUME_SET_ABSOLUTE":    {"requires_target": True,  "target_type": "instrument"},
    "VOLUME_ADJUST_RELATIVE": {"requires_target": True,  "target_type": "instrument"},
    "TEMPO_SET_ABSOLUTE":     {"requires_target": False, "target_type": None},
    "TEMPO_ADJUST_RELATIVE":  {"requires_target": False, "target_type": None},
    "TEMPO_FACTOR":           {"requires_target": False, "target_type": None},
    "TEMPO_STYLE_MARKING":    {"requires_target": False, "target_type": None},
    "TEMPO_GRADUAL":          {"requires_target": False, "target_type": None},
    "JUMP_TO_BAR":            {"requires_target": False, "target_type": None},
    "JUMP_TO_SECTION":        {"requires_target": True,  "target_type": "section"},
    "JUMP_RELATIVE":          {"requires_target": False, "target_type": None},
    "LOOP_BARS":              {"requires_target": False, "target_type": None},
    "LOOP_SECTION":           {"requires_target": True,  "target_type": "section"},
    "LOOP":                   {"requires_target": True,  "target_type": "section"},
    "STOP_LOOP":              {"requires_target": False, "target_type": None},
    "NOT_A_COMMAND":          {"requires_target": False, "target_type": None},
}


# ── Entity extraction ─────────────────────────────────────────────────────────
def extract_entities(text, intent, slot_results):
    text_lower = text.lower()

    # Detect language
    locale = "es" if any(w in text_lower for w in SPANISH_WORDS) else "en"

    # Try slot results first (from joint model)
    target = None
    for word, slot_label in slot_results.items():
        if "instrument" in slot_label and target is None:
            target = word
        elif "section" in slot_label and target is None:
            target = word

    # Fallback to keyword matching
    if target is None:
        metadata    = INTENT_METADATA.get(intent, {})
        target_type = metadata.get("target_type")
        if target_type == "instrument":
            for inst in INSTRUMENTS:
                if re.search(r"\b" + re.escape(inst) + r"\b", text_lower):
                    target = inst
                    break
        elif target_type == "section":
            for sec in SECTIONS:
                if re.search(r"\b" + re.escape(sec) + r"\b", text_lower):
                    target = sec
                    break
            if not target:
                for inst in INSTRUMENTS:
                    if re.search(r"\b" + re.escape(inst) + r"\b", text_lower):
                        target = inst
                        break

    # Extract numbers
    numbers = re.findall(r"\d+(?:\.\d+)?", text)
    params_json = {}

    if "TEMPO_SET_ABSOLUTE" in intent and numbers:
        params_json["bpm"] = int(float(numbers[0]))

    elif "TEMPO_ADJUST_RELATIVE" in intent:
        params_json["direction"] = "up" if any(
            w in text_lower for w in ["up", "increase", "faster", "sube", "más"]) else "down"
        params_json["bpm_change_percent"] = int(float(numbers[0])) if numbers else 10

    elif "TEMPO_FACTOR" in intent:
        if any(w in text_lower for w in ["double", "twice", "doble"]):
            params_json["multiplier"] = 2.0
        elif any(w in text_lower for w in ["half", "mitad"]):
            params_json["multiplier"] = 0.5
        elif any(w in text_lower for w in ["triple"]):
            params_json["multiplier"] = 3.0
        elif numbers:
            params_json["multiplier"] = float(numbers[0])
        else:
            params_json["multiplier"] = 1.0

    elif "TEMPO_STYLE_MARKING" in intent:
        markings = {
            "grave": "Grave", "largo": "Largo", "lento": "Lento",
            "adagio": "Adagio", "andante": "Andante", "moderato": "Moderato",
            "allegretto": "Allegretto", "allegro": "Allegro",
            "vivace": "Vivace", "presto": "Presto", "prestissimo": "Prestissimo",
        }
        for k, v in markings.items():
            if k in text_lower:
                params_json["style_marking"] = v
                break
        if "style_marking" not in params_json:
            params_json["style_marking"] = ""

    elif "TEMPO_GRADUAL" in intent:
        params_json["direction"] = "up" if any(
            w in text_lower for w in ["accelerando", "accel", "speed up", "faster"]) else "down"
        pct = re.search(r"(\d+\.?\d*)\s*(?:%|percent)", text_lower)
        params_json["bpm_change_percent"] = float(pct.group(1)) if pct else 10
        bars = re.search(r"(?:over|in|during|en|durante)\s*(\d+)\s*(?:bar|bars|compás)", text_lower)
        params_json["bars"] = int(bars.group(1)) if bars else (int(float(numbers[0])) if numbers else 4)

    elif "LOOP_BARS" in intent:
        rng = re.search(r"(\d+)\s*(?:-|to|through|hasta|a)\s*(\d+)", text_lower)
        if rng:
            params_json["start_bar"] = int(rng.group(1))
            params_json["end_bar"]   = int(rng.group(2))
        elif len(numbers) >= 2:
            params_json["start_bar"] = int(float(numbers[0]))
            params_json["end_bar"]   = int(float(numbers[1]))
        else:
            params_json["start_bar"] = None
            params_json["end_bar"]   = None

    elif "LOOP_SECTION" in intent:
        bars = re.search(r"(\d+)\s*(?:bar|bars|compás)", text_lower)
        params_json["bars"]      = int(bars.group(1)) if bars else None
        params_json["start_bar"] = None
        params_json["end_bar"]   = None

    elif "JUMP_TO_BAR" in intent and numbers:
        params_json["to_bar"] = int(float(numbers[0]))

    elif "JUMP_RELATIVE" in intent:
        params_json["direction"] = "up" if any(
            w in text_lower for w in ["ahead", "forward", "next", "adelante"]) else "down"
        bars = re.search(r"(\d+)\s*(?:bar|bars|compás)", text_lower)
        params_json["relative_bars"] = int(float(bars.group(1))) if bars else (
            int(float(numbers[0])) if numbers else 1)

    elif "VOLUME" in intent and numbers:
        params_json["vol"] = int(float(numbers[0]))
        if any(w in text_lower for w in ["up", "increase", "louder", "sube", "aumenta"]):
            params_json["direction"] = "up"
        elif any(w in text_lower for w in ["down", "decrease", "quieter", "baja", "reduce"]):
            params_json["direction"] = "down"
        else:
            params_json["direction"] = "up"

    elif "MUTE_ALL" in intent:
        params_json["mute_all"] = "off" if any(
            w in text_lower for w in ["unmute", "back", "enable", "on", "activar"]) else "on"

    elif "MUTE_INSTRUMENT" in intent:
        params_json["mute"] = "on"

    elif "UNMUTE_INSTRUMENT" in intent:
        params_json["mute"] = "off"

    return target, locale, params_json


# ── Classify endpoint ─────────────────────────────────────────────────────────
# MAX_LEN ya definido arriba (junto al warmup §11.8)

@app.route("/classify", methods=["POST"])
def classify():
    try:
        data      = request.json
        utterance = data.get("utterance", "")

        if not utterance:
            return jsonify({"error": "No utterance provided"}), 400

        tokens = utterance.lower().split()

        encoding = tokenizer(
            tokens,
            is_split_into_words=True,
            max_length=MAX_LEN,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )

        input_ids      = encoding["input_ids"].to(device)
        attention_mask = encoding["attention_mask"].to(device)

        with torch.no_grad():
            intent_logits, slot_logits = model(input_ids, attention_mask)

        # Intent
        predicted_class = torch.argmax(intent_logits, dim=1).item()
        confidence      = torch.softmax(intent_logits, dim=1)[0][predicted_class].item()
        intent_label    = id2intent[predicted_class]

        # Slots
        word_ids     = encoding.word_ids(batch_index=0)
        slot_preds   = torch.argmax(slot_logits, dim=2)[0]
        slot_results = {}
        for j, word_id in enumerate(word_ids):
            if word_id is None or word_id in slot_results:
                continue
            slot_label = id2slot.get(slot_preds[j].item(), "O")
            if slot_label != "O" and word_id < len(tokens):
                slot_results[tokens[word_id]] = slot_label

        # Intent upgrade (heredado de v5 — red de seguridad por si el joint model
        # devuelve un intent genérico cuando hay target específico en la utterance)
        text_lower     = utterance.lower()
        has_section    = any(re.search(r"\b" + re.escape(s) + r"\b", text_lower) for s in SECTIONS)
        has_instrument = any(re.search(r"\b" + re.escape(i) + r"\b", text_lower) for i in INSTRUMENTS)
        has_bar_nav    = any(w in text_lower for w in ["go to", "jump to", "ir a"]) and \
                         any(w in text_lower for w in ["bar", "measure", "compás"])
        has_volume     = any(w in text_lower for w in ["volume", "volumen", "loud", "quiet"])

        if intent_label == "LOOP" and has_section:
            intent_label = "LOOP_SECTION"
        elif intent_label == "MUTE" and has_instrument:
            intent_label = "MUTE_INSTRUMENT"
        elif intent_label == "SOLO" and has_instrument:
            intent_label = "SOLO_INSTRUMENT"
        elif intent_label == "UNMUTE" and has_instrument:
            intent_label = "UNMUTE_INSTRUMENT"
        elif intent_label in ["VOLUME_SET_ABSOLUTE", "VOLUME_ADJUST_RELATIVE"] and has_bar_nav:
            intent_label = "JUMP_TO_BAR"
        elif intent_label == "TEMPO_ADJUST_RELATIVE" and has_volume:
            intent_label = "VOLUME_ADJUST_RELATIVE"

        target, locale, params_json = extract_entities(utterance, intent_label, slot_results)

        return jsonify({
            "intent_label": intent_label,
            "locale":       locale,
            "target":       target,
            "params_json":  params_json,
            "confidence":   round(confidence, 4),
            "model":        "DistilBERT-Joint-v6",
            "slots":        slot_results,
        })

    except Exception as e:
        print(f"Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({
            "intent_label": "UNKNOWN",
            "locale":       "en",
            "target":       None,
            "params_json":  {"error": str(e)},
            "model":        "DistilBERT-Joint-v6",
        }), 500


# ── Health check ──────────────────────────────────────────────────────────────
@app.route("/health", methods=["GET"])
def health():
    return jsonify({
        "status":      "healthy",
        "model":       "DistilBERT-Joint-v6",
        "hf_model":    HF_MODEL_NAME,
        "num_intents": NUM_INTENTS,
        "num_slots":   NUM_SLOTS,
    })


if __name__ == "__main__":
    # Puerto adaptado para HF Spaces (default 7860). Debug off para producción.
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False)