# (Paste your full Gradio app.py code here)
!pip install -q gradio soundfile

import gradio as gr
from pathlib import Path

# Use your project brief path from the notebook
PROJECT_BRIEF_LOCAL = "/mnt/data/Project_description.pdf"


def normalize_audio_path(audio_file):
    """Return a filesystem path (or None) from Gradio's audio input."""
    if audio_file is None:
        return None
    if isinstance(audio_file, str):
        return audio_file
    if hasattr(audio_file, "name"):
        return audio_file.name
    if isinstance(audio_file, (tuple, list)) and len(audio_file) > 0:
        # In some cases Gradio gives (filepath, ...other stuff)
        return audio_file[0]
    return str(audio_file)


# Extended cuisine map: starters / main / desserts, separated veg / non-veg
CUISINE_MAP = {
    "kerala": {
        "region": "Kerala",
        "starters_veg": ["Banana Chips", "Vegetable Uzhunnu Vada"],
        "starters_nonveg": ["Kerala Fish Fry", "Chicken 65 (Kerala Style)"],
        "main_veg": ["Appam with Vegetable Stew", "Puttu with Kadala Curry"],
        "main_nonveg": ["Kerala Fish Curry", "Beef Ularthiyathu", "Chicken Stew"],
        "desserts": ["Palada Payasam", "Ada Pradhaman"],
    },

    "hindi": {
        "region": "North India (Hindi-speaking belt)",
        "starters_veg": ["Paneer Tikka", "Veg Samosa"],
        "starters_nonveg": ["Chicken 65", "Mutton Seekh Kabab"],
        "main_veg": ["Paneer Butter Masala", "Dal Makhani", "Aloo Paratha"],
        "main_nonveg": ["Butter Chicken", "Rogan Josh", "Mutton Biryani"],
        "desserts": ["Gulab Jamun", "Kheer"],
    },

    "telugu": {
        "region": "Andhra / Telangana",
        "starters_veg": ["Mirchi Bajji", "Pesarattu"],
        "starters_nonveg": ["Gongura Chicken Fry", "Chepala Vepudu"],
        "main_veg": ["Pappu (dal) with rice", "Pesara Pappu", "Gutti Vankaya"],
        "main_nonveg": ["Andhra Chicken Biryani", "Kodi Pulusu (Chicken Curry)"],
        "desserts": ["Bobbatlu (Puran Poli)", "Pootharekulu"],
    },

    "tamil": {
        "region": "Tamil Nadu",
        "starters_veg": ["Medu Vada", "Sundal"],
        "starters_nonveg": ["Chicken 65 (TN style)", "Fish Fry"],
        "main_veg": ["Idli & Sambar", "Masala Dosa", "Rasam Rice"],
        "main_nonveg": ["Chettinad Chicken", "Fish Curry (Meen Kuzhambu)"],
        "desserts": ["Payasam", "Kesari"],
    },

    "malayalam": {
        "region": "Kerala",
        "starters_veg": ["Banana Chips", "Vegetable Uzhunnu Vada"],
        "starters_nonveg": ["Kerala Fish Fry", "Chicken 65 (Kerala)"],
        "main_veg": ["Appam with Vegetable Stew", "Puttu and Kadala Curry"],
        "main_nonveg": ["Kerala Fish Curry", "Beef Ularthiyathu", "Chicken Stew"],
        "desserts": ["Palada Payasam", "Ada Pradhaman"],
    },

    "kannada": {
        "region": "Karnataka",
        "starters_veg": ["Maddur Vada", "Bonda"],
        "starters_nonveg": ["Kheema Cutlet"],
        "main_veg": ["Bisi Bele Bath", "Ragi Mudde", "Akki Roti"],
        "main_nonveg": ["Mangalorean Chicken Curry", "Neer Dosa with Fish Curry"],
        "desserts": ["Mysore Pak", "Kesari"],
    },

    "bengali": {
        "region": "West Bengal",
        "starters_veg": ["Vegetable Chop"],
        "starters_nonveg": ["Fish Kabiraji", "Kolkata-style Chicken Pakora"],
        "main_veg": ["Shukto", "Cholar Dal with Luchi"],
        "main_nonveg": ["Machher Jhol (Fish Curry)", "Ilish Bhapa (Hilsa)"],
        "desserts": ["Mishti Doi", "Rasgulla"],
    },

    # Jharkhand-specific mapping (use label "Jharkhand" -> "jharkhand")
    "jharkhand": {
        "region": "Jharkhand",
        "starters_veg": ["Dhuska", "Chilka Roti"],
        "starters_nonveg": ["Sohari Chicken Fry", "Kodo Chicken", "Jhal Murg"],
        "main_veg": ["Thekua with Chana Sabzi", "Litti Chokha"],
        "main_nonveg": [
            "Chicken Jhol",
            "Mutton with Kodo/Kutki Millet",
            "Bamboo Shoot Chicken",
        ],
        "desserts": ["Dudh Pitha", "Gur ki Roti"],
    },

    "gujarati": {
        "region": "Gujarat",
        "starters_veg": ["Dhokla", "Khandvi"],
        "starters_nonveg": ["(Typically veg cuisine; pick local non-veg if needed)"],
        "main_veg": ["Undhiyu", "Khichdi", "Thepla"],
        "main_nonveg": ["(Typically veg cuisine; pick local non-veg if needed)"],
        "desserts": ["Basundi", "Shrikhand"],
    },

    # default fallback
    "default": {
        "region": "Unknown / Other",
        "starters_veg": ["Local vegetarian starters"],
        "starters_nonveg": ["Local non-veg starters"],
        "main_veg": ["Local vegetarian mains"],
        "main_nonveg": ["Local non-veg mains"],
        "desserts": ["Local desserts"],
    },
}

