from transformers import pipeline
import gradio as gr
import matplotlib.pyplot as plt

# device=0 uses GPU if available (Runtime -> Change runtime type -> GPU, then rerun this)
classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None, device=0)
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=0)

lap_numbers = [1, 2, 3, 4, 5, 6]
lap_times = [82.1, 81.8, 82.0, 85.4, 84.9, 82.3]

mood_map = {
    'anger': 'stressed', 'fear': 'stressed', 'disgust': 'stressed', 'surprise': 'stressed',
    'sadness': 'tired',
    'joy': 'calm', 'neutral': 'calm'
}

mood_colors = {
    'stressed': ('#ff4444', '#3a1414'),
    'tired': ('#ffcc00', '#3a3314'),
    'calm': ('#00cc66', '#143a24')
}

emotion_colors = {
    'anger': '#ff4444', 'fear': '#ff8800', 'disgust': '#aa44ff', 'surprise': '#00ccff',
    'sadness': '#4488ff', 'joy': '#00cc66', 'neutral': '#888888'
}

def get_advice(driver_mood, lap_times):
    pace_drop = max(lap_times) - min(lap_times)
    if driver_mood == 'stressed' and pace_drop > 2:
        return "⚠️ Driver is stressed AND losing pace. Consider a radio check-in or box call — this combo often precedes a mistake."
    elif driver_mood == 'stressed':
        return "🟠 Driver sounds stressed but pace is holding. Keep monitoring, no action needed yet."
    elif driver_mood == 'tired':
        return "🟡 Fatigue signs detected. Watch for late braking or missed apexes in the next few laps."
    else:
        return "✅ Driver sounds calm and in control. No intervention needed."

# Build the lap chart ONCE, not on every click - it doesn't change per audio clip
plt.style.use('dark_background')
_lap_fig, _ax = plt.subplots(figsize=(5, 4), dpi=80)
_worst_idx = lap_times.index(max(lap_times))
_ax.plot(lap_numbers, lap_times, marker='o', color='#00d4ff', linewidth=2, markersize=8, zorder=2)
_ax.fill_between(lap_numbers, lap_times, min(lap_times) - 1, color='#00d4ff', alpha=0.1)
_ax.scatter(lap_numbers[_worst_idx], lap_times[_worst_idx], color='#ff4444', s=150, zorder=3, label='Slowest lap')
_ax.set_xlabel("Lap Number", fontsize=11)
_ax.set_ylabel("Lap Time (s)", fontsize=11)
_ax.set_title("Lap Performance", fontsize=13, fontweight='bold')
_ax.grid(True, alpha=0.2)
_ax.legend()
_lap_fig.tight_layout()

def make_emotion_pie(emotion_scores):
    fig, ax = plt.subplots(figsize=(5, 4), dpi=80)
    labels = [e['label'] for e in emotion_scores]
    scores = [e['score'] for e in emotion_scores]
    colors = [emotion_colors.get(l, '#666666') for l in labels]
    ax.pie(scores, labels=labels, colors=colors, autopct='%1.1f%%',
           textprops={'fontsize': 9}, wedgeprops={'edgecolor': '#111', 'linewidth': 1})
    ax.set_title("Emotion Breakdown", fontsize=13, fontweight='bold')
    fig.tight_layout()
    return fig

def mood_badge_html(driver_mood, raw_label, score):
    fg, bg = mood_colors[driver_mood]
    return f"""
    <div style="background:{bg}; border:2px solid {fg}; border-radius:10px; padding:14px 18px; text-align:center;">
        <span style="color:{fg}; font-size:22px; font-weight:bold;">{driver_mood.upper()}</span><br>
        <span style="color:#aaa; font-size:13px;">raw model: {raw_label} · confidence {score:.2f}</span>
    </div>
    """

def analyze_clip_gradio(audio_file):
    transcript = transcriber(audio_file)['text']
    emotion_scores = classifier(transcript)[0]
    top_emotion = max(emotion_scores, key=lambda x: x['score'])
    driver_mood = mood_map.get(top_emotion['label'], 'calm')

    badge = mood_badge_html(driver_mood, top_emotion['label'], top_emotion['score'])
    advice = get_advice(driver_mood, lap_times)
    pie = make_emotion_pie(emotion_scores)
    return transcript, badge, advice, _lap_fig, pie

