# /// 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(hide_code=True):
    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    import seaborn as sns
    import polars as pl
    import matplotlib.patches as mpatches
    from collections import Counter, defaultdict
    import warnings

    warnings.filterwarnings("ignore")

    import hmmlearn.vhmm
    from hmmlearn.base import ConvergenceMonitor

    SEED = 4

    TOKENS = ["A_try", "A_wait", "B_try", "B_wait", "lift", "fail"]

    TOKEN_COLORS = {
        "A_try": "#e67e22",
        "B_try": "#e67e22",
        "A_wait": "#95a5a6",
        "B_wait": "#95a5a6",
        "lift": "#2ecc71",
        "fail": "#e74c3c",
    }


@app.class_definition(hide_code=True)
class AgentPair:
    # A = A_competence * A_effort
    # A_effort ~ U(0,1) or = 1
    # B = B_competence * B_effort
    # B_effort ~ U(0,1) or = 1
    # for A_try,B_try,lift / B_try,A_try,lift:
    # easiest credit assignment is when Pr[(A+B)>5]>0 but |Pr[A>5]-Pr[B>5]|=~1 (e.g. Pr[A>5]=0 and Pr[B>5] is large. here, B is the one that is responsible)
    # impossible (i.e., not able to be better than chance) credit assignment is when Pr[(A+B)>5]>0 but |Pr[A>5]-Pr[B>5]|=~0

    def __init__(self, A_competence=None, B_competence=None, discriminability=0.5, discriminability_noise=0.5):
        """
        discriminability: 0 = impossible credit assignment, 1 = easiest
        discriminability_noise: 0 = no variability in effort, 1 = maximum variability in effort
        """

        min_lift_competence = 5

        if A_competence is None:
            self.A_competence = min_lift_competence * (1 - discriminability)
        else:
            self.A_competence = A_competence

        if B_competence is None:
            self.B_competence = min_lift_competence * (1 + discriminability)
        else:
            self.B_competence = B_competence

        self.discriminability = discriminability
        self.discriminability_noise = discriminability_noise

    def _generate_effort(self):
        # hand-crafted in desmos such that discriminability_noise=1 -> effort ~ U(0,1), while discriminability_noise=0 -> effort = 1, and it interpolates between those
        beta = 1.0
        alpha = 1.0 / np.max([self.discriminability_noise**2.0, 0.00001])
        return np.random.beta(alpha, beta)

    def A_effort(self):
        return self._generate_effort()

    def B_effort(self):
        return self._generate_effort()


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    Use (optionally) disjoint-windowed HMM to model box-lifting action sequences.

    Notes:
    - 2/5/26: it seems that maybe probability of outcomes are being underrated when they are unlikely in general, even when they are likely in context. i.e. **HMM seems sensitive to base rates and not just conditional probability/exhibits the base rate fallacy!** -- ~~IDEA don't have the input sequences be random; instead, just show all combinations the same amount of times~~; IDEA to diagnose if base rate fallacy is driving the weird results (see sanity checks output), calculate the frequency of each n-gram which the sanity checks are trying to model (as it varies slightly per each random seed).
    - 2/8/26 (currently, A=0 and B=10, and SEED=3): **backward model provides evidence of being able to de-alias outcomes (i.e., assign credit to/infer which agent tried)**! however, I'm also finding it's learning some weird statistical patterns (that I believe aren't there):
        - e.g., the backwards model is completely sure that in `?,B_wait,fail`, `?` is `A_try`, even though it should be equally likely as `A_wait`. maybe this is due to a base rate fallacy? (since WW is half as likely as TW).
        - similarly, for `?,fail` it should only infer `Pr[B_try] = 0`, and `Pr[A_wait]` should match `Pr[A_try]`, and `Pr[B_wait]` should be relatively high (not sure if should be higher than `Pr[A_try]`, though; would need to do the math)
        - there seems to be a relationship between A_try and fail, even though A_try/A_wait does not affect outcome; maybe it's due to some weird dependency like B_try-(+)>lift, A_try-(-)>B_try, thus A_try-(-)>lift; **there should be no relationship between A's and B's actions**. ~~***maybe I should feed in both agents' actions at once (i.e. encode the 4 unordered combinations A_try/wait,B_try/wait as four tokens)?**~~ no, that would actually just be making the problem for the HMM trivial. its job is to infer across time
    - 2/8/26: ~~should I be keeping the sequences as separate windows as I am currently doing? maybe I mean to/should have all the observations come in as a continuous sequence, though I'm not sure if this should make a meaningful difference, and if anything, feeding sequences as separate windows should make things easier for the HMM. _though maybe this is actually leading to the weird statistical patterns somehow (by not allowing the HMM to learn on its own what contingencies are not relevant)?_~~
        - 2/8/26: update, yeah I'll combine the sequences both to allow the HMM learn on its own what contingencies are irrelevant, and to make interpreting the predictions of illegal actions make some more sense (i.e., as I'm loosening the structure of the input data). the weird statistical patterns it's learning are still there, though

    ---

    - 2/10/26 found that was generating sequence incorrectly (needed to resample actions rather than shuffle them), and now all the above seems to be fixed! for now, reverting back to use flat=False to make things easier for model as a starting point. but still using variational HMM because there is some weird order effect when using basic categorical model
    """)
    return


@app.class_definition
class DebuggingMonitor(ConvergenceMonitor):
    class PROPS():
        NORM_POSTERIORS = ['startprob_', 'transmat_', 'emissionprob_']

    def __init__(self, model, tol, n_iter, verbose, save_props_while_fitting=[]):
        super().__init__(tol, n_iter, verbose)

        self.model: hmmlearn.vhmm.VariationalCategoricalHMM = model

        self.save_props_while_fitting_ = save_props_while_fitting
        self.model_history = []

    def report(self, log_prob):
        if len(self.save_props_while_fitting_) > 0:
            self.model_history.append(self._get_weights_record())

        super().report(log_prob)

    def _get_weights_record(self):
        return {
            name: getattr(self.model, name)
            for name in self.save_props_while_fitting_
        }


@app.class_definition
class WindowableHMM:
    """HMM that can use just batches of observations (if flat=False)."""

    def __init__(self, n_states=24, max_iter=150, flat=False,
                 save_props_while_fitting=[], custom_model_args={}):
        self.n_states = n_states
        self.token_to_idx = {t: i for i, t in enumerate(TOKENS)}

        self.model = hmmlearn.vhmm.VariationalCategoricalHMM(
            n_components=self.n_states,
            n_features=len(TOKENS),
            n_iter=max_iter,
            random_state=SEED,
            **custom_model_args
        )
        cur_monitor = self.model.monitor_
        self.model.monitor_ = DebuggingMonitor(self.model,
                                  cur_monitor.tol,
                                  cur_monitor.n_iter,
                                  cur_monitor.verbose,
                                  save_props_while_fitting = save_props_while_fitting)

        self.history = []
        self.flat = flat

    def _tokens_to_obs(self, tokens):
        return np.array([[self.token_to_idx[t]] for t in tokens])

    def add_to_history(self, tokens):
        obs = self._tokens_to_obs(tokens)
        self.history.append(obs)

    def fit(self):
        """Fit HMM to tokens"""
        sequences = np.array(self.history).reshape(-1, 1)

        if self.flat:
            self.model.fit(sequences)
        else:
            lengths = [len(seq) for seq in self.history]
            self.model.fit(sequences, lengths)

    def predict(self, tokens):
        """Get state distribution after input observations (accepts a single window)"""
        try:
            obs = self._tokens_to_obs(tokens)
            _, posteriors = self.model.score_samples(obs, [len(obs)])
            state_dist = posteriors[-1]

            # Compute next state prob dist and emission prob dist
            next_state_dist = state_dist @ self.model.transmat_
            pred_dist = next_state_dist @ self.model.emissionprob_
        except Exception as e:
            print(e)
            pred_dist = np.full(len(TOKENS), np.nan)

        pred_pretty = {t: pred_dist[i] for i, t in enumerate(TOKENS)}

        frequentist_pred = {
            next_token: self._get_conditional_frequency(tokens, next_token)
            for next_token in TOKENS
        }

        return {"hmm": pred_pretty, "frequentist": frequentist_pred}

    def score(self, tokens):
        """Compute log prob of tokens under the model (accepts a single window)"""
        obs = self._tokens_to_obs(tokens)
        return self.model.score(obs, [len(obs)])

    def accuracy(self, token_sequences, backward_learning=False):
        """
        If backward-learning:

        Probability that the model's 2 sampled tokens after token 1
        match actual tokens 2 and 3 (in any order; i.e. disjoint probability).

        If forward-learning:

        Probability that the model's 1 sampled token after tokens 2 and 3 (in either order; i.e. disjoint probability)
        matches actual token 3.
        """
        if backward_learning:
            # implementation for backward-learning
            accuracy_rates = []
            for tokens in token_sequences:
                obs = self._tokens_to_obs(tokens)
                _, posteriors = self.model.score_samples([obs[0]], [1])
                state_dist_0 = posteriors[-1]

                # Compute true joint P(O1=a, O2=b | O0) by conditioning t=2
                # state distribution on the actual emission at t=1 (don't just
                # use the product of marginals, which ignores the dependence
                # that arises from marginalizing over hidden states)
                def _joint_prob(_o1, _o2):
                    state_dist_1 = state_dist_0 @ self.model.transmat_
                    weighted = state_dist_1 * self.model.emissionprob_[:, _o1]
                    state_dist_2 = weighted @ self.model.transmat_
                    return (state_dist_2 * self.model.emissionprob_[:, _o2]).sum()

                o1, o2 = obs[1].item(), obs[2].item()
                p_actions_match = _joint_prob(o1, o2)
                p_actions_match_rev = _joint_prob(o2, o1)

                # These two events are disjoint (mutually exclusive), so
                # P(match in either order) = P(order 1) + P(order 2).
                # Guard: if o1 == o2, both orders are the same event.
                if o1 == o2:
                    p_actions_match_unordered = p_actions_match
                else:
                    p_actions_match_unordered = p_actions_match + p_actions_match_rev

                accuracy_rates.append(p_actions_match_unordered)

            return np.array(accuracy_rates)
        else:
            # implementation for forward-learning
            accuracy_rates = []
            for tokens in token_sequences:
                obs = self._tokens_to_obs(tokens)
                o_actual = obs[2].item()

                def _p_outcome_given_ordering(_a1, _a2):
                    action_obs = np.array([[_a1], [_a2]])
                    _, posteriors = self.model.score_samples(action_obs, [2])
                    state_dist = posteriors[-1]
                    next_state_dist = state_dist @ self.model.transmat_
                    p_correct = (next_state_dist * self.model.emissionprob_[:, o_actual]).sum()
                    lift_idx = self.token_to_idx["lift"]
                    fail_idx = self.token_to_idx["fail"]
                    p_outcomes = (
                        (next_state_dist * self.model.emissionprob_[:, lift_idx]).sum()
                        + (next_state_dist * self.model.emissionprob_[:, fail_idx]).sum()
                    )
                    return p_correct, p_outcomes

                a1, a2 = obs[0].item(), obs[1].item()
                p_correct_1, p_total_1 = _p_outcome_given_ordering(a1, a2)

                if a1 == a2:
                    p_accuracy = p_correct_1 / p_total_1 if p_total_1 > 0 else float('nan')
                else:
                    p_correct_2, p_total_2 = _p_outcome_given_ordering(a2, a1)
                    total_correct = p_correct_1 + p_correct_2
                    total_all = p_total_1 + p_total_2
                    p_accuracy = total_correct / total_all if total_all > 0 else float('nan')

                accuracy_rates.append(p_accuracy)

            return np.array(accuracy_rates)

    def accuracy_about_stronger_agent(self, token_sequences, backward_learning=False):
        """Probability that in the model's 2 sampled tokens after token 1, the one
        which is for the stronger agent matches the actual token 2 or 3 (whichever
        is for the stronger agent)

        If model is not backward-learning, not yet implemented.
        """
        if not backward_learning:
            raise Exception('not yet implemented for forward-learning model')

        B_ACTION_IDXS = {self.token_to_idx["B_try"], self.token_to_idx["B_wait"]}
        accuracy_rates = []
        for tokens in token_sequences:
            obs = self._tokens_to_obs(tokens)
            o1, o2 = obs[1].item(), obs[2].item()

            # Find B's actual action among positions 1 and 2
            if o1 in B_ACTION_IDXS:
                b = o1
            elif o2 in B_ACTION_IDXS:
                b = o2
            else:
                # B's action was at position 0 (the conditioning token);
                # nothing to predict about B
                accuracy_rates.append(np.nan)
                continue

            _, posteriors = self.model.score_samples([obs[0]], [1])
            state_dist_0 = posteriors[-1]

            # P(O1=b | O0)
            state_dist_1 = state_dist_0 @ self.model.transmat_
            p_o1_is_b = (state_dist_1 * self.model.emissionprob_[:, b]).sum()

            # P(O2=b | O0) — marginalizing over all O1
            state_dist_2 = state_dist_1 @ self.model.transmat_
            p_o2_is_b = (state_dist_2 * self.model.emissionprob_[:, b]).sum()

            # P(O1=b AND O2=b | O0)
            weighted = state_dist_1 * self.model.emissionprob_[:, b]
            state_dist_2_cond = weighted @ self.model.transmat_
            p_both_b = (state_dist_2_cond * self.model.emissionprob_[:, b]).sum()

            # Inclusion-exclusion
            p_any_b = p_o1_is_b + p_o2_is_b - p_both_b
            accuracy_rates.append(p_any_b)

        return np.array(accuracy_rates)

    def _get_conditional_frequency(self, conditional_tokens, next_token):
        n = len(conditional_tokens) + 1
        target_ngram = tuple(conditional_tokens) + (next_token,)
        context = tuple(conditional_tokens)
        match_count = 0
        context_count = 0
        if self.flat:
            # treat all windows as one continuous sequence
            all_tokens = [
                TOKENS[idx[0]] for window_obs in self.history for idx in window_obs
            ]
            for i in range(len(all_tokens) - n + 1):
                ngram = tuple(all_tokens[i : i + n])
                if ngram[: len(context)] == context:
                    context_count += 1
                    if ngram == target_ngram:
                        match_count += 1
        else:
            # treat each window independently, only checking the start of each window
            for window_obs in self.history:
                window_tokens = [TOKENS[idx[0]] for idx in window_obs]
                if (
                    len(window_tokens) >= n
                    and tuple(window_tokens[: len(context)]) == context
                ):
                    context_count += 1
                    if tuple(window_tokens[:n]) == target_ngram:
                        match_count += 1
        return match_count / context_count if context_count > 0 else np.nan

    @property
    def convergence_history(self):
        """Returns history of log likelihoods"""
        return np.array(self.model.monitor_.history)

    @property
    def fitting_history(self):
        """Returns history of hidden weights"""
        return np.array(self.model.monitor_.model_history)


@app.cell
def _():
    np.random.seed(SEED)
    pair = AgentPair(discriminability=0.5, discriminability_noise=0.5)
    windows = generate_windows(pair, num_windows=500)
    results = fit_predict_model(windows)
    return pair, results, windows


@app.function(hide_code=True)
def visualize_agent_contributions_sampdist(pair):
    # Simulate sampling distribution of A and B, then visualize
    _n_samples = 10000
    np.random.seed(SEED)

    _A_contributions = np.array(
        [pair.A_competence * pair.A_effort() for _ in range(_n_samples)]
    )
    _B_contributions = np.array(
        [pair.B_competence * pair.B_effort() for _ in range(_n_samples)]
    )
    _combined = _A_contributions + _B_contributions

    if np.max(_A_contributions) > np.max(_B_contributions):
        print(Warning('⚠️ Assumption violated! A is currently able to contribute more than B. B should always be the stronger agent, by convention'))

    _threshold = 5

    _fig, _axes = plt.subplots(1, 3, figsize=(14, 4))

    _axes[0].hist(_A_contributions, bins=50, color="#e67e22", alpha=0.7, density=True)
    _axes[0].axvline(
        _threshold, color="black", linestyle="--", label=f"threshold={_threshold}"
    )
    _axes[0].set_title(f"A contributions (competence={pair.A_competence:.1f})")
    _axes[0].set_xlabel("Contribution")
    _axes[0].set_xlim(0, 10)
    _pA = np.mean(_A_contributions >= _threshold)
    _axes[0].legend([f"P(A>={_threshold}) = {_pA:.2f}"])

    _axes[1].hist(_B_contributions, bins=50, color="#3498db", alpha=0.7, density=True)
    _axes[1].axvline(
        _threshold, color="black", linestyle="--", label=f"threshold={_threshold}"
    )
    _axes[1].set_title(f"B contributions (competence={pair.B_competence:.1f})")
    _axes[1].set_xlabel("Contribution")
    _axes[1].set_xlim(0, 10)
    _pB = np.mean(_B_contributions >= _threshold)
    _axes[1].legend([f"P(B>={_threshold}) = {_pB:.2f}"])

    _axes[2].hist(
        _A_contributions, bins=50, color="#e67e22", alpha=0.5, density=True, label="A"
    )
    _axes[2].hist(
        _B_contributions, bins=50, color="#3498db", alpha=0.5, density=True, label="B"
    )
    _axes[2].hist(
        _combined, bins=50, color="#2ecc71", alpha=0.5, density=True, label="A+B"
    )
    _axes[2].axvline(_threshold, color="black", linestyle="--")
    _pAB = np.mean(_combined >= _threshold)
    _axes[2].set_title(f"Combined: P(A+B>={_threshold}) = {_pAB:.2f}")
    _axes[2].set_xlabel("Contribution")
    _axes[2].legend()

    plt.suptitle(
        f"discriminability={pair.discriminability}, noise={pair.discriminability_noise}, A competence={pair.A_competence}",
        fontweight="bold",
    )
    plt.tight_layout()
    return _fig


@app.cell
def _(pair):
    visualize_agent_contributions_sampdist(pair)
    return


@app.cell(hide_code=True)
def _(results):
    mo.vstack(
        [
            plot_results(results),
            mo.md("""
        ~~**current interpretation of the double-descent in both models (when `flat=True`):** HMM starts learning statistical patterns from the start, but then after enough iterations it learns that, in forward model: the outcome does not predict the next token (first action of another sequence) at all; or in backward model: each second action token does not predict the next token (outcome of another sequence) at all. this was inferred based on the double-descent only appearing when fitting the HMM on continuous strings of tokens instead of discrete sequences (i.e. as separate windows)~~

        **one way to understand double-descent even when `flat=False` (happens sometime, at least with variational categorical HMM):** HMM first learns a spurious relationship between agents' actions/maybe cares about their order, but then learns that it's irrelevant and benefits from discarding/deprioritizing this information
              """),
        ]
    )
    return


@app.cell
def _(results):
    mo.vstack(
        [
            # before fitting
            plot_hmm_matrices(
                results["model_forward"].model.startprob_prior_.reshape(
                    results["model_forward"].n_states, 1
                ),
                results["model_forward"].model.transmat_prior_,
                results["model_forward"].model.emissionprob_prior_,
            ),
            # ???
            plot_hmm_matrices(
                results["model_forward"].model.startprob_.reshape(
                    results["model_forward"].n_states, 1
                ),
                results["model_forward"].model.transmat_,
                results["model_forward"].model.emissionprob_,
            ),
            # after fitting
            plot_hmm_matrices(
                results["model_forward"].model.startprob_posterior_.reshape(
                    results["model_forward"].n_states, 1
                ),
                results["model_forward"].model.transmat_posterior_,
                results["model_forward"].model.emissionprob_posterior_,
            ),
        ]
    )
    return


@app.cell
def _(results):
    for _i in range(50):
        _observed, _hidden = results["model_forward"].model.sample(3, random_state=_i)
        print([TOKENS[_i] for _i in _observed.reshape(-1)])

    # can only keep going if flat=True
    # _observed, _hidden = results['model_forward'].model.sample(3, currstate=_hidden[-1])
    # print([TOKENS[i] for i in _observed.reshape(-1)])
    return


@app.cell
def _(results):
    mo.vstack(
        [
            # before fitting
            plot_hmm_matrices(
                results["model_backward"].model.startprob_prior_.reshape(
                    results["model_backward"].n_states, 1
                ),
                results["model_backward"].model.transmat_prior_,
                results["model_backward"].model.emissionprob_prior_,
            ),
            # after fitting
            plot_hmm_matrices(
                results["model_backward"].model.startprob_.reshape(
                    results["model_backward"].n_states, 1
                ),
                results["model_backward"].model.transmat_,
                results["model_backward"].model.emissionprob_,
            ),
        ]
    )
    return


@app.cell
def _(results):
    for _i in range(50):
        _observed, _hidden = results["model_backward"].model.sample(3, random_state=_i)
        print(list(reversed([TOKENS[_i] for _i in _observed.reshape(-1)])))

    # can only keep going if flat=True
    # _observed, _hidden = results['model_backward'].model.sample(3, currstate=_hidden[-1])
    # print(list(reversed([TOKENS[i] for i in _observed.reshape(-1)])))
    return


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    ## tests to evaluate learning

    ### for forward model:
    - order of 2 actions shouldn't matter
    - TT should have higher lift likelihood than of TW (iff both agents are capable on their own)
    - TW should be sensitive to who's lifting (iff one agent is strong enough but the other isn't)
    - likelihood of fail for WW should be ~1
    - for T, lift and fail should each be impossible
    - agent shouldn't be able to take two actions in a row
    - ~~shouldn't be able to have two outcomes in a row?~~ behavior is not defined here

    > IDEA: visualize these tests/explicitly perform these tests

    ### for backward model:
    - should make smart inference about what happened in the two actions which preceded a lift or fail, based on the strengths of the agents (B_competence >= A_competence by my convention, so B is the stronger agent)
    """)
    return