# NEW: map HuBERT / MFCC labels (state names etc.) to cuisine keys above
LABEL_ALIAS = {
    # Andhra / Telangana
    "andhra_pradesh": "telugu",
    "ap": "telugu",
    "andhra": "telugu",
    "telugu": "telugu",
    "telangana": "telugu",

    # Kerala
    "kerala": "kerala",
    "malayalam": "malayalam",

    # Tamil Nadu
    "tamil_nadu": "tamil",
    "tamil": "tamil",

    # Karnataka
    "karnataka": "kannada",
    "kannada": "kannada",

    # Jharkhand
    "jharkhand": "jharkhand",
    "jharkhand_state": "jharkhand",

    # Gujarat
    "gujarat": "gujarati",
    "gujarati": "gujarati",

    # West Bengal
    "west_bengal": "bengali",
    "bengal": "bengali",
    "bengali": "bengali",

    # North India / Hindi belt
    "hindi": "hindi",
    "north_india": "hindi",
    "delhi": "hindi",
}


def format_cuisine_output(pred_label: str) -> str:
    """Format cuisine recommendations as a readable multiline string."""
    key = pred_label.lower().strip()
    key = key.replace(" ", "_")           # handle "Andhra Pradesh" → "andhra_pradesh"
    key = LABEL_ALIAS.get(key, key)       # map state-style labels → cuisine keys

    info = CUISINE_MAP.get(key, CUISINE_MAP["default"])

    lines = []
    lines.append(f"Inferred region: {info['region']}")
    lines.append("")
    lines.append("Starters (Veg): " + ", ".join(info["starters_veg"]))
    lines.append("Starters (Non-Veg): " + ", ".join(info["starters_nonveg"]))
    lines.append("")
    lines.append("Main Course (Veg): " + ", ".join(info["main_veg"]))
    lines.append("Main Course (Non-Veg): " + ", ".join(info["main_nonveg"]))
    lines.append("")
    lines.append("Desserts: " + ", ".join(info["desserts"]))
    return "\n".join(lines)


def ui_predict_multi(audio_file, feature_choice, hubert_layer_idx, use_trained_model):
    """Main function used by Gradio UI."""
    audio_path = normalize_audio_path(audio_file)
    if audio_path is None:
        return "No audio provided", ""

    feat = "hubert" if feature_choice == "HuBERT" else "mfcc"

    # If using trained model, pick the correct model info
    if use_trained_model:
        model_info = app_state.get("models", {}).get(feat)
        if not model_info:
            return (
                f"No trained {feat.upper()} model available on server. "
                "Train it or uncheck 'Use trained model'.",
                "",
            )

        clf = model_info["clf"]
        scaler = model_info["scaler"]
        le = model_info["le"]

        if feat == "hubert" and hubert_layer_idx != model_info.get(
            "layer", hubert_layer_idx
        ):
            return (
                f"Note: the HuBERT model was trained on layer "
                f"{model_info.get('layer')}. Set the slider to that layer or re-train.",
                "",
            )
    else:
        # Fallback path: only MFCC fallback is implemented
        if feat != "mfcc":
            return (
                "Fallback training is only available for MFCC. "
                "Choose MFCC or upload HuBERT artifacts.",
                "",
            )

        # Try to train a quick MFCC model using local data/ folder
        wav_paths = []
        if Path("data").exists():
            wav_paths = [str(p) for p in Path("data").rglob("*.wav")]

        model_info = train_fast_mfcc_fallback(wav_paths[:200])
        if model_info is None:
            return (
                "No fallback MFCC model could be trained (not enough data). "
                "Upload model artifacts or provide a data/ folder.",
                "",
            )

        clf = model_info["clf"]
        scaler = model_info["scaler"]
        le = model_info["le"]

    # Extract features
    try:
        if feat == "mfcc":
            x = extract_mfcc_pooled(audio_path)
        else:
            x = get_hubert_layer_embedding(audio_path, layer_idx=hubert_layer_idx)
    except Exception as e:
        return f"Feature extraction failed: {e}", ""

    # Predict
    try:
        Xs = scaler.transform(x.reshape(1, -1))
        pred_idx = clf.predict(Xs)[0]
        pred_label = le.inverse_transform([pred_idx])[0]
    except Exception as e:
        return f"Prediction failed (shape mismatch or model error): {e}", ""

    # Format cuisine recommendations
    cuisine_text = format_cuisine_output(pred_label)
    return pred_label, cuisine_text


# Build Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Accent Detection + Cuisine Recommendation Demo")
    gr.Markdown(f"[📄 Open Project Brief]({PROJECT_BRIEF_LOCAL})")

    with gr.Row():
        audio_in = gr.Audio(
            type="filepath",
            label="Upload audio (.wav/.flac/.mp3)",
        )

        with gr.Column():
            feature_choice = gr.Radio(
                ["MFCC", "HuBERT"],
                value="MFCC",
                label="Feature Type",
            )
            hubert_layer_idx = gr.Slider(
                minimum=0,
                maximum=24,
                value=11,
                step=1,
                label="HuBERT Layer",
            )
            use_trained_model = gr.Checkbox(
                value=True,
                label="Use trained model from notebook",
            )
            btn = gr.Button("Predict & Recommend")

    output_label = gr.Textbox(label="Predicted Accent")
    output_cuisines = gr.Textbox(
        label="Recommended Cuisines (starters, mains, desserts separated)",
        lines=12,
    )

    btn.click(
        ui_predict_multi,
        inputs=[audio_in, feature_choice, hubert_layer_idx, use_trained_model],
        outputs=[output_label, output_cuisines],
    )

demo.launch(share=True)