visualization.py

import gradio as gr
import plotly.graph_objs as go

def extract_data_and_explanation(text):
    speakers_data = {}
    current_speaker = None
    explanation = ""
    for line in text.split('\n'):
        line = line.strip()
        if line.startswith("-----------------------"):
            if current_speaker and explanation:
                speakers_data[current_speaker]["explanation"] = explanation.strip()
            explanation = ""
            current_speaker = None
            continue
        if line.startswith("Speaker"):
            current_speaker = line.strip()
            speakers_data[current_speaker] = {}
        elif ':' in line and current_speaker:
            key, value = line.split(':', 1)
            key = key.strip()
            value = value.strip()
            if key.lower() == "explanation":
                explanation += value + " "
            else:
                try:
                    speakers_data[current_speaker][key] = float(value)
                except ValueError:
                    speakers_data[current_speaker][key] = value
        elif line and current_speaker and not line.startswith("Explanation"):
            explanation += line + " "
    
    if current_speaker and explanation:
        speakers_data[current_speaker]["explanation"] = explanation.strip()

    return speakers_data

def create_bar_chart(data, title):
    if not data:
        return None
    fig = go.Figure(data=[go.Bar(
        x=list(data.keys()), 
        y=list(data.values()),
        text=list(data.values()),
        textposition='auto',
        marker_color=['red', 'green', 'blue', 'yellow', 'purple', 'orange', 'pink', 'cyan', 'magenta', 'brown'][:len(data)]
    )])
    fig.update_layout(title=title, xaxis_title="Traits", yaxis_title="Score")
    fig.update_xaxes(tickangle=45)
    return fig

def create_radar_chart(data, title):
    if not data:
        return None
    values = [data.get('Avoidance', 0), data.get('Self', 0), data.get('Anxiety', 0), data.get('Others', 0)]
    fig = go.Figure(data=go.Scatterpolar(
        r=values,
        theta=['Avoidance', 'Self', 'Anxiety', 'Others'],
        fill='toself'
    ))
    fig.update_layout(
        polar=dict(
            radialaxis=dict(visible=True, range=[0, max(values + [10])])
        ),
        showlegend=False,
        title=title
    )
    return fig

def update_visibility_and_charts(status, exec_time, lang, transcription, attachments, bigfive, personalities):
    outputs = [
        gr.update(value=status, visible=True),
        gr.update(value=exec_time, visible=True),
        gr.update(value=lang, visible=True),
        gr.update(value=transcription, visible=True),
    ]
    
    all_analyses = [
        ("Attachments", attachments),
        ("Big Five", bigfive),
        ("Personalities", personalities)
    ]
    
    all_speakers = set()
    for _, analysis_text in all_analyses:
        all_speakers.update(extract_data_and_explanation(analysis_text).keys())
    
    for speaker_index, speaker in enumerate(sorted(all_speakers)[:3]):  # Limit to 3 speakers
        speaker_outputs = []
        
        for analysis_type, analysis_text in all_analyses:
            speakers_data = extract_data_and_explanation(analysis_text)
            data = speakers_data.get(speaker, {})
            
            if data:
                if analysis_type == "Attachments":
                    chart_data = {k: v for k, v in data.items() if k in ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"] and isinstance(v, (int, float))}
                    speaker_outputs.append(gr.update(value=create_bar_chart(chart_data, f"{analysis_type} Analysis - {speaker}") if chart_data else None, visible=bool(chart_data)))
                    
                    radar_data = {k: v for k, v in data.items() if k in ["Anxiety", "Avoidance", "Self", "Others"] and isinstance(v, (int, float))}
                    speaker_outputs.append(gr.update(value=create_radar_chart(radar_data, f"Anxiety-Avoidance-Self-Others - {speaker}") if radar_data else None, visible=bool(radar_data)))
                else:
                    chart_data = {k: v for k, v in data.items() if k not in ["explanation"] and isinstance(v, (int, float))}
                    speaker_outputs.append(gr.update(value=create_bar_chart(chart_data, f"{analysis_type} Analysis - {speaker}") if chart_data else None, visible=bool(chart_data)))
                    speaker_outputs.append(gr.update(visible=False))  # Placeholder for consistency with Attachments
                
                explanation = data.get("explanation", "No explanation provided.")
                speaker_outputs.append(gr.update(value=explanation, visible=True, label=f"{analysis_type} Explanation - {speaker}"))
            else:
                speaker_outputs.extend([gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)])
        
        outputs.extend(speaker_outputs)
    
    # Hide unused speaker components
    for _ in range(3 - len(all_speakers)):
        outputs.extend([gr.update(visible=False)] * 9)  # 3 analyses * (2 charts + 1 explanation)
    
    return outputs