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
            explanation = ""
            continue
        if line.startswith("Speaker"):
            current_speaker = line.split(':')[0].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
            elif '|' in value:
                for part in value.split('|'):
                    part = part.strip()
                    if ':' in part:
                        sub_key, sub_value = part.split(':', 1)
                    else:
                        parts = part.split(None, 1)
                        if len(parts) == 2:
                            sub_key, sub_value = parts
                        else:
                            continue
                    try:
                        speakers_data[current_speaker][sub_key.strip()] = float(sub_value.strip())
                    except ValueError:
                        speakers_data[current_speaker][sub_key.strip()] = sub_value.strip()
            else:
                try:
                    speakers_data[current_speaker][key] = float(value)
                except ValueError:
                    speakers_data[current_speaker][key] = value
        elif line and current_speaker:
            speakers_data[current_speaker][line] = 0
    
    if current_speaker and explanation:
        speakers_data[current_speaker]["explanation"] = explanation

    return speakers_data

def create_bar_chart(data, title):
    fig = go.Figure(data=[go.Bar(
        x=list(data.keys()), 
        y=list(data.values()),
        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")
    return fig

def create_radar_chart(data, title):
    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, attachments, bigfive, personalities):
    outputs = [
        gr.update(value=status, visible=True),
        gr.update(value=exec_time, visible=True),
        gr.update(value=lang, visible=True),
    ]
    
    for analysis_text, analysis_type in [(attachments, "Attachments"), (bigfive, "Big Five"), (personalities, "Personalities")]:
        speakers_data = extract_data_and_explanation(analysis_text)
        for speaker, data in speakers_data.items():
            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"]}
                    if chart_data:
                        fig = create_bar_chart(chart_data, f"{analysis_type} Analysis - {speaker}")
                        outputs.append(gr.update(value=fig, visible=True))
                    else:
                        outputs.append(gr.update(visible=False))
                    
                    radar_data = {k: v for k, v in data.items() if k in ["Anxiety", "Avoidance", "Self", "Others"]}
                    if any(radar_data.values()):
                        radar_fig = create_radar_chart(radar_data, f"Anxiety-Avoidance-Self-Others - {speaker}")
                        outputs.append(gr.update(value=radar_fig, visible=True))
                    else:
                        outputs.append(gr.update(visible=False))
                else:
                    chart_data = {k: v for k, v in data.items() if k not in ["explanation"] and isinstance(v, (int, float))}
                    if chart_data:
                        fig = create_bar_chart(chart_data, f"{analysis_type} Analysis - {speaker}")
                        outputs.append(gr.update(value=fig, visible=True))
                    else:
                        outputs.append(gr.update(visible=False))
                    outputs.append(gr.update(visible=False))  # Placeholder for consistency
                
                explanation = data.get("explanation", "No explanation provided.")
                outputs.append(gr.update(value=explanation, visible=True, label=f"{analysis_type} Explanation - {speaker}"))
            else:
                outputs.extend([gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)])
    
    while len(outputs) < 21:
        outputs.append(gr.update(visible=False))
    
    print("Outputs generated:", outputs)
    return outputs