# /// 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.11"
app = marimo.App(width="full")

with app.setup:
    import numpy as np
    from mcp_hmm.hmm_xiang2023 import (
        generate_windows,
        generate_outcome,
        WindowableHMM,
        DebuggingMonitor,
        plot_results,
        plot_hmm_matrices,
        visualize_agent_contributions_sampdist,
        AgentPair,
        TOKENS,
        TOKEN_COLORS,
        make_null_actions_generator,
        generate_windows_with_actions_generator
    )
    import networkx as nx
    import polars as pl
    from polars import lit, col
    import marimo as mo

    from sklearn.model_selection import KFold, LeaveOneGroupOut

    import seaborn as sns
    from matplotlib import pyplot as plt
    from scipy.optimize import minimize


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    # Cross-validation
    """)
    return


@app.cell(hide_code=True)
def _():
    # discriminability = mo.ui.slider(start=0, stop=1, value=.5, step=.1, label="Discriminability")
    discriminability_noise = mo.ui.slider(start=0, stop=1, value=.2, step=.1, label="Discriminability Noise")
    A_competence = mo.ui.slider(start=0, stop=10, value=4, step=.5, label="A Competence")
    B_competence = mo.ui.slider(start=0, stop=10, value=6, step=.5, label="B Competence")

    mo.vstack([ discriminability_noise, A_competence, B_competence ])
    return A_competence, B_competence, discriminability_noise


@app.cell(hide_code=True)
def _(A_competence, B_competence, discriminability_noise):
    np.random.seed(4)
    pair = AgentPair(discriminability_noise=discriminability_noise.value, A_competence=A_competence.value, B_competence=B_competence.value)

    visualize_agent_contributions_sampdist(pair)
    return (pair,)


@app.cell(hide_code=True)
def _(accuracy_results):
    show_accuracy(accuracy_results)
    return


@app.function(hide_code=True)
def get_window_types(windows):
    _, unique_indices = np.unique(np.apply_along_axis(str, 1, windows), return_inverse=True)
    return unique_indices


@app.cell(hide_code=True)
def _(pair):
    n_windows = 100

    np.random.seed(4)
    fwd_windows = np.array(generate_windows(pair, num_windows=n_windows))
    bwd_windows = np.flip(fwd_windows)  # for the backwards model, we flip each window
    _bwd_windows_types = get_window_types(bwd_windows)

    test_scores = []
    train_scores = []
    test_scores_b = []
    train_scores_b = []
    for i, (train, test) in enumerate(LeaveOneGroupOut().split(bwd_windows, groups=_bwd_windows_types)):
        np.random.seed(i)

        train_windows = bwd_windows[train]
        test_windows = bwd_windows[test]

        bwd_model = WindowableHMM(n_states=10, max_iter=120)
        for win in train_windows:
            bwd_model.add_to_history(win)
        bwd_model.fit()

        print(f"{bwd_model.model.monitor_.iter} iterations out of {bwd_model.model.monitor_.n_iter} max")

        train_scores.append(bwd_model.accuracy(train_windows, backward_learning=True).reshape(-1))
        test_scores.append(bwd_model.accuracy(test_windows, backward_learning=True).reshape(-1))

        train_scores_b.append(bwd_model.accuracy_about_stronger_agent(train_windows, backward_learning=True).reshape(-1))
        test_scores_b.append(bwd_model.accuracy_about_stronger_agent(test_windows, backward_learning=True).reshape(-1))
    return (
        bwd_windows,
        n_windows,
        test_scores,
        test_scores_b,
        train_scores,
        train_scores_b,
    )


@app.function
def tidy_scores(folds):
    fold_idxs = []
    window_idxs = []
    scores = []
    for fold_idx, fold in enumerate(folds):
        for window_idx, score in enumerate(fold):
            fold_idxs.append(fold_idx)
            window_idxs.append(window_idx)
            scores.append(score)
    return pl.DataFrame({
        'fold': fold_idxs,
        'window_idx': window_idxs,
        'score': scores
    })


@app.cell(hide_code=True)
def _(test_scores, test_scores_b, train_scores, train_scores_b):
    accuracy_results = pl.concat([
        tidy_scores(train_scores).with_columns(label=lit('train'), target=lit('both agents')),
        tidy_scores(test_scores).with_columns(label=lit('test'), target=lit('both agents')),
        tidy_scores(train_scores_b).with_columns(label=lit('train'), target=lit('stronger agent')),
        tidy_scores(test_scores_b).with_columns(label=lit('test'), target=lit('stronger agent'))
    ])
    return (accuracy_results,)


@app.function(hide_code=True)
def show_accuracy(acc):
    return mo.vstack([
        sns.catplot(
            acc.to_pandas(),
            kind='bar',
            x='fold',
            hue='label',
            y='score',
            col='target'
        ).map_dataframe(
            sns.stripplot,
            x='fold',
            hue='label',
            y='score',
            dodge='auto',
            alpha=.1
        ).set(ylim=[0, 1]).map(plt.axhline, y=.5, ls='--', c='gray'),
        acc.group_by('label', 'target', maintain_order=True)
            .agg(mean_score = col('score').drop_nans().mean())
            .pivot(values='mean_score', on='target', maintain_order=True)
    ])


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    # Inspecting trajectory of hidden weights during fit
    """)
    return


