"""Functions for translating a system dynamics model into a PyMC equivalent.

Because an SDM run here is a sequential simulation, the core enabling
functionality is the ``scan`` function in PyTensor:
https://pytensor.readthedocs.io/en/latest/library/scan.html
An example for how ``scan`` ties into PyMC:
https://www.pymc.io/projects/examples/en/latest/time_series/Time_Series_Generative_Graph.html

The output generated by these translations follows a template that looks
something like:

.. code-block:: python

    def step_function(*all_stocks_flows_variables):
        # this is generated by pt_sim_step() or pt_sim_step_str()

        s, ... = all_stocks_flows_variables

        # difference/update equations for each stock/flow/variable
        s_next = s + ...
        ...

        # return all updated values (pymc/pytensor has specific boilerplate to use)
        return [s_next, ...], ...

    with pymc.Model() as model:
        # define all pymc distributions/variables and initial stock/flow values
        # e.g. for stock "s":
        s_init = pm.Deterministic("s_init", ...)
        ...

        # run the scan function to generate full sequence data for simulation
        [s_seq, ...] = pytensor.scan(fn=step_function, [s_init, ...], ...)

        # create pymc containers for the full sequence data (this way they get
        # added into any output arviz datasets from pymc sampling operations)
        s = pm.Deterministic("s", pt.concatenate([s_init, s_seq], ...)

Reno's underlying math framework works fairly similarly to PyTensor, (symbolic
tree and numpy-ish) so translation is done by calling the ``pt()`` or ``pt_str()``
function that every EquationPart has, which returns the PyTensor equivalent of
that operation/thing.

Note that there are functions to directly create pymc models/operations and functions
to create strings of PyMC/PyTensor code (suffixed with _str). These exist for two reasons:

1. Debugging - figuring out errors or issues with PyMC models can be an adventure,
    and doubly so if you can't directly see or modify the model code. ``to_pymc_model_str()``
    should _in theory_ generate a string of code equivalent to what's used in
    ``to_pymc_model()``, but can be edited and run separately as needed to test
    changes.
2. Flexibility - if trying to do something beyond what Reno's PyMC naive translation
    can support, the string code output gives a starting point that can be directly
    incorporated or modified to do more advanced things.
"""

import warnings
from collections.abc import Callable

import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt

import reno


def pt_sim_step(model: "reno.model.Model") -> Callable:
    """Returns a target function for pytensor scan, this equates to a single step in the
    system dynamics model simulation. This probably shouldn't be used anywhere
    outside of the context of the pymc model creation function."""

    # general flow is to:
    # 1. run all stock update equations (which are based on previous values,
    #   being passed in through args)
    # 2. run dependency ordered non-static flow and variable equations
    #   (which are based on current values, hence the need for dependency
    #   ordering)

    def step(*args):
        original_refs = resolve_pt_scan_args(model, args)
        # keeping a non-modified version of refs to help with typechecking
        # at the end of this function
        refs = resolve_pt_scan_args(model, args)

        stock_nexts = []
        other_nexts = {}
        # other_nexts is dictionary because we need to reorder to default
        # flows + vars for returns, see note:
        # NOTE: assumed order of variables whenever they're being passed or returned
        # is order found in model.stocks + model.flows + model.vars

        for stock in model.all_stocks():
            full_stock_eq = stock._implied_eq(stock + stock.compute_diff_eq())
            stock_nexts.append(full_stock_eq.pt(**refs))

        # make sure further references to stocks (in the flows and vars) use
        # this timestep's stock computation result
        for index, stock in enumerate(model.all_stocks()):
            refs[stock.qual_name()] = stock_nexts[index]

        # dependency ordering for non-static flow/var eqs
        ref_compute_order = model.dependency_compute_order(inits_order=False)

        for obj in ref_compute_order:
            if isinstance(obj, reno.components.Stock):
                # no need to do stocks, already handled above
                continue
            if not obj.is_static():
                pt_eq = obj._implied_eq().pt(**refs)
                other_nexts[obj] = pt_eq
                refs[obj.qual_name()] = pt_eq

        # do any type checking to ensure consistency (if ints become floats without
        # an explicit conversion, PyTensor gets real unhappy)
        for index, stock in enumerate(model.all_stocks()):
            if stock_nexts[index].dtype != original_refs[stock.qual_name()].dtype:
                stock_nexts[index] = stock_nexts[index].astype(
                    original_refs[stock.qual_name()].dtype
                )
        for obj in ref_compute_order:
            if isinstance(obj, reno.components.Stock):
                # already handled in for loop above
                continue
            if not obj.is_static():
                if other_nexts[obj].dtype != original_refs[obj.qual_name()].dtype:
                    other_nexts[obj] = other_nexts[obj].astype(
                        original_refs[obj.qual_name()].dtype
                    )

        # order returns according to model.stocks + model.flows + model.vars
        returns = stock_nexts
        for obj in model.all_flows() + model.all_vars():
            if obj in other_nexts:
                returns.append(other_nexts[obj])

        return list(returns), pm.pytensorf.collect_default_updates(
            inputs=args, outputs=returns
        )

    return step