@app.cell(hide_code=True)
def _(results):
    results["pred"]
    return


@app.function
def invisible(arg):
    return None


@app.function
def show_results_accuracy(results, windows):
    fwd_acc = results["model_forward"].accuracy(windows)
    bwd_acc = results["model_backward"].accuracy(
        [list(reversed(w)) for w in windows],
        backward_learning=True
    )

    # fwd_acc_b = results["model_forward"].accuracy_about_stronger_agent(windows)
    bwd_acc_b = results["model_backward"].accuracy_about_stronger_agent(
        [list(reversed(w)) for w in windows],
        backward_learning=True
    )

    acc_results = pl.concat([
        pl.DataFrame({
            'Forward': fwd_acc.reshape(-1),
            'Backward': bwd_acc.reshape(-1),
        })
        .unpivot(variable_name='model', value_name='accuracy')
        .with_columns(target=pl.lit('both agents')),

        pl.DataFrame({
            # 'Forward': fwd_acc_b.reshape(-1),
            'Backward': bwd_acc_b.reshape(-1),
        })
        .unpivot(variable_name='model', value_name='accuracy')
        .with_columns(target=pl.lit('stronger agent'))
    ])

    return mo.vstack([
        mo.hstack([
            mo.md(f"""
            accurate = sampled tokens 2 & 3 match actual, any order:  
            Forward model: {fwd_acc.mean():.1%}, Backward model: {bwd_acc.mean():.1%}
            """),
            mo.md(f"""
            accurate = matches stronger agent's token:  
            {invisible("Forward model: {np.nanmean(fwd_acc_b):.1%}")}, Backward model: {np.nanmean(bwd_acc_b):.1%}
            """)
        ]),
        sns.catplot(
            acc_results.to_pandas(),
            kind='point',
            x='model',
            y='accuracy',
            col='target'
        ).map_dataframe(
            sns.stripplot,
            x='model',
            y='accuracy',
            alpha=.1
        )
    ])


