"""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.
"""

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", steps: int) -> 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, steps)
        # keeping a non-modified version of refs to help with typechecking
        # at the end of this function
        refs = resolve_pt_scan_args(model, args, steps)
        refs["__PT_SEQ_LEN__"] = steps
        refs["__TIMESERIES__"] = "scan"

        historical_refs_from_tracked = find_tracked_ref_historical_ref(model)

        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]

            # ensure a full history correctly has the new value at the end.
            if stock in historical_refs_from_tracked:
                refs[f"{stock.qual_name()}_h"] = pt.concatenate(
                    [
                        refs[f"{stock.qual_name()}_h"][1:],
                        pt.as_tensor([stock_nexts[index]]),
                    ]
                )

        # dependency ordering for non-static flow/var eqs
        ref_compute_order = model.dependency_compute_order(inits_order=False)
        # find all the per-timestep distributions because we have to skip those in
        # several places
        per_timestep_dist_objs = []
        for obj in ref_compute_order:
            if (
                isinstance(obj, reno.components.Variable)
                and isinstance(obj.eq, reno.components.Distribution)
                and obj.eq.per_timestep
            ):
                per_timestep_dist_objs.append(obj)

        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() and obj not in per_timestep_dist_objs:
                pt_eq = obj._implied_eq().pt(**refs)
                other_nexts[obj] = pt_eq
                refs[obj.qual_name()] = pt_eq

                # ensure a full history correctly has the new value at the end.
                if obj in historical_refs_from_tracked:
                    refs[f"{obj.qual_name()}_h"] = pt.concatenate(
                        [refs[f"{obj.qual_name()}_h"][1:], pt.as_tensor([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() and obj not in per_timestep_dist_objs:
                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",
        "import pymc as pm",
    ]
    return "\n".join(import_strs)


def pt_sim_step_str(model: "reno.model.Model", steps: int) -> 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
        >>> 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"
    for ref_name, arg_indices in expected_pt_scan_arg_names(
        model, max_taps=steps
    ).items():
        if isinstance(arg_indices, slice):
            code += f"\t{ref_name} = pt.stack(args[{arg_indices}], axis=0)\n"
        else:
            code += f"\t{ref_name} = args[{arg_indices}]\n"

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

    refs_dict["__TIMESERIES__"] = "scan"

    historical_refs_from_tracked = find_tracked_ref_historical_ref(model)

    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"

        # ensure a full history correctly has the new value at the end
        if stock in historical_refs_from_tracked:
            code += f"\t{stock.qual_name()}_h = pt.concatenate([{stock.qual_name()}_h[1:], pt.as_tensor([{stock.qual_name()}_next])])\n"
            # refs_dict[f"{stock.qual_name()}_h"] = f"pt.concatenate([{refs_dict[f'{stock.qual_name()}_h']}, pt.as_tensor([{f'{stock.qual_name()}_next'}])"

    # dependency ordering for non-static flow/var eqs
    ref_compute_order = model.dependency_compute_order(inits_order=False)
    # find all the per-timestep distributions because we have to skip those in
    # several places
    per_timestep_dist_objs = []
    for obj in ref_compute_order:
        if (
            isinstance(obj, reno.components.Variable)
            and isinstance(obj.eq, reno.components.Distribution)
            and obj.eq.per_timestep
        ):
            per_timestep_dist_objs.append(obj)
    for obj in ref_compute_order:
        if isinstance(obj, reno.components.Stock):
            # no need to do stocks, already handled above
            continue
        # we also don't compute statics or per-timestep distributions
        if not obj.is_static() and obj not in per_timestep_dist_objs:
            pt_eq = obj._implied_eq().pt_str(**refs_dict)
            code += f"\t{obj.qual_name()}_next = {pt_eq}\n"
            refs_dict[obj.qual_name()] = f"{obj.qual_name()}_next"

            # ensure a full history correctly has the new value at the end
            if obj in historical_refs_from_tracked:
                code += f"\t{obj.qual_name()}_h = pt.concatenate([{obj.qual_name()}_h[1:], pt.as_tensor([{obj.qual_name()}_next])])\n"
                # refs_dict[f"{obj.qual_name()}_h"] = f"pt.concatenate([{refs_dict[f'{obj.qual_name()}_h']}, pt.as_tensor([{pt_eq}])"

    # type checking
    code += "\n\t# Type checks\n"
    for obj in ref_compute_order:
        if not obj.is_static() and obj not in per_timestep_dist_objs:
            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() and obj not in per_timestep_dist_objs
        ]
    )

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

    return code.replace("\t", "    ")


def to_pymc_model(
    model: "reno.model.Model",
    observations: list["reno.ops.Observation"] = None,
    steps: int = 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.
    """
    model._find_all_extended_op_implicit_components()

    # 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.

    if steps is None:
        steps = model.steps
    coords = {"t": range(steps)}
    for obj in model.all_refs():
        if obj.shape > 1:
            coords[obj.qual_name() + "_vec"] = range(obj.shape)
    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
        per_step_dists_by_obj = {}
        # variables that hold a distribution that's per-timestep get special
        # treatment too - they're already fully computed after initialization
        # and just need to be passed in to the "sequences" portion of the scan

        # refs["__PT_SEQ_LEN__"] = pt.as_tensor(steps)
        refs["__PT_SEQ_LEN__"] = 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

        refs["__TIMESERIES__"] = "init"
        # tell any timeseries operations to refer to the variable "_h" name.

        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:
            try:
                dims = None
                if obj.shape > 1:
                    dims = f"{obj.qual_name()}_vec"
                if isinstance(obj, reno.components.Stock) and obj.init is None:
                    if obj.shape == 1:
                        pt_eq = pm.Deterministic(
                            f"{obj.qual_name()}_init", pt.as_tensor(np.array(0.0))
                        )
                    else:
                        pt_eq = pm.Deterministic(
                            f"{obj.qual_name()}_init",
                            pt.as_tensor(np.array([0.0] * obj.shape)),
                            dims=dims,
                        )
                    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(),
                            __DIM__=obj.shape,
                            __DIMNAME__=obj.qual_name() + "_vec",
                            **refs,
                        )
                    # NOTE: slice has no shape and has to assume dim is specified
                    # correctly
                    elif obj.shape > obj.eq.shape and not isinstance(
                        obj.eq, reno.ops.slice
                    ):
                        pt_eq = pm.Deterministic(
                            obj.qual_name(),
                            obj._implied_eq(obj.eq).pt(**refs).repeat(obj.shape),
                            dims=dims,
                        )
                    else:
                        pt_eq = pm.Deterministic(
                            obj.qual_name(),
                            obj._implied_eq(obj.eq).pt(**refs),
                            dims=dims,
                        )
                    statics_by_obj[obj] = pt_eq
                    refs[obj.qual_name()] = pt_eq
                elif (
                    isinstance(obj, reno.components.Variable)
                    and isinstance(obj.eq, reno.components.Distribution)
                    and obj.eq.per_timestep
                ):
                    pt_eq = obj.eq.pt(
                        __PTNAME__=obj.qual_name(),
                        __DIM__=obj.shape,
                        __DIMNAME__=obj.qual_name() + "_vec",
                        **refs,
                    )
                    per_step_dists_by_obj[obj] = pt_eq
                    refs[obj.qual_name()] = pt_eq[0]
                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",
                            __DIM__=obj.shape,
                            __DIMNAME__=obj.qual_name() + "_vec",
                            **refs,
                        )
                    else:
                        # NOTE: slice has no shape and has to assume dim is specified
                        # correctly
                        if obj.shape > obj.init.shape and not isinstance(
                            obj.init, reno.ops.slice
                        ):
                            pt_eq = pm.Deterministic(
                                f"{obj.qual_name()}_init",
                                obj._implied_eq(obj.init).pt(**refs).repeat(obj.shape),
                                dims=dims,
                            )
                        else:
                            pt_eq = pm.Deterministic(
                                f"{obj.qual_name()}_init",
                                obj._implied_eq(obj.init).pt(**refs),
                                dims=dims,
                            )
                    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)
                    # NOTE: slice has no shape and has to assume dim is specified
                    # correctly
                    if obj.shape > obj.eq.shape and not isinstance(
                        obj.eq, reno.ops.slice
                    ):
                        pt_eq = pm.Deterministic(
                            f"{obj.qual_name()}_init",
                            obj._implied_eq(obj.eq).pt(**refs).repeat(obj.shape),
                            dims=dims,
                        )
                    else:
                        pt_eq = pm.Deterministic(
                            f"{obj.qual_name()}_init",
                            obj._implied_eq(obj.eq).pt(**refs),
                            dims=dims,
                        )
                    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
                    or None in historical_ref_taps[obj]
                ):
                    # static "historical references" get constructed on the fly
                    # because it wouldn't make sense to have a separate
                    # outputs_info and handling for a thing that doesn't change
                    if obj.is_static():
                        continue
                    inner_array_value = 0.0
                    if obj.shape > 1:
                        inner_array_value = [0.0] * obj.shape
                    history_size = 0
                    if None in historical_ref_taps[obj]:
                        history_size = steps - 1
                    else:
                        history_size = min(historical_ref_taps[obj]) * -1

                    # need to separately check for an inited dynamic variable vs
                    # static because a timeseries of a static won't show up in
                    # inits_by_obj
                    inner_array_final_value = None
                    if obj in inits_by_obj:
                        inner_array_final_value = inits_by_obj[obj]
                    elif obj in statics_by_obj:
                        inner_array_final_value = statics_by_obj[obj]

                    inner_array = [inner_array_value] * history_size + [
                        inner_array_final_value
                    ]
                    hist_var = pm.Deterministic(
                        f"{obj.qual_name()}_hist", pt.stack(inner_array)
                    )
                    hist_vars_by_obj[obj] = hist_var
                    if None in historical_ref_taps[obj]:
                        refs[f"{obj.qual_name()}_h"] = hist_var
            except Exception as e:
                e.add_note(
                    f"Was trying to set up initial PYMC equation for {obj.qual_name()}"
                )
                e.add_note(f"\tEquation: {obj.eq}")
                e.add_note(f"\tType: {obj.dtype}")
                e.add_note(f"\tShape: {obj.shape}")
                raise

        # 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, steps))
            sequences.append(timestep_seq)
        for obj in per_step_dists_by_obj:
            sequences.append(per_step_dists_by_obj[obj])

        # 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:
                taps = historical_ref_taps[obj]
                if None in taps:
                    # handle passing the full timestep range
                    taps = list(range(-steps, 0))
                inits.append(dict(initial=hist_vars_by_obj[obj], taps=taps))
            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, steps),
            sequences=sequences,
            non_sequences=statics,
            outputs_info=inits,
            strict=True,
            n_steps=steps - 1,
        )

        refs["__TIMESERIES__"] = "metric"
        # at this point, any timeseries refs will be in metric equations which
        # should _not_ go based on the historical variable

        # 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() and obj not in per_step_dists_by_obj
            ]
        ):
            dims = "t"
            if obj.shape > 1:
                dims = ("t", f"{obj.qual_name()}_vec")
            full_seq = pm.Deterministic(
                obj.qual_name(),
                pt.concatenate([pt.as_tensor([inits_by_obj[obj]]), seqs[i]]),
                dims=dims,
            )

            # 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",
        "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,
    steps: int = 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
        >>> import pymc as pm
        >>> import numpy as np

    (Alternatively get the string code for these imports with ``pymc_model_imports()``)
    """
    model._find_all_extended_op_implicit_components()

    if steps is None:
        steps = model.steps

    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, steps) + "\n"
    code += f'coords = {{\n\t"t": range({steps}),\n'
    for obj in model.all_refs():
        if obj.shape > 1:
            code += f'\t"{obj.qual_name()}_vec": range({obj.shape}),\n'
    code += "}\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__"] = 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

    initial_refs_dict["__TIMESERIES__"] = "init"
    # tell any timeseries operations to refer to the variable "_hist" name.

    per_step_dist_names = []
    # use these in passing seqeunces to scan and ignore them in the others

    ref_compute_order = model.dependency_compute_order(inits_order=True)
    code += "\t# Initial values/timestep 0 equations\n"
    for obj in ref_compute_order:
        try:
            dims_text = ""
            if obj.shape > 1:
                dims_text = f', dims="{obj.qual_name()}_vec"'

            if isinstance(obj, reno.components.Stock) and obj.init is None:
                if obj.shape == 1:
                    inner_eq = "pt.as_tensor(np.array(0.0))"
                else:
                    inner_eq = f"pt.as_tensor(np.array([0.0]*{obj.shape}))"
                code += f'\t{obj.qual_name()}_init = pm.Deterministic("{obj.qual_name()}_init", {inner_eq}{dims_text})\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(),
                        __DIM__=obj.shape,
                        __DIMNAME__=obj.qual_name() + "_vec",
                        **initial_refs_dict,
                    )
                    code += f"\t{obj.qual_name()} = {inner_eq}\n"
                # NOTE: slice has no shape and has to assume dim is specified
                # correctly
                elif obj.shape > obj.eq.shape and not isinstance(
                    obj.eq, reno.ops.slice
                ):
                    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}{dims_text}).repeat({obj.shape})\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}{dims_text})\n'
                initial_refs_dict[obj.qual_name()] = obj.qual_name()
            elif (
                isinstance(obj, reno.components.Variable)
                and isinstance(obj.eq, reno.components.Distribution)
                and obj.eq.per_timestep
            ):
                inner_eq = obj.eq.pt_str(
                    __PTNAME__=obj.qual_name(),
                    __DIM__=obj.shape,
                    __DIMNAME__=obj.qual_name() + "_vec",
                    **initial_refs_dict,
                )
                code += f"\t{obj.qual_name()} = {inner_eq}\n"
                per_step_dist_names.append(obj.qual_name())
                # note the index 0 below, remember that the init for this
                # distribution is for the entire sequence, but anytime we'd want to
                # _use_ it in one of the other init equations would be just for the
                # first timestep. (the sequence pass to scan handles the rest)
                initial_refs_dict[obj.qual_name()] = f"{obj.qual_name()}[0]"
            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",
                        __DIM__=obj.shape,
                        __DIMNAME__=obj.qual_name() + "_vec",
                        **initial_refs_dict,
                    )
                    code += f"\t{obj.qual_name()}_init = {inner_eq}\n"
                else:
                    # NOTE: slice has no shape and has to assume dim is specified
                    # correctly
                    if obj.shape > obj.init.shape and not isinstance(
                        obj.init, reno.ops.slice
                    ):
                        inner_eq = f"{obj._implied_eq(obj.init).pt_str(**initial_refs_dict)}.repeat({obj.shape})"
                    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}{dims_text})\n'
                initial_refs_dict[obj.qual_name()] = f"{obj.qual_name()}_init"
            else:
                # flows and vars without a separate init definition
                # NOTE: slice has no shape and has to assume dim is specified
                # correctly
                if obj.shape > obj.eq.shape and not isinstance(obj.eq, reno.ops.slice):
                    inner_eq = f"{obj._implied_eq(obj.eq).pt_str(**initial_refs_dict)}.repeat({obj.shape})"
                else:
                    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}{dims_text})\n'
                initial_refs_dict[obj.qual_name()] = f"{obj.qual_name()}_init"

            # handle initial historical arrays if necessary (but not ones with just -1)
            # TODO: is the len > 1 still necessary?
            if obj in historical_ref_taps and (
                len(historical_ref_taps[obj]) > 1 or None in historical_ref_taps[obj]
            ):
                # static "historical references" get constructed on the fly
                # because it wouldn't make sense to have a separate
                # outputs_info and handling for a thing that doesn't change
                if obj.is_static():
                    continue
                inner_array_value = "0.0"
                if obj.shape > 1:
                    inner_array_value = f"[0.0]*{obj.shape}"
                history_size = 0
                if None in historical_ref_taps[obj]:
                    history_size = steps - 1
                    initial_refs_dict[f"{obj.qual_name()}_h"] = (
                        f"{obj.qual_name()}_hist"
                    )
                else:
                    history_size = min(historical_ref_taps[obj]) * -1

                inner_name = f"{obj.qual_name()}_init"
                if obj.is_static():
                    inner_name = obj.qual_name()

                inner_array = (
                    ", ".join([inner_array_value] * history_size) + f", {inner_name}"
                )
                inner_eq = f"pt.stack([{inner_array}])"
                code += f'\t{obj.qual_name()}_hist = pm.Deterministic("{obj.qual_name()}_hist", {inner_eq})\n'
        except Exception as e:
            e.add_note(
                f"Was trying to set up initial PyMC equation string for {obj.qual_name()}"
            )
            e.add_note(f"\tEquation: {obj.eq}")
            e.add_note(f"\tType: {obj.dtype}")
            e.add_note(f"\tShape: {obj.shape}")
            raise

    # add timeseries
    sequences_code_bits = []
    if model.find_timeref_name() is not None:
        code += f"\ttimestep_seq = pt.as_tensor(np.arange(1, {steps}))\n"
        sequences_code_bits.append("timestep_seq")
    for name in per_step_dist_names:
        sequences_code_bits.append(name)
    sequences_code = f"[{','.join(sequences_code_bits)}]"

    # call the scan function
    return_names = ", ".join(
        [
            f"{obj.qual_name()}_seq"
            for obj in model.all_refs()
            if not obj.is_static() and obj.qual_name() not in per_step_dist_names
        ]
    )
    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() and obj.qual_name() not in per_step_dist_names
            ]
        )
    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:
            if (
                obj in historical_ref_taps
                and (
                    len(historical_ref_taps[obj]) > 1
                    or None in historical_ref_taps[obj]
                )
                and not obj.is_static()
            ):
                taps = historical_ref_taps[obj]
                if None in taps:
                    # handle passing the full timestep range
                    taps = list(range(-steps, 0))
                init_names.append(f"dict(initial={obj.qual_name()}_hist, taps={taps})")
            elif not obj.is_static() and obj.qual_name() not in per_step_dist_names:
                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"
    code += f"\t\tsequences={sequences_code},\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={steps - 1}\n"
    code += "\t)\n"

    initial_refs_dict["__TIMESERIES__"] = "metric"
    # at this point, any timeseries refs will be in metric equations which
    # should _not_ go based on the historical variable

    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() and not obj.qual_name() in per_step_dist_names
    ]:
        dims = '"t"'
        if obj.shape > 1:
            dims = f'("t", "{obj.qual_name()}_vec")'
        code += f'\t{obj.qual_name()} = pm.Deterministic("{obj.qual_name()}", pt.concatenate([pt.as_tensor([{obj.qual_name()}_init]), {obj.qual_name()}_seq]), dims={dims})\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.replace("\t", "    ")


