import warnings
warnings.filterwarnings("ignore")

import os
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torchaudio
from captum.attr import Saliency
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, pipeline, logging as hf_logging

hf_logging.set_verbosity_error()

MODEL_ID = "matheetharanadshyan/ipd-submission-model"
LLM_ID = "Qwen/Qwen2.5-0.5B-Instruct"
SAMPLE_RATE = 16000

PRACTICE_PHRASES = {
    "Simple Words": [
        "Hello", "Carrot", "Apple", "Water", "Happy", "Good Morning"
    ],
}

ALL_PHRASES = []
for category, phrases in PRACTICE_PHRASES.items():
    ALL_PHRASES.extend(phrases)
ALL_PHRASES.append("Custom Phrase...")

device = torch.device(
    "mps" if torch.backends.mps.is_available() 
    else "cuda" if torch.cuda.is_available() 
    else "cpu"
)

print(f"Loading Models On {device}...")
processor = Wav2Vec2Processor.from_pretrained(MODEL_ID)
model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID).to(device).eval()
llm_pipe = pipeline("text-generation", model=LLM_ID, device=device)
print("Models Loaded Successfully!")


def preprocess_audio(audio_path):
    speech, sr = torchaudio.load(audio_path)
    
    if speech.shape[0] > 1:
        speech = torch.mean(speech, dim=0, keepdim=True)
    
    if sr != SAMPLE_RATE:
        speech = torchaudio.transforms.Resample(sr, SAMPLE_RATE)(speech)
    
    speech = trim_silence(speech)
    speech = speech / (torch.max(torch.abs(speech)) + 1e-8)
    speech = torch.cat([speech[:, :1], speech[:, 1:] - 0.97 * speech[:, :-1]], dim=1)
    
    return speech