@app.cell
def _(pair, results, windows):
    mo.vstack([
        mo.md("**Within-sample accuracy**"),
        show_results_accuracy(results, windows),
        mo.md("**Out-of-sample accuracy**"),
        show_results_accuracy(results, generate_windows(pair, num_windows=500)),
    ])
    return


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    ## Comparison to null model
    """)
    return


@app.cell
def _(windows):
    null_windows = shuffle(np.array(windows).reshape(-1)).reshape(len(windows), 3)
    # null model (same model fit to windows without consistent structure)
    results_null = fit_predict_model(null_windows, max_iter=100)
    return null_windows, results_null


@app.cell(hide_code=True)
def _(results_null):
    plot_results(results_null)
    return


@app.cell(hide_code=True)
def _(results_null):
    mo.vstack(
        [
            # before fitting
            plot_hmm_matrices(
                results_null["model_forward"].model.startprob_prior_.reshape(
                    results_null["model_forward"].n_states, 1
                ),
                results_null["model_forward"].model.transmat_prior_,
                results_null["model_forward"].model.emissionprob_prior_,
            ),
            # after fitting
            plot_hmm_matrices(
                results_null["model_forward"].model.startprob_.reshape(
                    results_null["model_forward"].n_states, 1
                ),
                results_null["model_forward"].model.transmat_,
                results_null["model_forward"].model.emissionprob_,
            ),
            # before fitting
            plot_hmm_matrices(
                results_null["model_backward"].model.startprob_prior_.reshape(
                    results_null["model_backward"].n_states, 1
                ),
                results_null["model_backward"].model.transmat_prior_,
                results_null["model_backward"].model.emissionprob_prior_,
            ),
            # after fitting
            plot_hmm_matrices(
                results_null["model_backward"].model.startprob_.reshape(
                    results_null["model_backward"].n_states, 1
                ),
                results_null["model_backward"].model.transmat_,
                results_null["model_backward"].model.emissionprob_,
            ),
        ]
    )
    return


@app.cell(hide_code=True)
def _(null_windows, pair, results_null, windows):
    _new_null_windows = shuffle(np.array(null_windows).reshape(-1)).reshape(len(windows), 3)
    mo.vstack([
        mo.md("**_Null model_ Within-sample accuracy (null windows)**"),
        show_results_accuracy(results_null, null_windows),
        mo.md("**_Null model_ Out-of-sample accuracy (null windows)**"),
        show_results_accuracy(results_null, _new_null_windows),
        mo.md("**_Null model_ Out-of-sample accuracy (real windows)**"),
        show_results_accuracy(results_null, generate_windows(pair, num_windows=500))
    ])
    return


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    ## Sequence generation & model visualization code
    """)
    return