def find_tracked_ref_historical_ref(
    model: "reno.model.Model",
) -> dict[reno.components.TrackedReference, reno.components.HistoricalValue]:
    """Get the set of tracked references that have historical values with
    non-simple index equations and the corresponding historical value components"""

    historical_values = []
    timeseries_values = []
    for ref in model.all_refs():
        historical_values.extend(ref.find_refs_of_type(reno.components.HistoricalValue))
        timeseries_values.extend(
            [
                (ref, ts.sub_equation_parts[0])
                for ts in ref.find_refs_of_type(reno.ops.orient_timeseries)
            ]
        )

    tracked_to_historical = {}

    for val in historical_values:
        if reno.utils.check_for_easy_static_time_eq(val.index_eq):
            index_offset = val.index_eq.eval(force=True)
            # the np.equal(..., np.roll) is a way of determining if all values
            # in array are the same
            if (
                isinstance(index_offset, np.ndarray)
                and len(index_offset) > 1
                and not np.equal(index_offset, np.roll(index_offset, 1))
            ):
                tracked_to_historical[val.tracked_ref] = val
        else:
            tracked_to_historical[val.tracked_ref] = val

    for ref, val in timeseries_values:
        tracked_to_historical[val] = ref

    return tracked_to_historical


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 = []
    timeseries_values = []
    for ref in model.all_refs():
        historical_values.extend(ref.find_refs_of_type(reno.components.HistoricalValue))
        timeseries_values.extend(
            [
                ts.sub_equation_parts[0]
                for ts in ref.find_refs_of_type(reno.ops.orient_timeseries)
            ]
        )

    ref_taps = {}
    for val in historical_values:
        if val.tracked_ref not in ref_taps:
            ref_taps[val.tracked_ref] = []

        if reno.utils.check_for_easy_static_time_eq(val.index_eq):
            # as long as output is singular value or all values are the same,
            # insert the singular tap
            index_offset = val.index_eq.eval(force=True)
            # the np.equal(..., np.roll) is a way of determining if all values
            # in array are the same
            if (
                isinstance(index_offset, np.ndarray)
                and len(index_offset) > 1
                and not np.equal(index_offset, np.roll(index_offset, 1))
            ):
                # otherwise also insert a None
                ref_taps[val.tracked_ref].append(None)
            else:
                # was either a singular value or all values were same
                if isinstance(index_offset, np.ndarray):
                    index_offset = int(index_offset[0])
                ref_taps[val.tracked_ref].append(index_offset)
        else:
            # need to insert a None
            ref_taps[val.tracked_ref].append(None)

        # # 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."
        #     )
        #
        # index_offset = val.index_eq.eval(t=0, force=True)
        # TODO: this doesn't work yet because of how tap names get resolved.
        # With multidim, multiple separate tap references need to be used inside
        # the step function as a single vector, meaning one will have to be
        # constructed somehow. This is an add later feature. (essentially
        # between arbitrary history and multidim, the step function will need to
        # construct a matrix that gets indexed within pymc.)
        # if isinstance(index_offset, np.ndarray):
        #     # in this case (multidim) we essentially "flatten" to get
        #     # all indices that will be required
        #     ref_taps[val.tracked_ref].extend(list(set(list(index_offset))))  # note we only get unique
        # else:
        # if isinstance(index_offset, np.ndarray):
        #     index_offset = int(index_offset[0])
        # ref_taps[val.tracked_ref].append(index_offset)

    for val in timeseries_values:
        if val not in ref_taps:
            ref_taps[val] = []
        ref_taps[val].append(None)

    # 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():
        if None not in taps:
            taps.sort()
        else:
            # you can't sort a list with None in it, so sort the non-None part
            # and then add back in
            ref_taps[ref] = [tap for tap in ref_taps[ref] if tap is not None]
            ref_taps[ref].sort()
            ref_taps[ref].append(None)

    return ref_taps