theme = gr.themes.Monochrome(primary_hue="red", secondary_hue="slate")

with gr.Blocks(title="The Silent Co-Driver", theme=theme) as demo:
    gr.Markdown("# 🏎️ The Silent Co-Driver")
    gr.Markdown("Upload a driver radio clip to detect stress and get race engineer advice.")

    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="Radio Clip")
        analyze_btn = gr.Button("🔍 Analyze", variant="primary")

    gr.Examples(examples=["driver_clip.wav"], inputs=audio_input, label="Try a sample clip")

    with gr.Row():
        with gr.Column():
            transcript_out = gr.Textbox(label="Transcript")
            mood_out = gr.HTML(label="Driver Mood")
            advice_out = gr.Textbox(label="Engineer Advice")
        with gr.Column():
            chart_out = gr.Plot(label="Lap Performance")
            pie_out = gr.Plot(label="Emotion Breakdown")

    analyze_btn.click(analyze_clip_gradio, inputs=audio_input, outputs=[transcript_out, mood_out, advice_out, chart_out, pie_out])

demo.launch(share=True)
import json
import streamlit as st

with open("data.json") as f:
    data = json.load(f)


# ---------- PAGE SETUP ----------
st.set_page_config(page_title="The Silent Co-Driver", page_icon="🏎️", layout="wide")

st.title("🏎️ The Silent Co-Driver")
st.write("Reading driver stress from radio calls.")
st.divider()

# ---------- LOAD SAMPLE DATA (from data.json) ----------
# This lets you demo instantly using pre-made clips before your AI model is fully wired in
with open("data.json") as f:
    sample_clips = json.load(f)

st.subheader("📻 Sample Radio Clips")
st.write("Pick a pre-loaded clip to see the analysis instantly:")

clip_names = [clip["clip"] for clip in sample_clips]
selected_clip_name = st.selectbox("Choose a clip", clip_names)

# Find the selected clip's data
selected_clip = next(c for c in sample_clips if c["clip"] == selected_clip_name)

col1, col2, col3 = st.columns(3)
col1.metric("Lap Number", selected_clip["lap"])
col2.metric("Mood", selected_clip["mood"].upper())
col3.metric("Clip File", selected_clip["clip"])

st.write("**Transcript:**")
st.info(selected_clip["transcript"])

st.divider()

# ---------- UPLOAD YOUR OWN CLIP ----------
st.subheader("🎙️ Or Upload Your Own Radio Clip")

audio_file = st.file_uploader("Upload a .wav or .mp3 file", type=["wav", "mp3"])

if audio_file:
    st.audio(audio_file)  # lets you play the clip on the page

    if st.button("Analyze Clip"):
        with st.spinner("Listening to the radio call..."):
            # Save uploaded file temporarily so the AI model can read it
            with open("temp_audio.wav", "wb") as f:
                f.write(audio_file.read())

            # 👇 This is where your teammate's Hugging Face code plugs in
            # Example (uncomment once the model functions are ready):
            #
            # from transformers import pipeline
            # speech_to_text = pipeline("automatic-speech-recognition", model="openai/whisper-base")
            # emotion_detector = pipeline("audio-classification", model="superb/wav2vec2-base-superb-er")
            #
            # transcript = speech_to_text("temp_audio.wav")["text"]
            # mood = emotion_detector("temp_audio.wav")[0]["label"]

            # Placeholder values until the model is connected
            transcript = "Transcript will appear here once AI model is connected."
            mood = "Unknown"

        st.write("**Transcript:**")
        st.info(transcript)
        st.write("**Detected Mood:**")
        st.warning(mood)

st.divider()

# ---------- STRESS VS LAP TIME CHART ----------
st.subheader("📊 Stress vs Lap Time")

# Replace this with real data once you have it (e.g. from all clips + lap times)
chart_data = {
    "Lap 10": 88, "Lap 11": 89, "Lap 12": 95, "Lap 13": 91, "Lap 14": 90
}
st.line_chart(chart_data)