@app.function
def shuffle(arr):
    return np.random.choice(arr, size=np.shape(arr), replace=False)


@app.function
def resample(arr):
    return np.random.choice(arr, size=np.shape(arr), replace=True)


@app.function
def generate_outcome(pair, action1, action2):
    competences = {
        "A_try": pair.A_competence,  # A's competence
        "B_try": pair.B_competence,  # B's competence
        "A_wait": 0,
        "B_wait": 0,
    }
    efforts = {
        "A_try": pair.A_effort(),  # A's effort
        "B_try": pair.B_effort(),  # B's effort
        "A_wait": 0,  # doesn't matter since multiplying 0
        "B_wait": 0,  # doesn't matter since multiplying 0
    }
    action1_contribution = competences[action1] * efforts[action1]
    if action1 == action2:
        combined_contribution = action1_contribution
    else:
        action2_contribution = competences[action2] * efforts[action2]
        combined_contribution = action1_contribution + action2_contribution

    minimum_contribution = 5  # weight of 5
    success = combined_contribution >= minimum_contribution
    outcome = "lift" if success else "fail"
    return outcome


@app.function
def generate_windows(pair, num_windows):
    windows = []
    for _ in range(num_windows):
        agent1, agent2 = shuffle(["A", "B"])
        agent1_action, agent2_action = resample(["try", "wait"])
        action1 = f"{agent1}_{agent1_action}"
        action2 = f"{agent2}_{agent2_action}"
        outcome = generate_outcome(pair, action1, action2)
        current_window = [action1, action2, outcome]
        windows.append(current_window)
    return windows