@app.cell(hide_code=True)
def _(bwd_windows):
    np.random.seed(2)
    _n_states = 10
    # TI = trajectory inspector
    TI_bwd_model = WindowableHMM(n_states=_n_states, max_iter=120, save_props_while_fitting=['startprob_', 'transmat_', 'emissionprob_'], custom_model_args={
        # randomize the transmat priors to try to help more varied solutions arise (randomizing the startprob_prior doesn't seem to have an impacts, and can only initialize all the emission probs to the same value)
        'transmat_prior': np.random.rand(_n_states, _n_states)
    })
    for _win in bwd_windows:
        TI_bwd_model.add_to_history(_win)
    TI_bwd_model.fit()
    return (TI_bwd_model,)


@app.function
def loss_plot(model, frame=0):
    return sns.relplot(data=pl.DataFrame({
        'Iteration': np.arange(1, len(model.convergence_history)),
        '-logP': -model.convergence_history[1:]
    }), x='Iteration', y='-logP', kind='line', height=3).map(plt.axvline, x=frame)


@app.function
def make_heatmap_trajectory_inspector(model):
    _order = order_hiddens_by_emission_probs(model, model.fitting_history[-1])
    # ^ define order based on the final iteration so that we see the trajectory towards that

    matrix_figs = [
        plot_hmm_matrices(
            model.model.startprob_prior_.reshape(model.n_states, 1),
            model.model.transmat_prior_,
            model.model.emissionprob_prior_,
            order=_order
        )
    ] + [
        plot_hmm_matrices(
            _record['startprob_'].reshape(model.n_states, 1),
            _record['transmat_'],
            _record['emissionprob_'],
            order=_order
        )
        for _record in model.fitting_history
    ]

    def show(frame_slider):
        return mo.vstack([
            mo.hstack([ frame_slider, mo.md(f"Frame {frame_slider.value}/{frame_slider.stop} (0 = prior)") ]),
            mo.hstack([ matrix_figs[frame_slider.value], loss_plot(model, frame_slider.value) ])
        ])

    frame_slider = mo.ui.slider(start=0, stop=len(matrix_figs) - 1, step=1, value = len(matrix_figs) - 1, label="Frame")

    return show, frame_slider


@app.cell
def _(TI_bwd_model):
    hmTI, hmTI_frame_slider = make_heatmap_trajectory_inspector(TI_bwd_model)
    return hmTI, hmTI_frame_slider


@app.cell
def _(hmTI, hmTI_frame_slider):
    hmTI(hmTI_frame_slider)
    return