def pt_sim_step_imports() -> str:
    """Get the string of code for the imports needed for ``pt_sim_step_str()`` to
    run correctly."""
    import_strs = [
        "import pytensor.tensor as pt",
        "from pytensor.ifelse import ifelse",
        "import pymc as pm",
    ]
    return "\n".join(import_strs)


def pt_sim_step_str(model: "reno.model.Model") -> str:
    """Construct a string of python code for the step function for this model that could
    be used by a pytensor scan operation. Should be a functional (string) equivalent of
    the ``pt_sim_step()`` function.

    Expected imports for the resulting code to run:
        >>> import pytensor.tensor as pt
        >>> from pytensor.ifelse import ifelse
        >>> import pymc as pm

    (use ``pt_sim_step_imports()`` to get these programatically)

    Note that this code is in no way PEP8 formatted, for complex models some of the lines
    will be exceptionally long.
    """
    name = "step" if model.name is None else f"{model.name}_step"
    code = f"def {name}(*args):\n"
    code += f"\t{', '.join(expected_pt_scan_arg_names(model))} = args\n"

    refs_dict = expected_pt_scan_arg_names(model, as_dict=True)
    # (in pt_str calls, variables passed in from args should simply be substituted
    # with the variable name itself)

    code += "\n\t# Difference/recurrence equations\n"

    for stock in model.all_stocks():
        code += (
            f"\t{stock.qual_name()}_next = {stock._implied_eq().pt_str(**refs_dict)}\n"
        )

    # update all stock names to _next for use in flow/var updates
    # (we don't do this above on the off chance we somehow end up with a
    # stock being able to directly ref another stock? We'd want it to still
    # be based on old value at that point)
    for stock in model.all_stocks():
        refs_dict[stock.qual_name()] = f"{stock.qual_name()}_next"

    # dependency ordering for non-static flow/var eqs
    ref_compute_order = model.dependency_compute_order(inits_order=False)
    for obj in ref_compute_order:
        if isinstance(obj, reno.components.Stock):
            # no need to do stocks, already handled above
            continue
        if not obj.is_static():
            code += (
                f"\t{obj.qual_name()}_next = {obj._implied_eq().pt_str(**refs_dict)}\n"
            )
            refs_dict[obj.qual_name()] = f"{obj.qual_name()}_next"

    # type checking
    code += "\n\t# Type checks\n"
    for obj in ref_compute_order:
        if not obj.is_static():
            code += f"\t{obj.qual_name()}_next = {obj.qual_name()}_next.astype({obj.qual_name()}.dtype) if {obj.qual_name()}_next.dtype != {obj.qual_name()}.dtype else {obj.qual_name()}_next\n"

    # order returns according to model.stocks + model.flows + model.vars
    return_names = ", ".join(
        [f"{obj.qual_name()}_next" for obj in model.all_refs() if not obj.is_static()]
    )

    code += f"\n\treturn [{return_names}], pm.pytensorf.collect_default_updates(inputs=args, outputs=[{return_names}])\n"

    return code