@app.function
def make_null_actions_generator():
    """
    For generating random windows where the transitions between actions are random.
    Useful to test test whether the HMM needs to learn temporal contingencies between actions (e.g., that each agent takes exactly one action, and they are unrelated) to predict the outcome.
    """
    # The first two tokens in regular windows can be interpreted as a two-state sequence from a Markov chain:
    # Start -(Uniform probability)-> A_try, A_wait, B_try, B_wait
    # A_try -(50%)-> B_try
    # A_try -(50%)-> B_wait
    # B_try -(50%)-> A_try
    # B_try -(50%)-> A_wait
    # A_wait -(50%)-> B_wait
    # A_wait -(50%)-> B_try
    # B_wait -(50%)-> A_wait
    # B_wait -(50%)-> A_try
    # i.e. transition matrix is 4x4:
    # A_try A_wait B_try B_wait
    # 0  0   .5 .5
    # 0  0   .5 .5
    # .5 .5  0  0
    # .5 .5  0  0
    # The first two tokens in windows will null actions can be interpreted as a state sequence from a two-state Markov chain:
    # the transition matrix is 4x4 but randomized ~~as though it's an Erdos-Renyi random graph with p=(8/16) and weight=unif(0,1)~~
    # should it be matched on more dimensions? ideally, it should be matched on all dimensions except the one I'm testing.
    # so let's have it matched in transitions being uniform (unbiased), and also have the matrix still be row-stochastic (so that an action always leads to a second action; i.e. so that the sequence never terminates)
    # ~~and let's also have it match in having a stationary uniform distribution~~ this is actually hard, so overkill rn.

    # To do this: for each row of the matrix, select n in 1,2,3,4 and randomly select n cells in the row.
    # Then fill each selected cell with 1/n (as the transition probability)
    trans_mat = np.zeros((4, 4))
    for i in range(4):
        num_indices = np.random.randint(0, 4) + 1
        indices = np.random.choice(np.arange(4), size=num_indices, replace=False)
        trans_mat[i, indices] = 1./num_indices

    start_prob = np.ones(4) / 4

    def _generate(max_length=2):
        trans_prob = start_prob
        for _ in range(max_length):
            state = np.random.choice(np.arange(trans_prob.shape[0]), p=trans_prob)
            yield state
            trans_prob = trans_mat[state]

    return _generate, start_prob, trans_mat