@app.function(hide_code=True)
def order_hiddens_by_emission_probs(model: WindowableHMM, reference_record: dict):
    """
    Example usage with plot_hmm_matrices:
    _record = TI_bwd_model.fitting_history[-1]; plot_hmm_matrices(_record['startprob_'].reshape(TI_bwd_model.n_states, 1), _record['transmat_'], _record['emissionprob_'], order=order_hiddens_by_emission_probs(TI_bwd_model, _record))
    """
    from functools import cmp_to_key

    def compare_by_emission_probs(hidden_a, hidden_b):
        emission_idxs = np.arange(0, model.model.n_features)
        hidden_a_emissions = reference_record['emissionprob_'][hidden_a]
        hidden_b_emissions = reference_record['emissionprob_'][hidden_b]
        correlation_a = np.correlate(hidden_a_emissions, emission_idxs)
        correlation_b = np.correlate(hidden_b_emissions, emission_idxs)

        # if np.std(hidden_a_emissions) < .001 or np.std(hidden_b_emissions) < .001:
        #     return 1
        # ^ doesn't always work as intended (was meant to push the states not used for emissions to the bottom of the matrix)
        if correlation_a > correlation_b:
            return 1
        elif correlation_a < correlation_b:
            return -1
        else:
            return 0

    order = sorted(
        np.arange(0, model.n_states),
        key=cmp_to_key(compare_by_emission_probs)
    )
    return order


@app.function(hide_code=True)
def visualize_markov_chain(startprob, transmat, emissionprob, model: WindowableHMM, threshold=0.05, order=None):
    """Visualize a fit WindowableHMM's hidden states as a Markov chain directed graph.

    Args:
        model: A fit WindowableHMM instance (to extract n_states).
        startprob, transmat, emissionprob: Matrices (cf. plot_hmm_matrices)
        threshold: Minimum transition probability to display an edge.
        order: List of hidden state indices to include (and their display order).
               If None, all states are shown.
    """

    if order is None:
        states = list(range(model.n_states))
    else:
        states = list(order)

    # Compute dominant emission token per state
    dominant_token = {}
    for s in states:
        dominant_token[s] = TOKENS[np.argmax(emissionprob[s])]

    # Build directed graph
    G = nx.DiGraph()
    for s in states:
        G.add_node(s)

    self_loop_edges = {}
    normal_edges = {}
    for i in states:
        for j in states:
            prob = transmat[i, j]
            if prob >= threshold:
                if i == j:
                    self_loop_edges[(i, j)] = prob
                else:
                    normal_edges[(i, j)] = prob
                G.add_edge(i, j)

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

    # Node sizes scaled by start probability
    min_size, max_size = 2000, 5000
    sp = np.array([startprob[s] for s in G.nodes()])
    if sp.max() > sp.min():
        node_sizes = min_size + (max_size - min_size) * (sp - sp.min()) / (sp.max() - sp.min())
    else:
        node_sizes = np.full(len(sp), (min_size + max_size) / 2)

    node_colors = [TOKEN_COLORS.get(dominant_token[s], "#bdc3c7") for s in G.nodes()]

    nx.draw_networkx_nodes(G, pos, node_color=node_colors,
                           node_size=node_sizes, alpha=0.9, ax=ax)

    labels = {s: f"S{s}\n{dominant_token[s]}" for s in G.nodes()}
    nx.draw_networkx_labels(G, pos, labels=labels, font_size=10,
                            font_weight='bold', font_color='white', ax=ax)

    # Draw normal edges
    if normal_edges:
        nx.draw_networkx_edges(
            G, pos,
            edgelist=list(normal_edges.keys()),
            width=[5 * normal_edges[e] for e in normal_edges],
            alpha=0.7,
            edge_color='#555555',
            connectionstyle='arc3,rad=0.2',
            arrows=True,
            arrowsize=20,
            arrowstyle='-|>',
            node_size=max_size,
            ax=ax,
            min_source_margin=15,
            min_target_margin=15
        )

    # Draw self-loops with higher curvature
    if self_loop_edges:
        nx.draw_networkx_edges(
            G, pos,
            edgelist=list(self_loop_edges.keys()),
            width=[5 * self_loop_edges[e] for e in self_loop_edges],
            alpha=0.7,
            edge_color='#555555',
            arrows=True,
            arrowsize=20,
            arrowstyle='-|>',
            node_size=max_size,
            ax=ax,
        )

    # Edge labels
    all_edges = {**normal_edges, **self_loop_edges}
    edge_labels = {e: f"{p:.2f}" for e, p in all_edges.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=8,
        bbox=label_bbox,
        ax=ax,
        connectionstyle='arc3,rad=0.2'
    )

    ax.set_title(f"Hidden State Markov Chain (threshold={threshold:.2f})",
                 fontsize=14, fontweight='bold', pad=20)
    ax.axis('off')
    plt.tight_layout()

    return fig


