"""Load pathtracer-diff natively on this machine (Windows/SAURON included).

The published Hugging Face kernel carries `*-linux` variants only, so on
Windows `get_kernel` has no matching build. The source is right here, though:
JIT-build it with `torch.utils.cpp_extension` and expose the identical
`Scene` / `Camera` / `render` API.

    import load_local
    ptd = load_local.load()
"""
import os
import sys
import types

_cached = None


def load(verbose=False):
    """Return the pathtracer_diff module, JIT-built from the local source."""
    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__))
    cuda_dir = os.path.join(root, "pathtracer_diff_cuda")
    ext = _jit(name="pathtracer_diff_ext",
               sources=[os.path.join(root, "local_bind.cpp"),
                        os.path.join(cuda_dir, "pathtracer.cu"),
                        os.path.join(cuda_dir, "geometry_grad.cu")],
               extra_include_paths=[cuda_dir],
               verbose=verbose)
    ops_mod = types.ModuleType("pathtracer_diff._ops")
    ops_mod.ops = type("_Ops", (), {
        "pt_forward": staticmethod(ext.pt_forward),
        "pt_backward": staticmethod(ext.pt_backward),
        "pt_geometry_grad": staticmethod(ext.pt_geometry_grad),
    })()
    sys.modules["pathtracer_diff._ops"] = ops_mod
    sys.path.insert(0, os.path.join(root, "torch-ext"))
    import pathtracer_diff
    _cached = pathtracer_diff
    return pathtracer_diff