@app.function
def generate_windows_with_actions_generator(pair, num_windows, generator):
    # The outcome is still dependent on the actions.
    # ~~I think I'll still allow double-counting actions to keep constant the outcome-decoding process (as a function of the actions).~~
    # Actually, we get a confound if we allow double-counting-- we just get more diverse training data, which on its own may make it easier for the HMM to learn the generate_outcome function.
    # So I modified generate_outcome to reflect this; our target generate_outcome function need not be designed to allow double-counting, anyway.
    windows = []
    for _ in range(num_windows):
        actions = list(generator(max_length=2))
        action1 = TOKENS[actions[0]]
        action2 = TOKENS[actions[1]]
        outcome = generate_outcome(pair, action1, action2)
        current_window = [action1, action2, outcome]
        windows.append(current_window)
    return windows


@app.cell(disabled=True)
def _(pair):
    _gen, _start_prob, _trans_mat = make_null_actions_generator()
    _null_windows = generate_windows_with_actions_generator(pair, 1000, _gen)

    # validating that the startprob frequences are good
    _first_states = []
    for _win in _null_windows:
        _first_states.append(_win[0])

    # validating that the ngram frequences match the transition frequencies
    _ngrams = []
    for _win in _null_windows:
        _ngrams.append(f"{_win[0]}->{_win[1]}")

    print(_start_prob)
    sns.barplot(_first_states)
    print(_trans_mat)
    sns.barplot(_ngrams)
    return


