Update visualization.py

visualization.py

@@ -1,91 +1,74 @@
-import plotly.graph_objects as go
-
-def create_placeholder_chart(title="Error"):
-    fig = go.Figure()
-    fig.add_annotation(text="Error generating chart", xref="paper", yref="paper", showarrow=False, font=dict(size=20, color="red"))
-    fig.update_layout(title=title)
-    return fig
+import plotly.graph_objs as go
+from plotly.subplots import make_subplots
 
 def create_charts(results):
     charts = {}
     explanations = {}
     
-    for i, speaker_data in enumerate(results['speaker_analyses']):
-        charts[i] = {}
-        explanations[i] = {}
+    for speaker_id, data in results.items():
+        charts[speaker_id] = {}
+        explanations[speaker_id] = {}
         
         # Attachment Styles
-        try:
-            attachment_data = speaker_data['attachment_styles']
-            labels = ['Secured', 'Anxious-Preoccupied', 'Dismissive-Avoidant', 'Fearful-Avoidant']
-            values = [attachment_data[label] for label in labels]
-            colors = ['blue', 'orange', 'green', 'red']
-            
-            fig = go.Figure(data=[go.Bar(x=labels, y=values, marker_color=colors)])
-            fig.update_layout(title=f'Speaker {i+1}: Attachment Styles', yaxis_range=[0, 1])
-            
-            charts[i]['attachment'] = fig
-            explanations[i]['attachment'] = attachment_data['Explanation']
-        except Exception as e:
-            charts[i]['attachment'] = create_placeholder_chart(f'Speaker {i+1}: Attachment Styles')
-            explanations[i]['attachment'] = f"Error in attachment data: {str(e)}"
-
-        # Attachment Dimensions (Radar Chart)
-        try:
-            dimensions_data = speaker_data['attachment_styles']
-            labels = ['Self', 'Others', 'Anxiety', 'Avoidance']
-            values = [dimensions_data[label] for label in labels]
-            
-            fig = go.Figure(data=go.Scatterpolar(
-              r=values,
-              theta=labels,
-              fill='toself'
-            ))
-            fig.update_layout(
-              polar=dict(
-                radialaxis=dict(visible=True, range=[0, 10])
-              ),
-              showlegend=False,
-              title=f'Speaker {i+1}: Attachment Dimensions'
-            )
-            
-            charts[i]['dimensions'] = fig
-        except Exception as e:
-            charts[i]['dimensions'] = create_placeholder_chart(f'Speaker {i+1}: Attachment Dimensions')
-
+        attachment_data = data['attachments']
+        labels = ['Secured', 'Anxious-Preoccupied', 'Dismissive-Avoidant', 'Fearful-Avoidant']
+        values = [getattr(attachment_data, 'Secured', 0), 
+                  getattr(attachment_data, 'Anxious_Preoccupied', 0), 
+                  getattr(attachment_data, 'Dismissive_Avoidant', 0), 
+                  getattr(attachment_data, 'Fearful_Avoidant', 0)]
+        colors = ['blue', 'orange', 'green', 'red']
+        
+        fig = go.Figure(data=[go.Bar(x=labels, y=values, marker_color=colors)])
+        fig.update_layout(title=f'{speaker_id}: Attachment Styles', yaxis_range=[0, 1])
+        charts[speaker_id]['attachment'] = fig
+        explanations[speaker_id]['attachment'] = attachment_data.Explanation
+        
         # Big Five Traits