@app.function
def make_markov_chain_trajectory_inspector(model):
    _order = order_hiddens_by_emission_probs(
        model, model.fitting_history[-1]
    )

    def make_figs(threshold):
        chain_figs = [
            visualize_markov_chain(
                model.model.startprob_prior_.reshape(model.n_states, 1),
                model.model.transmat_prior_,
                model.model.emissionprob_prior_,
                model,
                threshold=threshold,
                order=_order,
            )
        ] + [
            visualize_markov_chain(
                _record["startprob_"].reshape(model.n_states, 1),
                _record["transmat_"],
                _record["emissionprob_"],
                model,
                threshold=threshold,
                order=_order,
            )
            for _record in model.fitting_history
        ]
        return chain_figs

        # order dictates the drawing order, which seems to make it match the heatmap order better
        # Note that should only follow the path for two hops, since the model's predictions after two steps are not relevant (not trained, and it's not defined how they'd be useful i.e. undefined behavior)
        # And note that nodes are labeled with their **dominant** (i.e., most likely) emission, even though they may have multiple possible emissions (and that ambiguity could potentially be informative).

    def show(frame_slider, threshold_slider):
        chain_figs = make_figs(threshold_slider.value)
        return mo.vstack([
            threshold_slider,
            mo.hstack([ frame_slider, mo.md(f"Frame {frame_slider.value}/{frame_slider.stop} (0 = prior)") ]),
            mo.hstack([ chain_figs[frame_slider.value], loss_plot(model, frame_slider.value) ])
        ])

    frame_slider = mo.ui.slider(start=0, stop=len(model.fitting_history) - 1 + 1, step=1, label="Frame", value = len(model.fitting_history) - 1 + 1)

    threshold_slider = mo.ui.slider(start=0.01, stop=0.5, step=0.01, value=0.15, label="Edge threshold")

    return show, frame_slider, threshold_slider


@app.cell
def _(TI_bwd_model):
    mcTI, _, mcTI_threshold_slider = make_markov_chain_trajectory_inspector(TI_bwd_model)
    return mcTI, mcTI_threshold_slider


