From 29df6a4bd9357598a8de53184382a6abae363adb Mon Sep 17 00:00:00 2001
From: Luigi Maiorano <luigi.maiorano@qumo.io>
Date: Mon, 2 Feb 2026 18:37:45 +0100
Subject: [PATCH] og traits

---
 plots.py     | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 reference.py | 38 ++++++++++++++++++++++++++++++++
 utils.py     |  1 +
 3 files changed, 101 insertions(+)
 create mode 100644 reference.py

diff --git a/plots.py b/plots.py
index a85cae2..c58c541 100644
--- a/plots.py
+++ b/plots.py
@@ -1039,3 +1039,65 @@ class QualtricsPlotsMixin:
             print(f"Word cloud saved to: {filepath}")
         
         return fig
+    def plot_character_trait_frequency(
+        self,
+        data: pl.LazyFrame | pl.DataFrame | None = None,
+        title: str = "Trait Frequency per Brand Character",
+        x_label: str = "Trait",
+        y_label: str = "Frequency (Times Chosen)",
+        height: int | None = None,
+        width: int | str | None = None,
+    ) -> alt.Chart:
+        """Create a grouped bar plot showing how often each trait is chosen per character.
+        
+        Original request: "I need a bar plot that shows the frequency of the times 
+        each trait is chosen per brand character"
+        
+        Expects data with columns: Character, Trait, Count (as produced by 
+        transform_character_trait_frequency).
+        """
+        df = self._ensure_dataframe(data)
+        
+        # Ensure we have the expected columns
+        required_cols = {'Character', 'Trait', 'Count'}
+        if not required_cols.issubset(set(df.columns)):
+            return alt.Chart(pd.DataFrame({'text': ['Data must have Character, Trait, Count columns']})).mark_text().encode(text='text:N')
+        
+        # Convert to pandas for Altair
+        plot_df = df.to_pandas() if hasattr(df, 'to_pandas') else df
+        
+        # Calculate total per trait for sorting (traits with highest overall frequency first)
+        trait_totals = plot_df.groupby('Trait')['Count'].sum().sort_values(ascending=False)
+        trait_order = trait_totals.index.tolist()
+        
+        # Get unique characters for color mapping
+        characters = plot_df['Character'].unique().tolist()
+        
+        # Interactive legend selection - click to filter
+        selection = alt.selection_point(fields=['Character'], bind='legend')
+        
+        chart = alt.Chart(plot_df).mark_bar().encode(
+            x=alt.X('Trait:N', 
+                    title=x_label, 
+                    sort=trait_order,
+                    axis=alt.Axis(labelAngle=-45, labelLimit=200)),
+            y=alt.Y('Count:Q', title=y_label),
+            xOffset='Character:N',
+            color=alt.Color('Character:N',
+                           scale=alt.Scale(domain=characters,
+                                         range=ColorPalette.CATEGORICAL[:len(characters)]),
+                           legend=alt.Legend(orient='top', direction='horizontal', title='Character')),
+            opacity=alt.condition(selection, alt.value(1), alt.value(0.2)),
+            tooltip=[
+                alt.Tooltip('Character:N', title='Character'),
+                alt.Tooltip('Trait:N', title='Trait'),
+                alt.Tooltip('Count:Q', title='Frequency')
+            ]
+        ).add_params(selection).properties(
+            title=self._process_title(title),
+            width=width or 900,
+            height=height or getattr(self, 'plot_height', 400)
+        )
+        
+        chart = self._save_plot(chart, title)
+        return chart
\ No newline at end of file
diff --git a/reference.py b/reference.py
new file mode 100644
index 0000000..a6ba358
--- /dev/null
+++ b/reference.py
@@ -0,0 +1,38 @@
+ORIGINAL_CHARACTER_TRAITS = {
+    "the_familiar_friend": [
+        "Warm",
+        "Friendly",
+        "Approachable",
+        "Familiar",
+        "Casual",
+        "Appreciative",
+        "Benevolent",
+    ],
+    "the_coach": [
+        "Empowering",
+        "Encouraging",
+        "Caring",
+        "Positive",
+        "Optimistic",
+        "Guiding",
+        "Reassuring",
+    ],
+    "the_personal_assistant": [
+        "Forward-thinking",
+        "Progressive",
+        "Cooperative",
+        "Intentional",
+        "Resourceful",
+        "Attentive",
+        "Adaptive",
+    ],
+    "the_bank_teller": [
+        "Patient",
+        "Grounded",
+        "Down-to-earth",
+        "Stable",
+        "Formal",
+        "Balanced",
+        "Efficient",
+    ]
+}
\ No newline at end of file
diff --git a/utils.py b/utils.py
index 68a668a..64a6501 100644
--- a/utils.py
+++ b/utils.py
@@ -1073,6 +1073,7 @@ class QualtricsSurvey(QualtricsPlotsMixin):
         return self._get_subset(q, QIDs, rename_cols=True), None
 
 
+
 def process_speaking_style_data(
     df: Union[pl.LazyFrame, pl.DataFrame],
     trait_map: dict[str, str]