-        try:
-            bigfive_data = speaker_data['big_five_traits']
-            labels = ['Extraversion', 'Agreeableness', 'Conscientiousness', 'Neuroticism', 'Openness']
-            values = [bigfive_data[label] for label in labels]
-            
-            fig = go.Figure(data=[go.Bar(x=labels, y=values, marker_color='rgb(55, 83, 109)')])
-            fig.update_layout(title=f'Speaker {i+1}: Big Five Traits', yaxis_range=[0, 10])
-            
-            charts[i]['bigfive'] = fig
-            explanations[i]['bigfive'] = bigfive_data['Explanation']
-        except Exception as e:
-            charts[i]['bigfive'] = create_placeholder_chart(f'Speaker {i+1}: Big Five Traits')
-            explanations[i]['bigfive'] = f"Error in big five data: {str(e)}"
-
+        bigfive_data = data['bigfive']
+        labels = ['Extraversion', 'Agreeableness', 'Conscientiousness', 'Neuroticism', 'Openness']
+        values = [bigfive_data.Extraversion, bigfive_data.Agreeableness, 
+                  bigfive_data.Conscientiousness, bigfive_data.Neuroticism, bigfive_data.Openness]
+        colors = ['blue', 'green', 'red', 'purple', 'orange']
+        
+        fig = go.Figure(data=[go.Bar(x=labels, y=values, marker_color=colors)])
+        fig.update_layout(title=f'{speaker_id}: Big Five Traits', yaxis_range=[0, 10])
+        charts[speaker_id]['bigfive'] = fig
+        explanations[speaker_id]['bigfive'] = bigfive_data.Explanation
+        
         # Personality Disorders
-        try:
-            personality_data = speaker_data['personality_disorders']
-            labels = [k for k in personality_data.keys() if k != 'Explanation']
-            values = [personality_data[label] for label in labels]
-            
-            fig = go.Figure(data=[go.Bar(x=labels, y=values, marker_color='rgb(26, 118, 255)')])
-            fig.update_layout(
-                title=f'Speaker {i+1}: Personality Disorders',
-                yaxis_range=[0, 5],
-                xaxis_tickangle=-45
-            )
-            
-            charts[i]['personality'] = fig
-            explanations[i]['personality'] = personality_data['Explanation']
-        except Exception as e:
-            charts[i]['personality'] = create_placeholder_chart(f'Speaker {i+1}: Personality Disorders')
-            explanations[i]['personality'] = f"Error in personality data: {str(e)}"
-
-    return charts, explanations
+        personality_data = data['personalities']
+        labels = ['Antisocial', 'Narcissistic', 'Depressed', 'Anxious-Avoidant',
+                  'Obsessive', 'Paranoid', 'Borderline', 'Dependent', 'Schizoid-Schizotypal']
+        values = [personality_data.Antisocial_Psychopathic, personality_data.Narcissistic,
+                  personality_data.Depressed, personality_data.Anxious_Avoidant,
+                  personality_data.Obsessional, personality_data.Paranoid,
+                  personality_data.Borderline_Dysregulated, personality_data.Dependent_Victimized,
+                  personality_data.Schizoid_Schizotypal]
+        colors = ['black', 'orange', 'gray', 'green', 'brown', 'purple', 'red', 'cyan', 'gold']
+        
+        fig = go.Figure(data=[go.Bar(x=labels, y=values, marker_color=colors)])
+        fig.update_layout(title=f'{speaker_id}: Personality Disorders', yaxis_range=[0, 5])
+        charts[speaker_id]['personality'] = fig
+        explanations[speaker_id]['personality'] = personality_data.Explanation
 
+        # Attachment Dimensions (Radar Chart)
+        dimensions_data = data['attachments']
+        labels = ['Self', 'Others', 'Anxiety', 'Avoidance']
+        values = [dimensions_data.Self, dimensions_data.Others, 
+                  dimensions_data.Anxiety, dimensions_data.Avoidance]
+        
+        fig = go.Figure(data=go.Scatterpolar(
+          r=values,
+          theta=labels,
+          fill='toself'
+        ))
+        fig.update_layout(
+          polar=dict(
+            radialaxis=dict(visible=True, range=[0, 10])
+          ),
+          showlegend=False,
+          title=f'{speaker_id}: Attachment Dimensions'
+        )
+        charts[speaker_id]['dimensions'] = fig
+    
+    return charts, explanations