@app.cell
def _(hmTI_frame_slider, mcTI, mcTI_threshold_slider):
    mcTI(hmTI_frame_slider, mcTI_threshold_slider)
    return


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    # Inferring competence from learned transition probabilities

    Our model has that:
    - P(lift | X_try, Y_try) ~ (Competence(X) + Competence(Y)) >= 5
    - P(lift | X_try, Y_wait) ~ Competence(X) >= 5
    - P(lift | X_wait, Y_wait) ~ 0
    - P(fail | ...) ~ 1 - P(lift | ...)

    with X and Y being A and B in any order,  with the conditioning events in no particular order, and ... being any of the three conditions. Importantly, `discriminability_noise` doesn't change these relationships, it just introduces noise.

    So then let Competence(X) be inferred by Bayesian inference on the probabilities learned by the HMM (where X can be A or B). If the HMM is backwards-learning (i.e., outcome comes first), I think we can just apply Bayes' rule to compute P(lift | ...) and P(fail | ...).
    """)
    return


@app.function
# Step 1: Extract P(lift | actions) from the backward HMM via Bayes' rule
# The backward model sees [outcome, action1, action2].
#
# To compute joint P(O1=a1, O2=a2 | O0=outcome), we must propagate through
# hidden states properly (cf. _joint_prob in hmm_xiang2023.py), NOT multiply
# marginal predict() outputs which would ignore hidden-state dependence.

def infer_lift_prob(bwd_model, pair):
    _model = bwd_model.model
    _tok2idx = bwd_model.token_to_idx

    # P(outcome) from the HMM's learned start/emission matrices
    # (first token in backward sequences is the outcome)
    _p_outcome = {}
    for _outcome in ["lift", "fail"]:
        _oidx = _tok2idx[_outcome]
        _p_outcome[_outcome] = (_model.startprob_ * _model.emissionprob_[:, _oidx]).sum()

    # Compute joint P(O1=a1, O2=a2 | O0=outcome) properly through hidden states:
    #   1. Get state dist after observing outcome: score_samples([outcome])
    #   2. Transition to t=1 states, weight by P(emit a1)
    #   3. Transition to t=2 states, weight by P(emit a2)
    #   (This is the _joint_prob pattern from hmm_xiang2023.py)
    def _joint_prob_given_outcome(outcome, a1, a2):
        _obs = bwd_model._tokens_to_obs([outcome])
        _, _posteriors = _model.score_samples(_obs, [len(_obs)])
        _state_dist_0 = _posteriors[-1]

        _a1_idx = _tok2idx[a1]
        _a2_idx = _tok2idx[a2]

        _state_dist_1 = _state_dist_0 @ _model.transmat_
        _weighted = _state_dist_1 * _model.emissionprob_[:, _a1_idx]
        _state_dist_2 = _weighted @ _model.transmat_
        return (_state_dist_2 * _model.emissionprob_[:, _a2_idx]).sum()

    # Conditions: unordered action pairs
    _conditions = [
        ("A_try", "B_try"),
        ("A_try", "B_wait"),
        ("B_try", "A_wait"),
        ("A_wait", "B_wait"),
    ]

    _hmm_p_lift = {}
    for _a1, _a2 in _conditions:
        _label = f"{_a1}, {_a2}"
        _p_actions_given = {}
        for _outcome in ["lift", "fail"]:
            # ordering 1: a1 then a2
            _p_order1 = _joint_prob_given_outcome(_outcome, _a1, _a2)

            if _a1 != _a2:  # (defensive programming)
                # ordering 2: a2 then a1
                _p_order2 = _joint_prob_given_outcome(_outcome, _a2, _a1)
                _p_actions_given[_outcome] = _p_order1 + _p_order2
            else:
                _p_actions_given[_outcome] = _p_order1

        # Bayes' rule: P(lift | {a1, a2}) = P({a1, a2} | lift) * P(lift) / P({a1, a2})
        _p_actions = sum(_p_actions_given[o] * _p_outcome[o] for o in ["lift", "fail"])
        if _p_actions > 0:
            _hmm_p_lift[_label] = _p_actions_given["lift"] * _p_outcome["lift"] / _p_actions
        else:
            _hmm_p_lift[_label] = float("nan")

    # Ground truth via Monte Carlo
    _n_mc = 10000
    _gt_p_lift = {}
    for _a1, _a2 in _conditions:
        _label = f"{_a1}, {_a2}"
        _lifts = 0
        for _ in range(_n_mc):
            _lifts += 1 if generate_outcome(pair, _a1, _a2) == "lift" else 0
        _gt_p_lift[_label] = _lifts / _n_mc

    lift_prob = {
        "conditions": [f"{a1}, {a2}" for a1, a2 in _conditions],
        "hmm_p_lift": _hmm_p_lift,
        "gt_p_lift": _gt_p_lift,
        "p_outcome": _p_outcome,
        "noise": 0.5,
        # ^ don't just feed in discriminability_noise.value -- it should be inferred by the participant model from the data, OR a free parameter set to optimize fit to real participant data. it's cheating if the model of the participant knows generation parameters! should only have access to observations
    }
    return lift_prob


@app.cell
def _(TI_bwd_model, pair):
    lift_prob = infer_lift_prob(TI_bwd_model, pair)
    return (lift_prob,)


@app.cell
def _(lift_prob, pair):
    # Step 2: Comparison table + competence inversion
    _lp = lift_prob

    _comparison_df = pl.DataFrame({
        "Condition": _lp["conditions"],
        "HMM P(lift)": [_lp["hmm_p_lift"][c] for c in _lp["conditions"]],
        "Ground Truth P(lift)": [_lp["gt_p_lift"][c] for c in _lp["conditions"]],
    }).with_columns(
        Difference=(col("HMM P(lift)") - col("Ground Truth P(lift)"))
    )

    # Invert solo-try conditions to estimate competence
    # Effort ~ Beta(alpha, 1), so p_lift = P(C*E >= 5) = 1 - (5/C)^alpha when C > 5
    # Inverting: C = 5 / (1 - p_lift)^(1/alpha)
    _noise = _lp["noise"]
    _alpha = 1.0 / max(_noise**2, 0.00001)

    def _invert_competence(p_lift):
        if p_lift == 0:
            return None  # can only say C <= 5
        if p_lift == 1:
            return float("inf")  # always succeeds solo, C unbounded above
        return 5.0 / (1.0 - p_lift) ** (1.0 / _alpha)

    _agents = ["A", "B"]
    _true_comp = [pair.A_competence, pair.B_competence]
    _solo_conditions = ["A_try, B_wait", "B_try, A_wait"]
    _hmm_inferred = [_invert_competence(_lp["hmm_p_lift"][c]) for c in _solo_conditions]
    _gt_inferred = [_invert_competence(_lp["gt_p_lift"][c]) for c in _solo_conditions]

    _competence_rows = []
    for _i, _agent in enumerate(_agents):
        _competence_rows.append({
            "Agent": _agent,
            "True Competence": _true_comp[_i],
            "GT-Inferred": _gt_inferred[_i] if _gt_inferred[_i] is not None else float("nan"),
            "HMM-Inferred": _hmm_inferred[_i] if _hmm_inferred[_i] is not None else float("nan"),
            "Solo P(lift) [HMM]": _lp["hmm_p_lift"][_solo_conditions[_i]],
            "Solo P(lift) [GT]": _lp["gt_p_lift"][_solo_conditions[_i]],
        })
    competence_results = pl.DataFrame(_competence_rows)

    mo.vstack([
        mo.md(f"""
    ### P(lift | actions): HMM vs Ground Truth

    HMM marginal: P(lift) = {_lp['p_outcome']['lift']:.3f}, P(fail) = {_lp['p_outcome']['fail']:.3f}

    Effort distribution: Beta(alpha={_alpha:.1f}, 1), noise={_noise}
        """),
        _comparison_df,
        mo.md("### Competence inversion from solo-try conditions"),
        mo.md(f"Formula: C = 5 / (1 - P(lift))^(1/{_alpha:.1f}). NaN means P(lift) = 0, so C <= 5."),
        competence_results,
    ])
    return


@app.function
# Step 3: Joint competence inference using solo AND joint-try conditions
# Solo: P(lift | X_try, Y_wait) = 1 - (5/C_X)^alpha  (when C_X > 5, else 0)
# Joint: P(lift | A_try, B_try) = P(C_A*E_A + C_B*E_B >= 5)  (no closed form)
#
# We find (C_A, C_B) that best fits all three observed lift probabilities.

def infer_competence_from_lift_prob(_lp, pair):
    _noise = _lp["noise"]
    _alpha = 1.0 / max(_noise**2, 0.00001)
    _n_mc = 10000

    def _predict_p_lift_solo(c, alpha):
        if c <= 5:
            return 0.0
        return 1.0 - (5.0 / c) ** alpha

    def _predict_p_lift_joint(c_a, c_b, alpha, n_mc):
        _e_a = np.random.beta(alpha, 1, size=n_mc)
        _e_b = np.random.beta(alpha, 1, size=n_mc)
        return np.mean(c_a * _e_a + c_b * _e_b >= 5)

    _p_lift_a_solo = _lp["hmm_p_lift"]["A_try, B_wait"]
    _p_lift_b_solo = _lp["hmm_p_lift"]["B_try, A_wait"]
    _p_lift_joint = _lp["hmm_p_lift"]["A_try, B_try"]

    def _loss(params):
        _c_a, _c_b = params
        if _c_a < 0 or _c_b < 0:
            return 1e6
        _pred_a = _predict_p_lift_solo(_c_a, _alpha)
        _pred_b = _predict_p_lift_solo(_c_b, _alpha)
        np.random.seed(0)
        _pred_joint = _predict_p_lift_joint(_c_a, _c_b, _alpha, _n_mc)
        return (
            (_pred_a - _p_lift_a_solo) ** 2
            + (_pred_b - _p_lift_b_solo) ** 2
            + (_pred_joint - _p_lift_joint) ** 2
        )

    # Try multiple starting points to avoid local minima
    _best = None
    for _c_a_init in [2, 4, 6, 8]:
        for _c_b_init in [2, 4, 6, 8]:
            _result = minimize(_loss, [_c_a_init, _c_b_init], method="Nelder-Mead")
            if _best is None or _result.fun < _best.fun:
                _best = _result

    _c_a_best, _c_b_best = _best.x

    joint_competence_results = pl.DataFrame({
        "Agent": ["A", "B"],
        "True Competence": [pair.A_competence, pair.B_competence],
        "Solo+Joint-Inferred": [_c_a_best, _c_b_best],  # i.e., inferred from both solo and joint attempts
    })

    original_predicted_p = {
        'A_solo': _p_lift_a_solo,
        'B_solo': _p_lift_b_solo,
        'joint': _p_lift_joint
    }

    optimized_predicted_p = {
        'A_solo': _predict_p_lift_solo(_c_a_best, _alpha),
        'B_solo': _predict_p_lift_solo(_c_b_best, _alpha),
        'joint': _predict_p_lift_joint(_c_a_best, _c_b_best, _alpha, _n_mc)
    }
    return joint_competence_results, original_predicted_p, optimized_predicted_p


@app.cell
def _(lift_prob, pair):
    np.random.seed(0)
    joint_competence_results, original_predicted_p, optimized_predicted_p = infer_competence_from_lift_prob(lift_prob, pair)

    # Show predicted vs observed lift probs for the joint fit
    mo.vstack([
        mo.md(f"""
    ### Joint competence inference (using solo + both-try conditions)

    Jointly optimizes (C_A, C_B) to fit three HMM-derived lift probabilities:
    - P(lift | A_try, B_wait) = {original_predicted_p['A_solo']:.3f} (predicted: {optimized_predicted_p['A_solo']:.3f})
    - P(lift | B_try, A_wait) = {original_predicted_p['B_solo']:.3f} (predicted: {optimized_predicted_p['B_solo']:.3f})
    - P(lift | A_try, B_try) = {original_predicted_p['joint']:.3f} (predicted: {optimized_predicted_p['joint']:.3f})
        """),
        joint_competence_results,
    ])
    return (joint_competence_results,)


@app.function
# Step 4: Bar chart comparing true vs inferred competence (solo and solo+joint)

def visualize_competence_results(joint_competence_results, pair):
    _plot_df = joint_competence_results.unpivot(
        on=["True Competence", "Solo+Joint-Inferred"],
        index="Agent",
        variable_name="Source",
        value_name="Competence",
    )

    _fig, _ax = plt.subplots(figsize=(7, 4))
    _x = np.arange(2)
    _width = 0.25
    _sources = ["True Competence", "Solo+Joint-Inferred"]
    _colors = ["#2ecc71", "#e67e22"]

    for _i, (_source, _color) in enumerate(zip(_sources, _colors)):
        _vals = _plot_df.filter(col("Source") == _source)["Competence"].to_numpy()
        _ax.bar(_x + _i * _width, _vals, _width, label=_source, color=_color)

    _ax.axhline(y=5, ls="--", color="gray", label="Lift threshold")
    _ax.set_xticks(_x + _width)
    _ax.set_xticklabels(["A", "B"])
    _ax.set_ylabel("Competence")
    _ax.set_title(f"Competence Inference (A={pair.A_competence}, B={pair.B_competence})")
    _ax.legend()
    _fig.tight_layout()
    return _fig


@app.cell
def _(joint_competence_results, pair):
    visualize_competence_results(joint_competence_results, pair)
    return


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    # Cross-prediction from null to intact data

    Null data so defined that this tests that the model needs to learn the temporal contingencies of the actions.
    """)
    return