@app.function
def fit_predict_model(windows, n_states=24, max_iter=100):
    hmm_model_forward = WindowableHMM(n_states=n_states, max_iter=max_iter)
    hmm_model_backward = WindowableHMM(n_states=n_states, max_iter=max_iter)

    for window in windows:
        hmm_model_forward.add_to_history(window)
        hmm_model_backward.add_to_history(list(reversed(window)))

    hmm_model_forward.fit()
    hmm_model_backward.fit()

    pred_forward = {}

    # consider both orderings (should be the same in this case, since HMM should not learn that there's any interdependence/importance of order)
    pred_forward["TT?"] = {
        "A_try,B_try,?": hmm_model_forward.predict(["A_try", "B_try"]),
        "B_try,A_try,?": hmm_model_forward.predict(["B_try", "A_try"]),
    }

    pred_forward["TW?"] = {
        "A_try,B_wait,?": hmm_model_forward.predict(["A_try", "B_wait"]),
        "B_try,A_wait,?": hmm_model_forward.predict(["B_try", "A_wait"]),
        "B_wait,A_try,?": hmm_model_forward.predict(["B_wait", "A_try"]),
        "A_wait,B_try,?": hmm_model_forward.predict(["A_wait", "B_try"]),
    }

    pred_forward["WW?"] = {
        "A_wait,B_wait,?": hmm_model_forward.predict(["A_wait", "B_wait"]),
        "B_wait,A_wait,?": hmm_model_forward.predict(["B_wait", "A_wait"]),
    }

    pred_forward["T?"] = {
        "A_try,?": hmm_model_forward.predict(["A_try"]),
        "B_try,?": hmm_model_forward.predict(["B_try"]),
    }

    pred_backward = {}

    pred_backward["?Outcome"] = {
        "?,lift": hmm_model_backward.predict(["lift"]),
        "?,fail": hmm_model_backward.predict(["fail"]),
    }

    pred_backward["?TOutcome"] = {
        "?,A_try,lift": hmm_model_backward.predict(["lift", "A_try"]),
        "?,B_try,lift": hmm_model_backward.predict(["lift", "B_try"]),
        "?,A_try,fail": hmm_model_backward.predict(["fail", "A_try"]),
        "?,B_try,fail": hmm_model_backward.predict(["fail", "B_try"]),
    }
    pred_backward["?WOutcome"] = {
        "?,A_wait,lift": hmm_model_backward.predict(["lift", "A_wait"]),
        "?,B_wait,lift": hmm_model_backward.predict(["lift", "B_wait"]),
        "?,A_wait,fail": hmm_model_backward.predict(["fail", "A_wait"]),
        "?,B_wait,fail": hmm_model_backward.predict(["fail", "B_wait"]),
    }

    return {
        "windows": windows,
        "pred": {
            "pred_forward": pred_forward,
            "pred_backward": pred_backward,
        },
        "model_forward": hmm_model_forward,
        "model_backward": hmm_model_backward,
    }