def to_pymc_model(
    model: "reno.model.Model", observations: list["reno.ops.Observation"] = None
) -> pm.model.core.Model:
    """Generate a pymc model for bayesian analysis of this system dynamics model. The general
    idea is that this creates corresponding pymc variables (or distributions as relevant) for
    each stock/flow/var in the model, and sets up the full simulation sequence computations
    based on the generated step function from ``pt_sim_step()``.

    Sampling with priors should be equivalent to running the system dynamics model normally
    (this is essentially "forward simulation mode".) Add observations to the pymc model
    variables and sample from posterior to run bayesian analysis/determine how distributions
    of any other variables may be affected.
    """

    # general flow is to:
    # 1. set up all initial value computations in pymc deterministics (or
    #   appropriate distributions for reno.ops.distribution values) Note
    #   that _all_ initial computations need to be dependency ordered
    # 2. apply the pytensor scan function targeting ``pt_sim_step()``, any
    #   static references are passed as non-sequences, time is passed as
    #   a sequence, and all dynamic refs (stocks, flows, vars) are passed
    #   as outputs info, representing the initial values for each sequence
    #   and the set of expected return sequences
    # 3. create a deterministic with dim=timesteps for each output sequence
    #   from the scan operation.

    timesteps = model.steps
    coords = {"t": range(timesteps)}
    with pm.Model(coords=coords) as m:
        ref_compute_order = model.dependency_compute_order(inits_order=True)

        refs = {}
        inits_by_obj = {}
        # same logic as pt_sim_step for these being dictionaries, ordering is tricky with these,
        # we always assume arg/return orders are in default stocks + flows + vars order
        statics_by_obj = {}
        # note that we separate "statics" because the pytensor scan
        # equivalent is to pass them as "non-sequences" (they're passed in a
        # separate portion of the *args, and assumed to not update/not
        # requiring a return)
        hist_vars_by_obj = {}
        # any variables that require special treatment for historical
        # values have their starting historical array deterministics
        # stored in here

        refs["__PT_SEQ_LEN__"] = pt.as_tensor(model.steps)
        # there are a few operations that need to know the total length
        # of the simulation/sequence in order to create the correct
        # pytensor equivalent output, see ops.py

        historical_ref_taps = find_historical_tracked_refs(model)

        # have an initial t=0 variable if necessary
        if model.find_timeref_name() is not None:
            refs[model.find_timeref_name()] = pt.as_tensor(0)

        # create symbolic computations for all initial values for stocks/flows/vars
        for obj in ref_compute_order:
            if isinstance(obj, reno.components.Stock) and obj.init is None:
                pt_eq = pm.Deterministic(
                    f"{obj.qual_name()}_init", pt.as_tensor(np.array(0.0))
                )
                inits_by_obj[obj] = pt_eq
                refs[obj.qual_name()] = pt_eq
            elif obj.is_static():
                if isinstance(obj.eq, reno.components.Distribution):
                    pt_eq = obj.eq.pt(__PTNAME__=obj.qual_name(), **refs)
                else:
                    pt_eq = pm.Deterministic(
                        obj.qual_name(), obj._implied_eq(obj.eq).pt(**refs)
                    )
                statics_by_obj[obj] = pt_eq
                refs[obj.qual_name()] = pt_eq
            elif obj.init is not None:
                # stocks, flows, vars with a defined init
                if isinstance(obj.init, reno.components.Distribution):
                    pt_eq = obj.init.pt(__PTNAME__=f"{obj.qual_name()}_init", **refs)
                else:
                    pt_eq = pm.Deterministic(
                        f"{obj.qual_name()}_init",
                        obj._implied_eq(obj.init).pt(**refs),
                    )
                inits_by_obj[obj] = pt_eq
                refs[obj.qual_name()] = pt_eq
            else:
                # flows and vars without a separate init definition
                # (just run a single normal calculation of their equation)
                pt_eq = pm.Deterministic(
                    f"{obj.qual_name()}_init", obj._implied_eq(obj.eq).pt(**refs)
                )
                inits_by_obj[obj] = pt_eq
                refs[obj.qual_name()] = pt_eq

            # handle initial historical arrays if necessary (but not ones with just -1)
            if obj in historical_ref_taps and len(historical_ref_taps[obj]) > 1:
                # TODO: doesn't handle static but still unclear if that
                # makes sense anyway
                inner_array = [0.0] * (min(historical_ref_taps[obj]) * -1) + [
                    inits_by_obj[obj]
                ]
                hist_var = pm.Deterministic(
                    f"{obj.qual_name()}_hist", pt.stack(inner_array)
                )
                hist_vars_by_obj[obj] = hist_var

        # set up a timestep sequence to pass to scan if any equations are using a TimeRef
        # e.g. t = [1, 2, 3, 4, 5, 6, 7, 8, 9]
        sequences = []
        if model.find_timeref_name() is not None:
            timestep_seq = pt.as_tensor(np.arange(1, timesteps))
            sequences = [timestep_seq]

        # ensure all outputs_info are in correct order
        # NOTE: assumed order of variables whenever they're being passed or returned
        # is order found in model.stocks + model.flows + model.vars
        inits = []
        for obj in model.all_refs():
            if obj in hist_vars_by_obj:
                inits.append(
                    dict(initial=hist_vars_by_obj[obj], taps=historical_ref_taps[obj])
                )
            elif obj in inits_by_obj:
                inits.append(inits_by_obj[obj])

        statics = []  # statics get passed into scan target as non_sequences
        for obj in model.all_flows() + model.all_vars():
            if obj in statics_by_obj:
                statics.append(statics_by_obj[obj])

        seqs, updates = pytensor.scan(
            fn=pt_sim_step(model),
            sequences=sequences,
            non_sequences=statics,
            outputs_info=inits,
            strict=True,
            n_steps=model.steps - 1,
        )

        # handle weird error when only a single stock returned (not wrapping
        # in a list correctly for some reason)
        if len(inits) == 1:
            seqs = [seqs]

        # assign the fully simulated sequences of all non-static refs to deterministics
        # (order is as above, always model.stocks + model.flows + model.vars)
        for i, obj in enumerate(
            model.all_stocks()
            + [
                obj
                for obj in model.all_flows() + model.all_vars()
                if not obj.is_static()
            ]
        ):
            full_seq = pm.Deterministic(
                obj.qual_name(),
                pt.concatenate([pt.as_tensor([inits_by_obj[obj]]), seqs[i].flatten()]),
                dims="t",
            )

            # set pt ref to full sequence for any subsequent metric equations
            refs[obj.qual_name()] = full_seq

        # dependency sort metric equations
        metrics = reno.utils.dependency_compute_order([*model.all_metrics()])
        for metric in metrics:
            metric_var = pm.Deterministic(metric.qual_name(), metric.eq.pt(**refs))
            refs[metric.qual_name()] = metric_var

    # add any likelihoods + observations if relevant
    if observations is not None:
        for observation in observations:
            observation.add_tensors(m)

    return m