@app.cell
def _(bwd_windows, n_windows, pair):
    null_actions_gen, _, _ = make_null_actions_generator()

    np.random.seed(4)
    null_fwd_windows = np.array(generate_windows_with_actions_generator(pair, n_windows, null_actions_gen))
    null_bwd_windows = np.flip(null_fwd_windows)  # for the backwards model, we flip each window
    _null_bwd_windows_types = get_window_types(null_bwd_windows)

    _test_scores = []
    _train_scores = []
    _test_scores_b = []
    _train_scores_b = []
    for _i, (_train, _test) in enumerate(LeaveOneGroupOut().split(null_bwd_windows, groups=_null_bwd_windows_types)):
        np.random.seed(_i)

        _train_windows = null_bwd_windows[_train]

        _rng = np.random.default_rng()
        _test_scores_iteration_means = []
        _test_scores_b_iteration_means = []
        for _j in range(10):
            print(f"Fold {_i}, Iteration {_j}")
            _test_windows = _rng.choice(bwd_windows, size=len(_test), replace=False)
            # ^ could use ALL intact data, but for the sake of matching the precision of the scores to the regular version, sample the same number as in the selected test set.

            _bwd_model = WindowableHMM(n_states=10, max_iter=120)
            for _win in _train_windows:
                _bwd_model.add_to_history(_win)
            _bwd_model.fit()

            print(f"{_bwd_model.model.monitor_.iter} iterations out of {_bwd_model.model.monitor_.n_iter} max")

            _test_scores_iteration_means.append(_bwd_model.accuracy(_test_windows, backward_learning=True).reshape(-1).mean())
            _test_scores_b_iteration_means.append(_bwd_model.accuracy_about_stronger_agent(_test_windows, backward_learning=True).reshape(-1).mean())

        _train_scores.append(_bwd_model.accuracy(_train_windows, backward_learning=True).reshape(-1))
        _test_scores.append(_test_scores_iteration_means)

        _train_scores_b.append(_bwd_model.accuracy_about_stronger_agent(_train_windows, backward_learning=True).reshape(-1))
        _test_scores_b.append(_test_scores_b_iteration_means)

    null_cp_accuracy_results = pl.concat([
        tidy_scores(_train_scores).with_columns(label=lit('train'), target=lit('both agents')),
        tidy_scores(_test_scores).with_columns(label=lit('test'), target=lit('both agents')),
        tidy_scores(_train_scores_b).with_columns(label=lit('train'), target=lit('stronger agent')),
        tidy_scores(_test_scores_b).with_columns(label=lit('test'), target=lit('stronger agent'))
    ])
    return (null_cp_accuracy_results,)


@app.cell
def _(null_cp_accuracy_results):
    show_accuracy(null_cp_accuracy_results)
    return


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