"""Load physarum natively on hosts without a published variant (e.g. a Pi).

JIT-builds the local source with `torch.utils.cpp_extension` and exposes the
identical `physarum` module API.

    import load_local
    physarum = load_local.load()
    sim = physarum.Physarum(width=1024, height=1024, agents=200000)
"""
import os
import sys
import types

_cached = None


def load(verbose=False):
    global _cached
    if _cached is not None:
        return _cached
    from torch.utils.cpp_extension import load as _jit

    root = os.path.dirname(os.path.abspath(__file__))
    ext = _jit(name="physarum_ext",
               sources=[os.path.join(root, "local_bind.cpp"),
                        os.path.join(root, "physarum_csrc", "physarum_cpu.cpp"),
                        os.path.join(root, "physarum_csrc", "flow_cpu.cpp")],
               extra_cflags=["-O3"],
               verbose=verbose)
    ops_mod = types.ModuleType("physarum._ops")
    ops_mod.ops = type("_Ops", (), {"physarum_step": staticmethod(ext.physarum_step),
                                    "flow_cg": staticmethod(ext.flow_cg)})()
    sys.modules["physarum._ops"] = ops_mod
    sys.path.insert(0, os.path.join(root, "torch-ext"))
    import physarum
    _cached = physarum
    return physarum