def pymc_model_imports() -> str:
    """Get the string of code for the imports needed for ``to_pymc_model_str()`` to
    run correctly."""
    import_strs = [
        "import pytensor",
        "import pytensor.tensor as pt",
        "from pytensor.ifelse import ifelse",
        "import pymc as pm",
        "import numpy as np",
    ]
    return "\n".join(import_strs)


# TODO: wrap in function option (string function would return model)
# TODO: option to add in necessary imports
# TODO: can you run black formatting programmatically on a string?
def to_pymc_model_str(
    model: "reno.model.Model", observations: list[reno.ops.Observation] = None
) -> str:
    """Construct a string of python code to create a pymc model wrapping this system dynamics
    model. Should be a functional (string) equivalent of the ``to_pymc_model()`` function. Includes
    the output from ``pt_sim_step_str()``.

    Expected imports for the resulting code to run:
        >>> import pytensor
        >>> import pytensor.tensor as pt
        >>> from pytensor.ifelse import ifelse
        >>> import pymc as pm
        >>> import numpy as np

    (Alternatively get the string code for these imports with ``pymc_model_imports()``)
    """
    name = "pymc_m" if model.name is None else f"{model.name}_pymc_m"
    step_name = "step" if model.name is None else f"{model.name}_step"
    code = pt_sim_step_str(model) + "\n"
    code += f'coords = {{"t": range({model.steps})}}\n'
    code += f"with pm.Model(coords=coords) as {name}:\n"

    historical_ref_taps = find_historical_tracked_refs(model)

    initial_refs_dict = {}
    if model.find_timeref_name() is not None:
        initial_refs_dict[model.find_timeref_name()] = "pt.as_tensor(0)"

    initial_refs_dict["__PT_SEQ_LEN__"] = f"pt.as_tensor({model.steps})"
    # there are a few operations that need to know the total length
    # of the simulation/sequence in order to create the correct
    # pytensor equivalent output, see ops.py

    ref_compute_order = model.dependency_compute_order(inits_order=True)
    code += "\t# Initial values/timestep 0 equations\n"
    for obj in ref_compute_order:
        if isinstance(obj, reno.components.Stock) and obj.init is None:
            inner_eq = "pt.as_tensor(np.array(0.0))"
            code += f'\t{obj.qual_name()}_init = pm.Deterministic("{obj.qual_name()}_init", {inner_eq})\n'
            initial_refs_dict[obj.qual_name()] = f"{obj.qual_name()}_init"
        elif obj.is_static():
            if isinstance(obj.eq, reno.components.Distribution):
                inner_eq = obj.eq.pt_str(
                    __PTNAME__=obj.qual_name(), **initial_refs_dict
                )
                code += f"\t{obj.qual_name()} = {inner_eq}\n"
            else:
                inner_eq = obj._implied_eq(obj.eq).pt_str(**initial_refs_dict)
                code += f'\t{obj.qual_name()} = pm.Deterministic("{obj.qual_name()}", {inner_eq})\n'
            initial_refs_dict[obj.qual_name()] = obj.qual_name()
        elif obj.init is not None:
            # stocks, flows, vars with a defined init
            if isinstance(obj.init, reno.components.Distribution):
                inner_eq = obj.init.pt_str(
                    __PTNAME__=obj.qual_name() + "_init", **initial_refs_dict
                )
                code += f"\t{obj.qual_name()}_init = {inner_eq}\n"
            else:
                inner_eq = obj._implied_eq(obj.init).pt_str(**initial_refs_dict)
                code += f'\t{obj.qual_name()}_init = pm.Deterministic("{obj.qual_name()}_init", {inner_eq})\n'
            initial_refs_dict[obj.qual_name()] = f"{obj.qual_name()}_init"
        else:
            # flows and vars without a separate init definition
            inner_eq = obj._implied_eq(obj.eq).pt_str(**initial_refs_dict)
            code += f'\t{obj.qual_name()}_init = pm.Deterministic("{obj.qual_name()}_init", {inner_eq})\n'
            initial_refs_dict[obj.qual_name()] = f"{obj.qual_name()}_init"

        # handle initial historical arrays if necessary (but not ones with just -1)
        if obj in historical_ref_taps and len(historical_ref_taps[obj]) > 1:
            # TODO: doesn't handle static but still unclear if that makes sense anyway
            inner_array = (
                ", ".join(["0.0"] * (min(historical_ref_taps[obj]) * -1))
                + f", {obj.qual_name()}_init"
            )
            inner_eq = f"pt.stack([{inner_array}])"
            code += f'\t{obj.qual_name()}_hist = pm.Deterministic("{obj.qual_name()}_hist", {inner_eq})\n'

    # add timeseries
    if model.find_timeref_name() is not None:
        code += f"\ttimestep_seq = pt.as_tensor(np.arange(1, {model.steps}))\n"

    # call the scan function
    return_names = ", ".join(
        [f"{obj.qual_name()}_seq" for obj in model.all_refs() if not obj.is_static()]
    )
    if (
        len(
            model.stocks
            + [
                obj
                for obj in model.all_flows() + model.all_vars()
                if not obj.is_static()
            ]
        )
        > 1
    ):
        return_names = f"[{return_names}]"

    # if we don't have to worry about historical values, the outputs_info is
    # just the array of inits
    if len(historical_ref_taps) == 0:
        init_names = ", ".join(
            [
                f"{obj.qual_name()}_init"
                for obj in model.all_refs()
                if not obj.is_static()
            ]
        )
    else:
        init_names = []
        for obj in model.all_refs():
            # TODO: (2025.02.11) there's some mixed up logic here/not
            # handling if we have a historical value on a static variable
            # (which doesn't truely make sense but also shouldn't be
            # explicitly disallowed?)
            if obj in historical_ref_taps and len(historical_ref_taps[obj]) > 1:
                taps = historical_ref_taps[obj]
                init_names.append(f"dict(initial={obj.qual_name()}_hist, taps={taps})")
            elif not obj.is_static():
                init_names.append(f"{obj.qual_name()}_init")
        init_names = ", ".join(init_names)
    code += "\n\t# Run autoregressive step function/timestep-wise updates to fill sequences\n"
    code += f"\t{return_names}, updates = pytensor.scan(\n"
    code += f"\t\tfn={step_name},\n"
    if model.find_timeref_name() is not None:
        code += "\t\tsequences=[timestep_seq],\n"
    code += f"\t\tnon_sequences=[{', '.join([obj.qual_name() for obj in model.all_flows() + model.all_vars() if reno.utils.is_static(obj)])}],\n"
    code += f"\t\toutputs_info=[{init_names}],\n"
    code += "\t\tstrict=True,\n"
    code += f"\t\tn_steps={model.steps - 1}\n"
    code += "\t)\n"

    code += "\n\t# Collect full sequence data for all stocks/flows/vars into pymc variables\n"
    for obj in model.all_stocks() + [
        obj for obj in model.all_flows() + model.all_vars() if not obj.is_static()
    ]:
        code += f'\t{obj.qual_name()} = pm.Deterministic("{obj.qual_name()}", pt.concatenate([pt.as_tensor([{obj.qual_name()}_init]), {obj.qual_name()}_seq.flatten()]), dims="t")\n'

        # set pt ref to full sequence for any metrics
        initial_refs_dict[obj.qual_name()] = obj.qual_name()

    # dependency sort metric equations
    metrics = reno.utils.dependency_compute_order([*model.metrics])
    for metric in metrics:
        code += f'\t{metric.qual_name()} = pm.Deterministic("{metric.qual_name()}", {metric.eq.pt_str(**initial_refs_dict)})\n'
        initial_refs_dict[metric.qual_name()] = metric.qual_name()

    # add any likelihoods + observations if relevant
    if observations is not None:
        for observation in observations:
            code += f"\t{observation.ref.qual_name()}_likelihood = {observation.pt_str(**initial_refs_dict)}\n"

    return code


