# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "marimo>=0.19.0",
#     "pyzmq",
# ]
# ///

# NOTE: Use the command palette (cmd+shift+p) and search "marimo start in editor" to generate the live preview.

import marimo

__generated_with = "0.19.6"
app = marimo.App()


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    Compare sliding-window HMM vs online Bayesian model for trait estimation.

    Focus: How do trait estimates (warmth, competence, effort) evolve over time?

    - Bayesian: Accumulates all evidence, converges to global average
    - HMM: Uses recent window, tracks local behavioral "mode"
    """)
    return


@app.cell(hide_code=True)
def _():
    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from collections import Counter, defaultdict
    import warnings
    warnings.filterwarnings('ignore')

    from hmmlearn import hmm
    import networkx as nx
    return Counter, defaultdict, hmm, mo, mpatches, np, nx, plt


@app.cell(hide_code=True)
def _():
    TOKENS = ["try_success", "try_fail", "help", "harm", "wait"]
    TOKEN_COLORS = {
        "try_success": "#2ecc71",
        "try_fail": "#e74c3c", 
        "help": "#3498db",
        "harm": "#9b59b6",
        "wait": "#95a5a6"
    }

    TRAIT_COLORS = {
        "effort": "#e67e22",
        "competence": "#27ae60", 
        "warmth": "#2980b9"
    }

    mode_colors = {
        "prosocial": "violet",
        "antisocial": "orange", 
        "withdrawn": "black"
    }

    outcome_colors = {
        "success": "#2ecc71",
        "fail": "#e74c3c",
        "help": "#3498db",
        "harm": "#9b59b6",
        "wait": "#95a5a6",
        "try": "#e67e22"
    }
    return TOKENS, TOKEN_COLORS, mode_colors, outcome_colors


@app.cell(hide_code=True)
def _(np):
    from dataclasses import dataclass, field
    from abc import ABC, abstractmethod

    @dataclass
    class BaseAgent(ABC):
        """Base class for behavioral agents with mode-based action generation."""
        modes: dict = field(default_factory=dict)
        competence: dict = field(default_factory=dict)
        transitions: dict = field(default_factory=dict)
        current_mode: str = "prosocial"
        initial_mode: str = "prosocial"

        def reset(self):
            self.current_mode = self.initial_mode

        def generate_token(self):
            probs = self.modes[self.current_mode]
            actions = list(probs.keys())
            p = list(probs.values())
            action = np.random.choice(actions, p=p)

            if action == "try":
                comp = self.competence[self.current_mode]
                outcome = "success" if np.random.random() < comp else "fail"
                token = f"try_{outcome}"
                self._transition(outcome)
            else:
                token = action
                self._transition(action)

            return token, self.current_mode

        def _transition(self, outcome):
            key = (self.current_mode, outcome)
            if key in self.transitions:
                trans = self.transitions[key]
                modes = list(trans.keys())
                probs = list(trans.values())
                self.current_mode = np.random.choice(modes, p=probs)


    @dataclass
    class VolatileAgent(BaseAgent):
        """Agent with distinct prosocial/antisocial/withdrawn phases."""

        def __post_init__(self):
            self.modes = {
                "prosocial": {"try": 0.4, "help": 0.5, "harm": 0.0, "wait": 0.1},
                "antisocial": {"try": 0.2, "help": 0.0, "harm": 0.6, "wait": 0.2},
                "withdrawn": {"try": 0.1, "help": 0.1, "harm": 0.1, "wait": 0.7},
            }
            self.competence = {"prosocial": 0.7, "antisocial": 0.7, "withdrawn": 0.4}
            self.transitions = {
                ("prosocial", "success"): {"prosocial": 0.85, "antisocial": 0.05, "withdrawn": 0.1},
                ("prosocial", "fail"): {"prosocial": 0.6, "antisocial": 0.2, "withdrawn": 0.2},
                ("prosocial", "help"): {"prosocial": 0.9, "antisocial": 0.0, "withdrawn": 0.1},
                ("prosocial", "harm"): {"prosocial": 0.35, "antisocial": 0.45, "withdrawn": 0.2},
                ("prosocial", "wait"): {"prosocial": 0.6, "antisocial": 0.1, "withdrawn": 0.3},
                ("antisocial", "success"): {"prosocial": 0.3, "antisocial": 0.6, "withdrawn": 0.1},
                ("antisocial", "fail"): {"prosocial": 0.1, "antisocial": 0.7, "withdrawn": 0.2},
                ("antisocial", "help"): {"prosocial": 0.35, "antisocial": 0.45, "withdrawn": 0.2},
                ("antisocial", "harm"): {"prosocial": 0.0, "antisocial": 0.9, "withdrawn": 0.1},
                ("antisocial", "wait"): {"prosocial": 0.2, "antisocial": 0.5, "withdrawn": 0.3},
                ("withdrawn", "success"): {"prosocial": 0.35, "antisocial": 0.35, "withdrawn": 0.3},
                ("withdrawn", "fail"): {"prosocial": 0.2, "antisocial": 0.2, "withdrawn": 0.6},
                ("withdrawn", "help"): {"prosocial": 0.5, "antisocial": 0.2, "withdrawn": 0.3},
                ("withdrawn", "harm"): {"prosocial": 0.1, "antisocial": 0.5, "withdrawn": 0.4},
                ("withdrawn", "wait"): {"prosocial": 0.2, "antisocial": 0.2, "withdrawn": 0.6},
            }
            self.current_mode = "prosocial"
            self.initial_mode = "prosocial"


    @dataclass
    class LimitedEffortAgent(BaseAgent):
        """Agent that is almost always prosocial but has decreasing effort over time."""
        initial_effort: float = 0.9
        effort_decay: float = 0.01
        min_effort: float = 0.1
        current_effort: float = field(init=False)
        step: int = field(init=False)

        def __post_init__(self):
            self.current_effort = self.initial_effort
            self.step = 0
            self.modes = {
                "prosocial": {"try": 0.5, "help": 0.4, "harm": 0.02, "wait": 0.08},
                "antisocial": {"try": 0.2, "help": 0.0, "harm": 0.6, "wait": 0.2},
                "withdrawn": {"try": 0.1, "help": 0.1, "harm": 0.1, "wait": 0.7},
            }
            self.competence = {"prosocial": 0.75, "antisocial": 0.7, "withdrawn": 0.4}
            self.transitions = {
                ("prosocial", "success"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("prosocial", "fail"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("prosocial", "help"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("prosocial", "harm"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("prosocial", "wait"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("antisocial", "success"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("antisocial", "fail"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("antisocial", "help"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("antisocial", "harm"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("antisocial", "wait"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("withdrawn", "success"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("withdrawn", "fail"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("withdrawn", "help"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("withdrawn", "harm"): {"prosocial": 1.0, "antisocial": 0.0, "withdrawn": 0.0},
                ("withdrawn", "wait"): {"prosocial": .2, "antisocial": 0.0, "withdrawn": 0.8},
            }
            self.current_mode = "prosocial"
            self.initial_mode = "prosocial"

        def reset(self):
            self.current_effort = self.initial_effort
            self.step = 0
            self.current_mode = self.initial_mode

        def generate_token(self):
            self.step += 1
            self.current_effort = max(
                self.min_effort,
                self.initial_effort - self.effort_decay * self.step
            )

            if self.current_effort < .2:
                self.current_mode = "withdrawn"

            # Use base class token generation
            probs = self.modes[self.current_mode]
            actions = list(probs.keys())
            p = list(probs.values())
            action = np.random.choice(actions, p=p)

            if action == "try":
                comp = self.competence[self.current_mode]
                outcome = "success" if np.random.random() < comp else "fail"
                token = f"try_{outcome}"
            else:
                token = action

            return token, self.current_mode
    return LimitedEffortAgent, VolatileAgent


@app.cell
def _(mo):
    agent_selector = mo.ui.dropdown(
        options=["VolatileAgent", "LimitedEffortAgent"],
        value="VolatileAgent",
        label="Current agent",
        full_width=False
    )
    agent_selector
    return (agent_selector,)


@app.cell
def _(LimitedEffortAgent, VolatileAgent, agent_selector):
    if agent_selector.value == "VolatileAgent":
        Agent = VolatileAgent
    elif agent_selector.value == "LimitedEffortAgent":
        Agent = LimitedEffortAgent
    else:
        raise ValueError("Unknown agent type selected.")
    return (Agent,)


@app.class_definition(hide_code=True)
class OnlineBayesian:
    """Bayesian model with conjugate Beta priors for trait estimation."""

    def __init__(self, prior_alpha=1.0, prior_beta=1.0):
        # Effort: P(not wait)
        self.effort_alpha = prior_alpha
        self.effort_beta = prior_beta
        # Competence: P(success | try)
        self.comp_alpha = prior_alpha
        self.comp_beta = prior_beta
        # Warmth components
        self.help_alpha = prior_alpha
        self.help_beta = prior_beta
        self.harm_alpha = prior_alpha
        self.harm_beta = prior_beta

        self.n_obs = 0

    def update(self, token):
        self.n_obs += 1

        # Update effort (not wait vs wait)
        if token == "wait":
            self.effort_beta += 1
        else:
            self.effort_alpha += 1

        # Update competence (success vs fail, conditional on try)
        if token == "try_success":
            self.comp_alpha += 1
        elif token == "try_fail":
            self.comp_beta += 1

        # Update warmth components
        if token == "help":
            self.help_alpha += 1
        else:
            self.help_beta += 1

        if token == "harm":
            self.harm_alpha += 1
        else:
            self.harm_beta += 1

    def get_traits(self):
        effort = self.effort_alpha / (self.effort_alpha + self.effort_beta)
        competence = self.comp_alpha / (self.comp_alpha + self.comp_beta)
        p_help = self.help_alpha / (self.help_alpha + self.help_beta)
        p_harm = self.harm_alpha / (self.harm_alpha + self.harm_beta)
        warmth = p_help - p_harm

        return {"effort": effort, "competence": competence, "warmth": warmth}


@app.cell(hide_code=True)
def _(Counter, TOKENS, hmm, np):
    class SlidingWindowHMM:
        """HMM that uses only recent observations to estimate traits."""

        def __init__(self, window_size=20, n_states=3):
            self.window_size = window_size
            self.n_states = n_states
            self.tokens = TOKENS
            self.n_tokens = len(TOKENS)
            self.token_to_idx = {t: i for i, t in enumerate(TOKENS)}
            self.history = []

        def update(self, token):
            self.history.append(token)

        def get_traits(self):
            """Get traits from HMM trained on recent window."""
            if len(self.history) < 5:
                return {"effort": 0.5, "competence": 0.5, "warmth": 0.0}

            # Use sliding window
            window = self.history[-self.window_size:] if len(self.history) > self.window_size else self.history

            # Fit HMM to window
            obs = np.array([[self.token_to_idx[t]] for t in window])

            model = hmm.CategoricalHMM(
                n_components=self.n_states,
                n_iter=50,
                random_state=42
            )
            model.n_features = self.n_tokens

            try:
                model.fit(obs)

                # Get current state distribution
                _, posteriors = model.score_samples(obs)
                current_state = posteriors[-1]

                # Predict next state and emission
                next_state = current_state @ model.transmat_
                pred = next_state @ model.emissionprob_

            except:
                # Fallback to empirical counts from window
                counts = Counter(window)
                total = len(window)
                pred = np.array([counts.get(t, 0) / total for t in TOKENS])

            # Extract traits from prediction
            p_try = pred[0] + pred[1]  # try_success + try_fail
            p_wait = pred[4]
            p_help = pred[2]
            p_harm = pred[3]

            effort = 1 - p_wait
            competence = pred[0] / p_try if p_try > 0.01 else 0.5
            warmth = p_help - p_harm

            return {"effort": effort, "competence": competence, "warmth": warmth}
    return (SlidingWindowHMM,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## State Machine Visualization (Mode-Centric)

    Nodes are **modes**, edges show transitions after specific outcomes.

    Treating "try" as a composite outcome (success/fail weighted by competence).
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    outcome_selector = mo.ui.dropdown(
        options=["success", "fail", "help", "harm", "wait", "try"],
        value="help",
        label="Outcome to visualize",
        full_width=False
    )

    threshold_slider = mo.ui.slider(
        start=0.0,
        stop=0.5,
        step=0.025,
        value=0.1,
        label="Minimum probability threshold",
        full_width=False
    )

    mo.hstack([outcome_selector, threshold_slider])
    return outcome_selector, threshold_slider


@app.cell(hide_code=True)
def _(outcome_selector, threshold_slider, visualize_state_machine):
    visualize_state_machine(
        outcome_filter=outcome_selector.value,
        threshold=threshold_slider.value
    )
    return


@app.cell(hide_code=True)
def _(Agent, defaultdict, mode_colors, nx, outcome_colors, plt):
    def visualize_state_machine(outcome_filter, threshold=0.1):
        """
        Visualize the Agent's state transitions as a directed graph.

        Edges point FROM current_mode TO next_mode after outcome occurs.
        Example: prosocial --[success]--> prosocial (self-loop with 0.85 probability)

        Args:
            outcome_filter: Which outcome to show ("success", "fail", "help", "harm", "wait", "try")
            threshold: Minimum probability to display (0-1)
        """
        agent = Agent()

        G = nx.DiGraph()

        for mode in ["prosocial", "antisocial", "withdrawn"]:
            G.add_node(mode)

        edge_data = defaultdict(list)

        if outcome_filter == "try":
            # "try" combines success and fail weighted by competence
            for from_mode in ["prosocial", "antisocial", "withdrawn"]:
                comp = agent.competence[from_mode]
                success_trans = agent.transitions[(from_mode, "success")]
                fail_trans = agent.transitions[(from_mode, "fail")]

                for to_mode in ["prosocial", "antisocial", "withdrawn"]:
                    # P(to_mode | try) = P(success) * P(to_mode | success) + P(fail) * P(to_mode | fail)
                    prob = comp * success_trans[to_mode] + (1 - comp) * fail_trans[to_mode]
                    if prob < threshold:
                        continue

                    edge = (from_mode, to_mode)
                    edge_data[edge].append(("try", prob, outcome_colors.get("try", "#7f8c8d"), prob * 5))
        else:
            for (from_mode, outcome), trans_probs in agent.transitions.items():
                if outcome != outcome_filter:
                    continue

                for to_mode, prob in trans_probs.items():
                    if prob < threshold:
                        continue

                    edge = (from_mode, to_mode)
                    edge_data[edge].append((outcome, prob, outcome_colors.get(outcome, "#7f8c8d"), prob * 5))

        for edge in edge_data:
            G.add_edge(*edge)

        fig, ax = plt.subplots(figsize=(14, 10))
        pos = nx.circular_layout(G)

        nx.draw_networkx_nodes(G, pos, node_color=[mode_colors[node] for node in G.nodes()], 
                               node_size=4000, alpha=0.9, ax=ax)

        nx.draw_networkx_labels(G, pos, font_size=12, font_weight='bold', 
                                font_color='white', ax=ax)

        edges_by_style = {}
        connectionstyle = 'arc3,rad=0.25'

        for edge, outcomes_list in edge_data.items():
            avg_width = sum(d[3] for d in outcomes_list) / len(outcomes_list)
            max_idx = max(range(len(outcomes_list)), key=lambda i: outcomes_list[i][1])
            edge_color = outcomes_list[max_idx][2]

            if connectionstyle not in edges_by_style:
                edges_by_style[connectionstyle] = []
            edges_by_style[connectionstyle].append((edge, avg_width, edge_color, outcomes_list))

        for connectionstyle, edge_list in edges_by_style.items():
            nx.draw_networkx_edges(
                G, pos, 
                edgelist=[e[0] for e in edge_list],
                width=[e[1] for e in edge_list], 
                alpha=0.7,
                edge_color=[e[2] for e in edge_list], 
                connectionstyle=connectionstyle,
                arrows=True, 
                arrowsize=30, 
                arrowstyle='-|>',
                node_size=4000,
                ax=ax,
                min_source_margin=20,
                min_target_margin=20
            )

        edge_labels = {}
        for edge, outcomes_list in edge_data.items():
            non_zero = [(outcome, prob) for outcome, prob, _, _ in outcomes_list if prob > 0]
            if non_zero:
                label = "\n".join([f"{prob:.2f}" for outcome, prob in non_zero])
                edge_labels[edge] = label

        label_bbox = dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.9, edgecolor='gray', linewidth=0.5)
        nx.draw_networkx_edge_labels(
            G, pos, 
            edge_labels=edge_labels,
            font_size=8,
            bbox=label_bbox,
            connectionstyle=connectionstyle,
            ax=ax
        )

        title = f"State Machine: '{outcome_filter}' Outcome (threshold={threshold:.3f})"
        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)

        ax.axis('off')
        ax.set_xlim(-1., 1.4)
        ax.set_ylim(-1.4, 1.4)
        plt.tight_layout()

        return fig
    return (visualize_state_machine,)


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Action Transition Graph (Action-Centric)

    Nodes are **actions/outcomes**, edges show action-to-action transitions.
    The mode filter controls which behavioral context to visualize.

    Treating "try" as an intermediate action (leads to success/fail based on competence).
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mode_selector = mo.ui.dropdown(
        options=["prosocial", "antisocial", "withdrawn"],
        value="prosocial",
        label="Mode",
        full_width=False
    )
    mode_selector
    return (mode_selector,)


@app.cell
def _(mode_selector, visualize_action_graph):
    visualize_action_graph(mode_filter=mode_selector.value)
    return


@app.cell(hide_code=True)
def _(Agent, mode_colors, nx, outcome_colors, plt):
    def visualize_action_graph(mode_filter="prosocial"):
        """
        Visualize actions as nodes for a single mode (assuming mode doesn't change).

        Edge weights are P(next_action | mode) - the probability of each action in the mode.
        All outgoing edges from any node have the same weights since mode is fixed.

        Args:
            mode_filter: Which mode to show ("prosocial", "antisocial", "withdrawn")
        """
        agent = Agent()

        mode_actions = agent.modes[mode_filter]
        p_success = agent.competence[mode_filter]

        actions = ["help", "harm", "try", "wait"]
        outcomes = ["success", "fail"]

        # Build edge data: (from, to) -> prob, filtering out zero probabilities
        edge_data = {}

        # try -> success/fail based on competence
        edge_data[("try", "success")] = p_success
        edge_data[("try", "fail")] = 1 - p_success

        # All other transitions: from any action/outcome to the next action
        for from_node in actions + outcomes:
            if from_node == "try":
                continue
            for next_action, prob in mode_actions.items():
                if prob > 0:
                    edge_data[(from_node, next_action)] = prob

        # Build graph with only non-zero edges
        G = nx.DiGraph()
        for action in actions + outcomes:
            G.add_node(action)
        for edge in edge_data:
            G.add_edge(*edge)

        fig, ax = plt.subplots(figsize=(14, 10))
        pos = nx.spring_layout(G, k=2, iterations=50, seed=42)

        nx.draw_networkx_nodes(G, pos, node_color=[outcome_colors[node] for node in G.nodes()],
                               node_size=3000, alpha=0.9, ax=ax)

        nx.draw_networkx_labels(G, pos, font_size=11, font_weight='bold',
                                font_color='white', ax=ax)

        # Draw all edges at once
        edge_color = mode_colors[mode_filter]
        connectionstyle = 'arc3,rad=0.2'

        nx.draw_networkx_edges(
            G, pos,
            edgelist=list(edge_data.keys()),
            width=[5 * edge_data[e] for e in edge_data],
            alpha=0.7,
            edge_color=edge_color,
            connectionstyle=connectionstyle,
            arrows=True,
            arrowsize=20,
            arrowstyle='-|>',
            node_size=3000,
            ax=ax
        )

        # Edge labels for non-zero edges only
        edge_labels = {e: f"{p:.2f}" for e, p in edge_data.items()}

        label_bbox = dict(boxstyle='round,pad=0.2', facecolor='white', alpha=0.8, edgecolor='gray', linewidth=0.5)
        nx.draw_networkx_edge_labels(
            G, pos,
            edge_labels=edge_labels,
            font_size=9,
            bbox=label_bbox,
            connectionstyle=connectionstyle,
            ax=ax
        )

        title = f"Action Transition Graph: {mode_filter.capitalize()} Mode"
        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)

        ax.axis('off')
        ax.set_xlim(-1.4, 1.4)
        ax.set_ylim(-1.4, 1.4)
        plt.tight_layout()

        return fig
    return (visualize_action_graph,)


@app.cell(hide_code=True)
def _(plot_comparison, results):
    # Plot
    plot_comparison(results)
    return


@app.cell(hide_code=True)
def _(Counter, results):
    # Count tokens
    counts = Counter(results['tokens'])
    print("\nSequence statistics:")
    for t, c in counts.most_common():
        print(f"  {t}: {c}")

    # Count modes
    mode_counts = Counter(results['true_modes'])
    print("\nTrue mode frequencies:")
    for m, c in mode_counts.most_common():
        print(f"  {m}: {c} ({c/len(results['true_modes'])*100:.1f}%)")
    return


@app.cell(hide_code=True)
def _(np, run_comparison):
    np.random.seed(4)
    results = run_comparison(seq_length=150, window_size=25)
    return (results,)


@app.cell
def _(Agent, Counter, SlidingWindowHMM, TOKEN_COLORS, mpatches, np, plt):
    def run_comparison(seq_length=150, window_size=20):    
        # Generate sequence
        agent = Agent()
        agent.reset()

        tokens = []
        true_modes = []
        for _ in range(seq_length):
            token, mode = agent.generate_token()
            tokens.append(token)
            true_modes.append(mode)

        # Initialize models
        bayesian = OnlineBayesian()
        hmm_model = SlidingWindowHMM(window_size=window_size, n_states=3)

        # Track estimates over time
        bayes_effort, bayes_comp, bayes_warmth = [], [], []
        hmm_effort, hmm_comp, hmm_warmth = [], [], []

        for token in tokens:
            bayesian.update(token)
            hmm_model.update(token)

            b_traits = bayesian.get_traits()
            h_traits = hmm_model.get_traits()

            bayes_effort.append(b_traits['effort'])
            bayes_comp.append(b_traits['competence'])
            bayes_warmth.append(b_traits['warmth'])

            hmm_effort.append(h_traits['effort'])
            hmm_comp.append(h_traits['competence'])
            hmm_warmth.append(h_traits['warmth'])

        # Compute "true" local traits (from sliding window empirical counts)
        true_effort, true_comp, true_warmth = [], [], []
        for i in range(len(tokens)):
            window = tokens[max(0, i-window_size+1):i+1]
            counts = Counter(window)
            n = len(window)

            n_wait = counts.get("wait", 0)
            n_try_s = counts.get("try_success", 0)
            n_try_f = counts.get("try_fail", 0)
            n_help = counts.get("help", 0)
            n_harm = counts.get("harm", 0)

            true_effort.append(1 - n_wait/n)
            true_comp.append(n_try_s / (n_try_s + n_try_f) if (n_try_s + n_try_f) > 0 else 0.5)
            true_warmth.append((n_help - n_harm) / n)

        return {
            'tokens': tokens,
            'true_modes': true_modes,
            'bayes': {'effort': bayes_effort, 'competence': bayes_comp, 'warmth': bayes_warmth},
            'hmm': {'effort': hmm_effort, 'competence': hmm_comp, 'warmth': hmm_warmth},
            'local_truth': {'effort': true_effort, 'competence': true_comp, 'warmth': true_warmth},
            'window_size': window_size
        }


    def plot_comparison(results):
        tokens = results['tokens']
        true_modes = results['true_modes']
        bayes = results['bayes']
        hmm = results['hmm']
        local = results['local_truth']
        window_size = results['window_size']

        fig, axes = plt.subplots(5, 1, figsize=(16, 13), sharex=True,
                                 gridspec_kw={'height_ratios': [0.6, 0.6, 1, 1, 1]})
        steps = np.arange(1, len(tokens) + 1)

        # Panel 1: Action sequence
        ax = axes[0]
        for i, token in enumerate(tokens):
            color = TOKEN_COLORS.get(token, "#bdc3c7")
            ax.barh(0, 1, left=i, color=color, edgecolor='none')
        ax.set_xlim(0, len(tokens))
        ax.set_ylim(-0.5, 0.5)
        ax.set_yticks([])
        ax.set_title("Observed Action Sequence", fontsize=11, fontweight='bold')
        patches = [mpatches.Patch(color=c, label=t) for t, c in TOKEN_COLORS.items()]
        ax.legend(handles=patches, loc='upper right', ncol=5, fontsize=8)
        # Add small markers on the action sequence for each "try" event
        try_indices = [i for i, token in enumerate(tokens) if token.startswith("try_")]
        ax.scatter(try_indices, [0.] * len(try_indices), 
                  marker='|', s=50, color='#e67e22', alpha=1., zorder=10, clip_on=False)

        # Panel 2: True mode
        ax = axes[1]
        mode_colors = {"prosocial": "#2ecc71", "antisocial": "#e74c3c", "withdrawn": "#95a5a6"}
        for i, mode in enumerate(true_modes):
            ax.barh(0, 1, left=i, color=mode_colors[mode], edgecolor='none')
        ax.set_xlim(0, len(tokens))
        ax.set_ylim(-0.5, 0.5)
        ax.set_yticks([])
        ax.set_title("Agent's TRUE Internal Mode (Hidden from Observer)", fontsize=11, fontweight='bold')
        patches = [mpatches.Patch(color=c, label=m) for m, c in mode_colors.items()]
        ax.legend(handles=patches, loc='upper right', ncol=3, fontsize=9)

        # Apply light smoothing to HMM for readability
        def smooth(x, w=3):
            return np.convolve(x, np.ones(w)/w, mode='same')

        hmm_effort_smooth = smooth(hmm['effort'], 5)
        hmm_comp_smooth = smooth(hmm['competence'], 5)
        hmm_warmth_smooth = smooth(hmm['warmth'], 5)

        # Panel 3: Effort
        ax = axes[2]
        ax.fill_between(steps, local['effort'], alpha=0.15, color='gray')
        ax.plot(steps, local['effort'], color='gray', linewidth=1, alpha=0.5, label=f'Local empirical (window={window_size})')
        ax.plot(steps, bayes['effort'], color='#c0392b', linewidth=2.5, label='Bayesian (cumulative)')
        ax.plot(steps, hmm_effort_smooth, color='#2980b9', linewidth=2.5, label=f'HMM (window={window_size})')
        ax.set_ylim(0.2, 1.05)
        ax.set_ylabel("Effort", fontsize=11)
        ax.set_title("EFFORT: 1 - P(wait)", 
                     fontsize=11, fontweight='bold')
        ax.legend(loc='lower right', fontsize=9)
        ax.grid(True, alpha=0.3)

        # Panel 4: Competence
        ax = axes[3]
        ax.fill_between(steps, local['competence'], alpha=0.15, color='gray')
        ax.plot(steps, local['competence'], color='gray', linewidth=1, alpha=0.5, label='Local empirical')
        ax.plot(steps, bayes['competence'], color='#c0392b', linewidth=2.5, label='Bayesian')
        ax.plot(steps, hmm_comp_smooth, color='#2980b9', linewidth=2.5, label='HMM')
        ax.set_ylim(0, 1.05)
        ax.set_ylabel("Competence", fontsize=11)
        ax.set_title("COMPETENCE: P(success | try)", 
                     fontsize=11, fontweight='bold')
        ax.legend(loc='lower right', fontsize=9)
        ax.axhline(0.5, color='gray', linestyle='--', alpha=0.3)
        ax.grid(True, alpha=0.3)

        # Panel 5: Warmth - most diagnostic!
        ax = axes[4]
        ax.fill_between(steps, local['warmth'], alpha=0.15, color='gray')
        ax.plot(steps, local['warmth'], color='gray', linewidth=1, alpha=0.5, label='Local empirical')
        ax.plot(steps, bayes['warmth'], color='#c0392b', linewidth=2.5, label='Bayesian')
        ax.plot(steps, hmm_warmth_smooth, color='#2980b9', linewidth=2.5, label='HMM')
        ax.set_ylim(-0.7, 0.7)
        ax.set_ylabel("Warmth", fontsize=11)
        ax.set_xlabel("Time step", fontsize=11)
        ax.set_title("WARMTH: P(help) - P(harm) ", 
                     fontsize=11, fontweight='bold')
        ax.legend(loc='lower right', fontsize=9)
        ax.axhline(0, color='gray', linestyle='--', alpha=0.5)
        ax.grid(True, alpha=0.3)

        # Add shaded regions for mode phases
        for ax in axes[2:]:
            for i, mode in enumerate(true_modes):
                if mode == "antisocial":
                    ax.axvspan(i, i+1, alpha=0.08, color='red', linewidth=0)
                elif mode == "withdrawn":
                    ax.axvspan(i, i+1, alpha=0.08, color='gray', linewidth=0)

        fig.suptitle("Sliding Window HMM vs Online Bayesian Model: Tracking Behavioral Phases", 
                     fontsize=13, fontweight='bold', y=0.99)
        plt.tight_layout()

        return fig
    return plot_comparison, run_comparison


if __name__ == "__main__":
    app.run()
