# /// 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
    )
    from mcp_hmm.hmm_xiang2023_bwd_meta import (
        show_accuracy,
        order_hiddens_by_emission_probs,
        visualize_markov_chain,
        make_heatmap_trajectory_inspector,
        make_markov_chain_trajectory_inspector,
        tidy_scores,
        infer_competence_from_lift_prob,
        visualize_competence_results,
        get_window_types
    )

    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.cell(hide_code=True)
def _(pair):
    n_windows = 100

    np.random.seed(4)
    fwd_windows = np.array(generate_windows(pair, num_windows=n_windows))
    _fwd_windows_types = get_window_types(fwd_windows)

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

        train_windows = fwd_windows[train]
        test_windows = fwd_windows[test]

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

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

        train_scores.append(fwd_model.accuracy(train_windows).reshape(-1))
        test_scores.append(fwd_model.accuracy(test_windows).reshape(-1))

        # train_scores_b.append(fwd_model.accuracy_about_stronger_agent(train_windows).reshape(-1))
        # test_scores_b.append(fwd_model.accuracy_about_stronger_agent(test_windows).reshape(-1))
    return fwd_windows, n_windows, test_scores, train_scores


@app.cell(hide_code=True)
def _(test_scores, train_scores):
    accuracy_results = pl.concat([
        tidy_scores(train_scores).with_columns(label=lit('train'), target=lit('outcome')),
        tidy_scores(test_scores).with_columns(label=lit('test'), target=lit('outcome')),
        # 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.cell(hide_code=True)
def _():
    mo.md(r"""
    # Inspecting trajectory of hidden weights during fit
    """)
    return


@app.cell(hide_code=True)
def _(fwd_windows):
    np.random.seed(2)
    _n_states = 10
    # TI = trajectory inspector
    TI_fwd_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 fwd_windows:
        TI_fwd_model.add_to_history(_win)
    TI_fwd_model.fit()
    return (TI_fwd_model,)


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


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


@app.cell(hide_code=True)
def _(TI_fwd_model):
    mcTI, a, mcTI_threshold_slider = make_markov_chain_trajectory_inspector(TI_fwd_model)
    return a, mcTI, mcTI_threshold_slider


@app.cell(hide_code=True)
def _(a, mcTI, mcTI_threshold_slider):
    mcTI(a, 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). Since the HMM is forwards-learning (i.e., sequences are [action1, action2, outcome]), we can directly estimate P(lift | actions) and P(fail | actions) by conditioning on the observed actions and reading off the predicted emission for the outcome token. No Bayes' rule inversion is needed (unlike the backwards-learning model).
    """)
    return


@app.function
# Step 1: Extract P(lift | actions) directly from the forward HMM
# The forward model sees [action1, action2, outcome].
# We condition on the two actions and read off P(outcome) directly.

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

    def _p_outcome_given_actions(a1, a2):
        """Compute P(outcome=lift) and P(outcome=fail) given observing [a1, a2]."""
        _obs = fwd_model._tokens_to_obs([a1, a2])
        _, _posteriors = _model.score_samples(_obs, [len(_obs)])
        _state_dist_after_actions = _posteriors[-1]

        # One more transition + emission to get P(next token = outcome)
        _next_state_dist = _state_dist_after_actions @ _model.transmat_
        _p = {}
        for _outcome in ["lift", "fail"]:
            _oidx = _tok2idx[_outcome]
            _p[_outcome] = (_next_state_dist * _model.emissionprob_[:, _oidx]).sum()
        return _p

    # 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}"

        # ordering 1: a1 then a2
        _p1 = _p_outcome_given_actions(_a1, _a2)

        if _a1 != _a2:
            # ordering 2: a2 then a1 (actions are unordered in the domain)
            _p2 = _p_outcome_given_actions(_a2, _a1)
            # Average P(lift) across both orderings, weighted by their total mass
            _total_lift = _p1["lift"] + _p2["lift"]
            _total_all = _total_lift + _p1["fail"] + _p2["fail"]
        else:
            _total_lift = _p1["lift"]
            _total_all = _total_lift + _p1["fail"]

        _hmm_p_lift[_label] = _total_lift / _total_all if _total_all > 0 else float("nan")

    # Marginal P(outcome) for display (average over all conditions)
    _p_outcome = {"lift": 0.0, "fail": 0.0}
    for _label in _hmm_p_lift:
        _p_outcome["lift"] += _hmm_p_lift[_label]
        _p_outcome["fail"] += 1.0 - _hmm_p_lift[_label]
    _total = _p_outcome["lift"] + _p_outcome["fail"]
    _p_outcome = {k: v / _total for k, v in _p_outcome.items()}

    # 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_fwd_model, pair):
    lift_prob = infer_lift_prob(TI_fwd_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.cell
def _(lift_prob, pair):
    # Show predicted vs observed lift probs for the joint fit
    np.random.seed(0)
    joint_competence_results, original_predicted_p, optimized_predicted_p = infer_competence_from_lift_prob(lift_prob, pair)

    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.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 _(fwd_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_fwd_windows_types = get_window_types(null_fwd_windows)

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

        _train_windows = null_fwd_windows[_train]

        _rng = np.random.default_rng()
        _test_scores_iteration_means = []
        for _j in range(10):
            print(f"Fold {_i}, Iteration {_j}")
            _test_windows = _rng.choice(fwd_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.

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

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

            _test_scores_iteration_means.append(_fwd_model.accuracy(_test_windows).reshape(-1).mean())

        _train_scores.append(_fwd_model.accuracy(_train_windows).reshape(-1))
        _test_scores.append(_test_scores_iteration_means)

    null_cp_accuracy_results = pl.concat([
        tidy_scores(_train_scores).with_columns(label=lit('train'), target=lit('outcome')),
        tidy_scores(_test_scores).with_columns(label=lit('test'), target=lit('outcome')),
        # 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


@app.cell
def _():
    return


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