@app.function
def plot_results(results):
    tokens = np.array(results["windows"]).reshape(-1)

    fig, axes = plt.subplots(
        3,
        1,
        figsize=(12, 7),
    )

    # Panel 1: Token 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_xlabel("Token Index")
    ax.set_title("Observed Token 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)

    # Panels 2-3: LogProb
    ax = axes[1]
    ax.plot(
        -results["model_forward"].convergence_history[1:],
        marker="o",
        linestyle="-",
        color="#3498db",
    )
    ax.set_title(
        "Forward Model Fit",
        fontsize=11,
        fontweight="bold",
    )
    ax.set_xlabel("Iteration")
    ax.set_ylabel("-logP")
    ax.set_xlim(0, 100)

    ax = axes[2]
    ax.plot(
        -results["model_backward"].convergence_history[1:],
        marker="o",
        linestyle="-",
        color="#3498db",
    )
    ax.set_title(
        "Backward Model Fit",
        fontsize=11,
        fontweight="bold",
    )
    ax.set_xlabel("Iteration")
    ax.set_ylabel("-logP")
    ax.set_xlim(0, 100)

    plt.tight_layout()

    return fig


@app.function
def plot_hmm_matrices(start_prob, trans_mat, emission_prob, order=None):
    _fig, _ax = plt.subplots(1, 3, figsize=(9, 3))

    assert np.shape(start_prob)[0] == np.shape(trans_mat)[0]
    assert np.shape(start_prob)[0] == np.shape(emission_prob)[0]
    if order is None:
        order = np.arange(0, np.shape(start_prob)[0])

    sns.heatmap(start_prob[order], vmin=0, vmax=1, cbar=False, yticklabels=order, xticklabels=False, ax=_ax[0])
    sns.heatmap(trans_mat[order], vmin=0, vmax=1, cbar=False, yticklabels=order, ax=_ax[1])
    sns.heatmap(emission_prob[order], vmin=0, vmax=1, yticklabels=order, xticklabels=TOKENS, ax=_ax[2])
    _fig.tight_layout()
    return _fig


@app.cell(hide_code=True)
def _():
    mo.md(r"""
    ---
    # Reversible HMM attempt (2/10/26)
    """)
    return


@app.cell
def _(SEED2):
    import deeptime

    class ReversibleHMM:
        """HMM that has a reversible transition matrix. Currently only models observations as one continuous sequence."""

        def __init__(self, n_states=20, flat=True):
            self.n_states = n_states
            self.token_to_idx = {t: i for i, t in enumerate(TOKENS)}
            self.history = []

            self.flat = flat

        def _tokens_to_obs(self, tokens):
            return np.array([self.token_to_idx[t] for t in tokens])

        def add_to_history(self, tokens):
            obs = self._tokens_to_obs(tokens)
            self.history.append(obs)

        def fit(self):
            if self.flat:
                sequences = np.array(self.history).reshape(-1)
            else:
                # sequences = np.array(self.history)
                raise Exception("flat=False not implemented")

            self.initial_model = deeptime.markov.hmm.init.discrete.random_guess(
                n_observation_states=len(TOKENS),
                n_hidden_states=self.n_states,
                seed=SEED2,
            )

            # Fit HMM to tokens
            self.model = deeptime.markov.hmm.MaximumLikelihoodHMM(
                self.initial_model, lagtime=1, reversible=True
            ).fit_fetch(sequences)

        def predict(self, tokens):
            # Get state distribution after input observations
            raise Exception("Not implemented")

        def _get_conditional_frequency(self, conditional_tokens, next_token):
            n = len(conditional_tokens) + 1
            target_ngram = tuple(conditional_tokens) + (next_token,)
            context = tuple(conditional_tokens)
            match_count = 0
            context_count = 0
            if self.flat:
                # treat all windows as one continuous sequence
                all_tokens = [
                    TOKENS[idx[0]] for window_obs in self.history for idx in window_obs
                ]
                for i in range(len(all_tokens) - n + 1):
                    ngram = tuple(all_tokens[i : i + n])
                    if ngram[: len(context)] == context:
                        context_count += 1
                        if ngram == target_ngram:
                            match_count += 1
            else:
                # treat each window independently, only checking the start of each window
                for window_obs in self.history:
                    window_tokens = [TOKENS[idx[0]] for idx in window_obs]
                    if (
                        len(window_tokens) >= n
                        and tuple(window_tokens[: len(context)]) == context
                    ):
                        context_count += 1
                        if tuple(window_tokens[:n]) == target_ngram:
                            match_count += 1
            return match_count / context_count if context_count > 0 else np.nan

        @property
        def convergence_history(self):
            return np.array(self.model.likelihoods)

    return (ReversibleHMM,)


@app.cell
def _():
    SEED2 = 6
    return (SEED2,)


@app.cell
def _(ReversibleHMM, SEED2, windows):
    np.random.seed(SEED2)
    hmm_model = ReversibleHMM(flat=True)

    for window in windows:
        hmm_model.add_to_history(window)

    hmm_model.fit()

    plt.plot(-hmm_model.convergence_history)
    return (hmm_model,)


@app.cell
def _(hmm_model):
    mo.vstack(
        [
            # before fitting
            plot_hmm_matrices(
                hmm_model.initial_model.initial_distribution.reshape(
                    hmm_model.n_states, 1
                ),
                hmm_model.initial_model.transition_model.transition_matrix,
                hmm_model.initial_model.output_probabilities,
            ),
            # after fitting
            plot_hmm_matrices(
                hmm_model.model.initial_distribution.reshape(hmm_model.n_states, 1),
                hmm_model.model.transition_model.transition_matrix,
                hmm_model.model.output_probabilities,
            ),
        ]
    )
    return


@app.cell
def _(hmm_model):
    _hidden, _observed = hmm_model.model.simulate(3)
    print([TOKENS[i] for i in _observed])

    _hidden, _observed = hmm_model.model.simulate(3, start=_hidden[-1])
    print([TOKENS[i] for i in _observed])

    _hidden, _observed = hmm_model.model.simulate(3, start=_hidden[-1])
    print([TOKENS[i] for i in _observed])

    _hidden, _observed = hmm_model.model.simulate(3, start=_hidden[-1])
    print([TOKENS[i] for i in _observed])

    _hidden, _observed = hmm_model.model.simulate(3, start=_hidden[-1])
    print([TOKENS[i] for i in _observed])
    return


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