def trim_silence(audio, threshold=0.01, frame_length=1024):
    audio_abs = torch.abs(audio.squeeze())
    start_idx, end_idx = 0, len(audio_abs)
    
    for i in range(0, len(audio_abs) - frame_length, frame_length // 2):
        if torch.mean(audio_abs[i:i+frame_length]) > threshold:
            start_idx = max(0, i - frame_length)
            break
    
    for i in range(len(audio_abs) - frame_length, 0, -frame_length // 2):
        if torch.mean(audio_abs[i:i+frame_length]) > threshold:
            end_idx = min(len(audio_abs), i + frame_length * 2)
            break
    
    return audio[:, start_idx:end_idx]


def get_alignment(target, actual):
    n, m = len(target), len(actual)
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    
    for i in range(n + 1):
        dp[i][0] = i
    for j in range(m + 1):
        dp[0][j] = j
    
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = 0 if target[i-1] == actual[j-1] else 1
            dp[i][j] = min(dp[i-1][j] + 1, dp[i][j-1] + 1, dp[i-1][j-1] + cost)
    
    align, i, j = [], n, m
    while i > 0 or j > 0:
        if i > 0 and j > 0 and target[i-1] == actual[j-1]:
            align.append((target[i-1], actual[j-1], "✓ Correct"))
            i, j = i - 1, j - 1
        elif i > 0 and j > 0 and dp[i][j] == dp[i-1][j-1] + 1:
            align.append((target[i-1], actual[j-1], "✗ Substitution"))
            i, j = i - 1, j - 1
        elif j > 0 and (i == 0 or dp[i][j] == dp[i][j-1] + 1):
            align.append(("—", actual[j-1], "⊕ Insertion"))
            j -= 1
        else:
            align.append((target[i-1], "—", "⊖ Deletion"))
            i -= 1
    
    return align[::-1]


def calculate_accuracy(alignment):
    total = len(alignment)
    correct = sum(1 for _, _, s in alignment if "Correct" in s)
    subs = sum(1 for _, _, s in alignment if "Substitution" in s)
    ins = sum(1 for _, _, s in alignment if "Insertion" in s)
    dels = sum(1 for _, _, s in alignment if "Deletion" in s)
    
    return {
        "accuracy": (correct / total * 100) if total > 0 else 0,
        "wer": min(((subs + ins + dels) / max(total - ins, 1)) * 100, 100),
        "correct": correct,
        "total": total
    }


def create_saliency_plot(grads, transcription=""):
    grads_norm = grads / (np.max(grads) + 1e-8)
    threshold = np.percentile(grads_norm, 85)
    grads_smooth = np.convolve(grads_norm, np.ones(5)/5, mode='same')
    
    fig, ax = plt.subplots(figsize=(10, 3), facecolor='#ffffff')
    time_axis = np.arange(len(grads_norm)) / 50
    
    ax.fill_between(time_axis, grads_smooth, color='#3b82f6', alpha=0.15)
    ax.plot(time_axis, grads_smooth, color='#3b82f6', alpha=0.9, linewidth=2)
    
    important_mask = grads_smooth > threshold
    ax.fill_between(time_axis, grads_smooth, where=important_mask, 
                    color='#3b82f6', alpha=0.4, label='High importance')
    
    ax.set_facecolor('#ffffff')
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.tick_params(colors='#6b7280', which='both', labelsize=9, length=0)
    ax.set_xlabel('Time (seconds)', color='#374151', fontsize=10, fontweight='500')
    ax.set_ylabel('Model Attention', color='#374151', fontsize=10, fontweight='500')
    ax.set_xlim(0, time_axis[-1] if len(time_axis) > 0 else 1)
    ax.set_ylim(0, 1.1)
    ax.grid(True, axis='y', alpha=0.2, linestyle='-', linewidth=0.5, color='#d1d5db')
    
    if transcription:
        ax.set_title(f'Speech Pattern Analysis', fontsize=11, fontweight='600', 
                     color='#111827', pad=10)
    
    plt.tight_layout(pad=1.5)
    return fig


def generate_feedback(target_txt, alignment, metrics):
    errors = []
    for target, actual, status in alignment:
        if "Substitution" in status:
            errors.append(f"said '{actual}' instead of '{target}'")
        elif "Deletion" in status:
            errors.append(f"omitted '{target}'")
        elif "Insertion" in status:
            errors.append(f"added extra '{actual}'")
    
    accuracy = metrics['accuracy']
    if accuracy >= 90:
        score_color = "excellent"
    elif accuracy >= 70:
        score_color = "good"
    elif accuracy >= 50:
        score_color = "improving"
    else:
        score_color = "keep practicing"
    
    score_text = f"Score: {accuracy:.0f}%\n**{metrics['correct']}/{metrics['total']}** words correct — {score_color}!\n\n"
    
    if not errors:
        return score_text + "**Perfect pronunciation!** All words were clearly articulated. Great job!"
    
    prompt = f"""<|im_start|>system
You are a friendly speech coach. Provide ONE specific, actionable tip to improve pronunciation. Focus on mouth position, tongue placement, or breathing. Keep under 40 words. Be encouraging.<|im_end|>
<|im_start|>user
Target phrase: "{target_txt}"
Errors made: {"; ".join(errors[:4])}
Accuracy: {metrics['accuracy']:.0f}%<|im_end|>
<|im_start|>assistant
"""
    
    try:
        response = llm_pipe(prompt, max_new_tokens=70, max_length=None, do_sample=True, temperature=0.7)
        feedback = response[0]['generated_text'].split("assistant\n")[-1].strip()
        if feedback and feedback[-1] not in '.!?':
            feedback = feedback.rsplit('.', 1)[0] + '.' if '.' in feedback else feedback
    except Exception:
        if any("Substitution" in e for _, _, e in alignment):
            feedback = "Try speaking more slowly and focus on each syllable. Exaggerate mouth movements to improve clarity."
        elif any("Deletion" in e for _, _, e in alignment):
            feedback = "Some words were missed. Try pausing briefly between words and speaking at a steady pace."
        else:
            feedback = "Focus on speaking clearly and at a comfortable pace. Record yourself and compare to improve."
    
    return score_text + f"**Tip:** {feedback}"


def get_accuracy_color(accuracy):
    """Return a color based on accuracy level."""
    if accuracy >= 90:
        return "#10b981"
    elif accuracy >= 70:
        return "#f59e0b"
    elif accuracy >= 50:
        return "#f97316"
    else:
        return "#ef4444"


def analyze(audio, selected, custom):
    target_txt = (custom.strip() if selected == "Custom Phrase..." and custom else selected).strip()
    
    if not target_txt or target_txt == "Custom Phrase...":
        return (
            "", 
            None, 
            "**Please select or enter a target phrase** to begin your practice session.",
            None
        )
    
    if not audio:
        return (
            "", 
            None, 
            f"**Ready to analyze:** \"{target_txt}\"\n\nClick the microphone button above to record yourself saying this phrase.",
            None
        )
    
    try:
        speech = preprocess_audio(audio)
        
        if speech.shape[1] < 1600:
            return (
                "", 
                None, 
                "**Recording too short**\n\nPlease record for at least 1 second. Speak the entire phrase clearly.",
                None
            )
        
        inputs = processor(
            speech.squeeze().numpy(), 
            sampling_rate=SAMPLE_RATE, 
            return_tensors="pt"
        ).input_values.to(device)
        inputs.requires_grad = True
        
        with torch.no_grad():
            logits = model(inputs).logits
        
        pred_ids = torch.argmax(logits, dim=-1)
        transcription = processor.batch_decode(pred_ids)[0].lower().strip()
        
        if not transcription:
            return (
                "", 
                None, 
                "**No speech detected**\n\nPlease speak louder and closer to the microphone. Make sure your microphone is working.",
                None
            )
        
        alignment = get_alignment(target_txt.lower().split(), transcription.split())
        metrics = calculate_accuracy(alignment)
        
        df = pd.DataFrame(alignment, columns=["Expected", "Heard", "Status"])
        
        inputs_grad = processor(
            speech.squeeze().numpy(),
            sampling_rate=SAMPLE_RATE,
            return_tensors="pt"
        ).input_values.to(device)
        inputs_grad.requires_grad = True
        
        def forward_fn(x):
            return torch.gather(model(x).logits, 2, pred_ids.unsqueeze(-1)).squeeze(-1).sum(dim=-1)
        
        grads = Saliency(forward_fn).attribute(inputs_grad).abs().squeeze().detach().cpu().numpy()
        fig = create_saliency_plot(grads, transcription)
        
        feedback = generate_feedback(target_txt, alignment, metrics)
        
        return transcription.upper(), fig, feedback, df
    
    except Exception as e:
        import traceback
        traceback.print_exc()
        return (
            "", 
            None, 
            f"**Analysis Error**\n\nSomething went wrong: {str(e)}\n\nPlease try recording again.",
            None
        )

css = """
* { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important; }
footer { visibility: hidden; }

.gradio-container {
    max-width: 640px !important;
    margin: 32px auto !important;
    padding: 32px !important;
}

.app-header { text-align: center; margin-bottom: 32px; }
.app-title { font-size: 1.5rem; font-weight: 600; color: #111; margin: 0 0 4px; }
.app-subtitle { font-size: 0.875rem; color: #666; margin: 0; }

.section-title {
    font-size: 0.7rem; font-weight: 600; color: #888;
    text-transform: uppercase; letter-spacing: 0.05em;
    margin: 24px 0 8px;
}

.analyze-button {
    background: #111 !important; color: #fff !important;
    font-weight: 500 !important; font-size: 0.875rem !important;
    padding: 10px 20px !important; border: none !important;
    border-radius: 6px !important; width: 100% !important;
    margin-top: 16px !important;
}
.analyze-button:hover { background: #333 !important; }

.transcription-output textarea {
    text-align: center !important; font-weight: 600 !important;
    font-size: 1.25rem !important; padding: 16px !important;
    background: #f9f9f9 !important; border: 1px solid #eee !important;
    border-radius: 6px !important;
}

.feedback-output {
    background: #fefce8 !important; border: 1px solid #fde047 !important;
    border-radius: 6px !important; padding: 16px !important;
    line-height: 1.6 !important; font-size: 0.875rem !important;
}

.dataframe-container table { font-size: 0.8rem !important; }
.dataframe-container thead th {
    background: #f5f5f5 !important; font-size: 0.7rem !important;
    text-transform: uppercase; padding: 8px 12px !important;
}
.dataframe-container tbody td { padding: 8px 12px !important; }

.gr-box, .gr-panel, .gr-form, .gr-group, .block, .wrap, .contain {
    background: transparent !important; border: none !important; box-shadow: none !important;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.HTML("""
        <div class="app-header">
            <h1 class="app-title">Vocalis</h1>
            <p class="app-subtitle">An Adaptive Explainable AI Framework For Pronunciation Training</p>
        </div>
    """)
    
    target_drop = gr.Dropdown(
        choices=ALL_PHRASES,
        value="Hello",
        label="Phrase to practice",
        interactive=True
    )
    custom_in = gr.Textbox(
        placeholder="Type your phrase...",
        label="Custom phrase",
        visible=False
    )
    
    audio_in = gr.Audio(
        sources=["microphone", "upload"],
        type="filepath",
        label="Recording"
    )

    run_btn = gr.Button("Analyze", elem_classes="analyze-button")
    
    gr.HTML('<div class="section-title">Results</div>')
    
    out_t = gr.Textbox(
        label="Transcription",
        elem_classes="transcription-output",
        interactive=False
    )
    
    out_f = gr.Markdown(elem_classes="feedback-output")
    
    out_p = gr.Plot(label="Speech Pattern", show_label=True)
    
    out_d = gr.Dataframe(
        label="Word Comparison",
        interactive=False,
        elem_classes="dataframe-container"
    )
    
    target_drop.change(
        lambda x: gr.update(visible=(x == "Custom Phrase...")),
        target_drop, custom_in
    )
    
    run_btn.click(
        fn=analyze,
        inputs=[audio_in, target_drop, custom_in],
        outputs=[out_t, out_p, out_f, out_d]
    )

if __name__ == "__main__":
    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)