def find_historical_tracked_refs(
    model: "reno.model.Model",
) -> dict[reno.components.TrackedReference, list[int]]:
    """Get a dictionary of all tracked references in the model used in historical values,
    and get the associated tap list for each. (Taps is what PyTensor calls indices
    for previous sequence values, so you this is used for when an equation calls for
    some stock value at t-5, etc.)

    This is primarily used for population of pymc scan step and scan arguments."""

    historical_values = []
    for ref in model.all_refs():
        historical_values.extend(ref.find_refs_of_type(reno.components.HistoricalValue))

    ref_taps = {}
    for val in historical_values:
        # for now we only directly support a `t - some_static_amount`
        # NOTE: expanding this would likely require always passing in the full
        # time series, much more annoying boilerplate I suspect.
        if not reno.utils.check_for_easy_static_time_eq(val.index_eq):
            warnings.warn(
                f"Historical value {val.qual_name()} has a non-basic time equation which isn't supported in pymc model conversion yet. Please index with a 't - [static_eq]' equation."
            )
        if val.tracked_ref not in ref_taps:
            ref_taps[val.tracked_ref] = []

        index_offset = val.index_eq.eval(t=0, force=True)
        ref_taps[val.tracked_ref].append(index_offset)

    # remove duplicate taps, ensure -1 is in there (we always want last val
    # of every ref in case it's used in a stock eq)
    ref_taps = {ref: list(set(taps + [-1])) for ref, taps in ref_taps.items()}

    # sort each taps list (order of taps is order they're passed to scan target
    # function, so we should keep consistent sorted ordering for the sanity of everyone)
    for ref, taps in ref_taps.items():
        taps.sort()

    return ref_taps