def expected_pt_scan_arg_names(
    model: "reno.model.Model",
    as_dict: bool = False,
    tap_1_only: bool = False,
    max_taps=None,
) -> 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

    names_to_args_indices = {}
    index = 0

    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)
        names_to_args_indices[timeref_name] = index
        index += 1
    # find any per-timestep distributions
    for var in model.all_vars():
        if isinstance(var.eq, reno.components.Distribution) and var.eq.per_timestep:
            names.append(var.qual_name())
            names_to_args_indices[var.qual_name()] = index
            index += 1

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

    # get last of each stock
    for stock in model.all_stocks():
        if stock in historical_ref_taps:
            # handle setting up individual tap names
            if None not in historical_ref_taps[stock]:
                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_to_args_indices[name] = index
                        index += 1
            else:
                # handle passing full sequence
                names.append(f"{stock.qual_name()}_h")
                names_to_args_indices[f"{stock.qual_name()}_h"] = slice(
                    index, index + max_taps
                )
                index += max_taps - 1  # -1 to handle v0 == v0[t-1]
        names.append(stock.qual_name())
        names_to_args_indices[stock.qual_name()] = index
        index += 1

    for flow in model.all_flows():
        if not flow.is_static():
            if flow in historical_ref_taps:
                # handle setting up individual tap names
                if None not in historical_ref_taps[flow]:
                    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_to_args_indices[name] = index
                            index += 1
                else:
                    # handle passing full sequence
                    names.append(f"{flow.qual_name()}_h")
                    names_to_args_indices[f"{flow.qual_name()}_h"] = slice(
                        index, index + max_taps
                    )
                    index += max_taps - 1  # -1 to handle v0 == v0[t-1]
            names.append(flow.qual_name())
            names_to_args_indices[flow.qual_name()] = index
            index += 1

    for var in model.all_vars():
        # skip vars that are just per-timestep distributions (already handled)
        if isinstance(var.eq, reno.components.Distribution) and var.eq.per_timestep:
            continue
        if not var.is_static():
            if var in historical_ref_taps:
                # handle setting up individual tap names
                if None not in historical_ref_taps[var]:
                    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_to_args_indices[name] = index
                            index += 1
                else:
                    # handle passing full sequence
                    names.append(f"{var.qual_name()}_h")
                    names_to_args_indices[f"{var.qual_name()}_h"] = slice(
                        index, index + max_taps
                    )
                    index += max_taps - 1  # -1 to handle v0 == v0[t-1]
            names.append(var.qual_name())
            names_to_args_indices[var.qual_name()] = index
            index += 1

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

    for flow in model.all_flows():
        if flow.is_static():
            names.append(flow.qual_name())
            names_to_args_indices[flow.qual_name()] = index
            index += 1

    for var in model.all_vars():
        if var.is_static():
            names.append(var.qual_name())
            names_to_args_indices[var.qual_name()] = index
            index += 1

    # 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
    return names_to_args_indices


def resolve_pt_scan_args(
    model: "reno.model.Model", args: list[pt.TensorVariable], max_taps: int = None
) -> 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_to_indices = expected_pt_scan_arg_names(model, max_taps=max_taps)
    for name, index in arg_names_to_indices.items():
        if isinstance(index, slice):
            refs[name] = pt.stack(args[index], axis=0)
        else:
            refs[name] = args[index]

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

    return refs