def expected_pt_scan_arg_names(
    model: "reno.model.Model", as_dict: bool = False, tap_1_only: bool = False
) -> list[str] | dict[str, str]:
    """Pytensor scan target functions are called with a very specific ordering of function
    arguments (see https://pytensor.readthedocs.io/en/latest/library/scan.html).

    Assuming a scan function signature of ``def scan(*args):``, this function
    returns a list of the appropriate reference names to put on the left side of
    a ``refname1, refname2, ... = args`` line.

    Note that with as_dict, it returns a dictionary of reference strings. This is used directly
    by pt_sim_step_str, but also includes refs for tap -1 references (to avoid being updated to
    the "current" timestep) so this is useful for pt_sim_step too

    so with as_dict False:
        [t, v0, v1, v2_h2, v2, ...]
    with as_dict True, tap_1_only False:
        {"t": "t", "v0": "v0", "v1_h1": "v1", "v1": "v1", "v2_h2": "v2_h2", "v2_h1": "v2", "v2": "v2", ...}
    with as_dict True, tap_1_only True:
        {"v1_h1": "v1", "v2_h1": "v2"}
    """

    # remember ordering is sequences/seq taps, output taps, non-sequence args
    names = []
    names_as_dict = {}
    # initially only filled with the mirrored -1 taps, if actually returning
    # as dict, extend with zipped names

    historical_ref_taps = find_historical_tracked_refs(model)

    # -----------------------
    # SEQUENCES/SEQUENCE TAPS
    # -----------------------

    timeref_name = model.find_timeref_name()
    if timeref_name is not None:
        names.append(timeref_name)

    # -----------------------
    # OUTPUT TAPS (non-static references)
    # -----------------------

    # get last of each stock
    for stock in model.all_stocks():
        if stock in historical_ref_taps:
            for tap in historical_ref_taps[stock]:
                str_tap = str(tap * -1)
                name = f"{stock.qual_name()}_h{str_tap}"
                if str_tap == "1":
                    # special handling for -1 tap
                    names_as_dict[name] = stock.qual_name()
                else:
                    names.append(name)
        names.append(stock.qual_name())

    for flow in model.all_flows():
        if not flow.is_static():
            if flow in historical_ref_taps:
                for tap in historical_ref_taps[flow]:
                    str_tap = str(tap * -1)
                    name = f"{flow.qual_name()}_h{str_tap}"
                    if str_tap == "1":
                        # special handling for -1 tap
                        names_as_dict[name] = flow.qual_name()
                    else:
                        names.append(name)
            names.append(flow.qual_name())

    for var in model.all_vars():
        if not var.is_static():
            if var in historical_ref_taps:
                for tap in historical_ref_taps[var]:
                    str_tap = str(tap * -1)
                    name = f"{var.qual_name()}_h{str_tap}"
                    if str_tap == "1":
                        # special handling for -1 tap
                        names_as_dict[name] = var.qual_name()
                    else:
                        names.append(name)
            names.append(var.qual_name())

    # -----------------------
    # NON-SEQUENCE ARGS (static references)
    # -----------------------

    for flow in model.all_flows():
        if flow.is_static():
            names.append(flow.qual_name())

    for var in model.all_vars():
        if var.is_static():
            names.append(var.qual_name())

    # NOTE: assumed order of variables whenever they're being passed or returned
    # is order found in model.stocks + model.flows + model.vars

    if as_dict and not tap_1_only:
        names_as_dict.update(dict(zip(names, names)))
        return names_as_dict
    elif as_dict and tap_1_only:
        return names_as_dict
    return names


def resolve_pt_scan_args(
    model: "reno.model.Model", args: list[pt.TensorVariable]
) -> dict[str, pt.TensorVariable]:
    """Take the set of arguments passed by pytensor to a target scan function,
    and turn it into a dictionary for passing into reno math components' `pt()`
    calls.

    (This is necessary to ensure the returned pytensor symbolic operations
    are referring to the correct tensorvariables passed into scan.)
    """
    # remember ordering is sequences/seq taps, output taps, non-sequence args
    refs = {}

    arg_names = expected_pt_scan_arg_names(model)
    for i, name in enumerate(arg_names):
        refs[name] = args[i]

    # handle mirrored -1 taps (v1_h1 == v1)
    tap1_mirrors = expected_pt_scan_arg_names(model, True, True)
    for hist1_name, target_name in tap1_mirrors.items():
        refs[hist1_name] = refs[target_name]

    return refs
