Source code for pykinbiont._core

from __future__ import annotations

import json
import warnings
from pathlib import Path
from typing import Optional

_CONFIG_FILE = Path.home() / ".config" / "pykinbiont" / "config.json"
_KINBIONT_UUID = "3b7c519f-be5e-4159-a208-6486ce96bb37"

_jl: Optional[object] = None


def _apply_dev_path(resolved: str) -> None:
    """Register a local KinBiont.jl clone as a dev package with juliapkg.

    juliapkg manages its own Julia environment and ignores JULIA_PROJECT, so
    the only reliable way to redirect it to a local source tree is via its
    own API.  This writes the dev spec to juliapkg's per-environment
    juliapkg.json; the actual Julia re-resolution happens the next time
    juliacall starts Julia (i.e. after a kernel / process restart).
    """
    try:
        import juliapkg  # type: ignore[import]

        juliapkg.add(
            juliapkg.PkgSpec(
                name="Kinbiont",
                uuid=_KINBIONT_UUID,
                dev=True,
                path=resolved,
            )
        )
    except Exception as exc:
        warnings.warn(
            f"Could not register dev Kinbiont path with juliapkg: {exc}\n"
            "Julia will use the registry-installed version.",
            stacklevel=3,
        )




def configure(project_path: str) -> None:
    """Point pykinbiont at a local KinBiont.jl development clone.

    Registers the local directory as a *dev* package with juliapkg so that
    Julia will load your source tree instead of the registry-installed version.
    The path is persisted to ``~/.config/pykinbiont/config.json`` and applied
    automatically on every subsequent Python / kernel start, so you only need
    to call this once per machine.

    .. important::
        If Julia is already running in the current kernel this call has no
        immediate effect — **restart the kernel** so juliacall can pick up the
        new environment before calling ``fit()`` or ``preprocess()``.

    Parameters
    ----------
    project_path:
        Path to a local KinBiont.jl source directory.
    """
    resolved = str(Path(project_path).resolve())

    _CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(_CONFIG_FILE, "w") as f:
        json.dump({"project_path": resolved}, f, indent=2)

    _apply_dev_path(resolved)

    if _jl is not None:
        print(
            f"KinBiont dev path saved: {resolved}\n"
            "Julia is already running — restart the kernel for this to take effect."
        )
    else:
        print(f"KinBiont dev path configured: {resolved}")



def _load_saved_config() -> Optional[str]:
    if _CONFIG_FILE.exists():
        with open(_CONFIG_FILE) as f:
            return json.load(f).get("project_path")
    return None




def init(project_path: Optional[str] = None) -> object:
    """Start Julia and load Kinbiont.

    Calling this explicitly is optional — all conversion and fitting functions
    trigger it automatically on first use.

    Parameters
    ----------
    project_path:
        Optional path to a local KinBiont.jl directory.
        Equivalent to calling ``configure(project_path)`` before ``init()``.
    """
    global _jl

    if _jl is not None:
        return _jl

    if project_path is not None:
        configure(project_path)

    from juliacall import Main as jl  # noqa: PLC0415  # Julia starts here

    # If a local dev path is configured, activate it so that `using Kinbiont`
    # loads from the source tree rather than the registry-installed version.
    # Pkg.activate(path) + using Kinbiont is the standard Julia dev workflow.
    #
    # We immediately restore the original (juliapkg) project afterwards so that
    # PythonCall and its dependents remain in the active environment.  This
    # allows Julia to precompile extension packages (SciMLBasePythonCallExt,
    # etc.) correctly and suppresses the otherwise-spurious "not installed"
    # warnings about PythonCall.
    dev_path = _load_saved_config()
    if dev_path:
        try:
            saved = jl.seval("Base.active_project()")
            jl.seval(f'import Pkg; Pkg.activate(raw"{dev_path}")')
            jl.seval("using Kinbiont")
            # Restore the juliapkg environment so PythonCall is in scope for
            # extension precompilation, then retry any extensions that failed.
            jl.seval(f'Pkg.activate(raw"{str(saved)}")')
            jl.seval("Base.retry_load_extensions()")
        except Exception as exc:
            warnings.warn(
                f"Dev-path Kinbiont activation failed: {exc}\n"
                "Falling back to the registry-installed Kinbiont.",
                stacklevel=2,
            )
    else:
        jl.seval("using Kinbiont")

    jl.seval("function _pykinbiont_to_matrix(arr); Matrix{Float64}(arr); end")

    # Helpers to wrap Python callables as proper Julia Function objects.
    # NLModel and ODEModel structs require func::Function; plain Py callables
    # cannot be stored there.  These wrappers create real Julia closures.
    jl.seval("""
    function _pykinbiont_make_nl_func(pyfunc)::Function
        return (p, t) -> PythonCall.pyconvert(Vector{Float64}, pyfunc(p, t))
    end
    function _pykinbiont_make_ode_func(pyfunc)::Function
        return (du, u, p, t) -> (pyfunc(du, u, p, t); nothing)
    end
    """)

    _jl = jl
    return _jl





def get_jl() -> object:
    """Return the Julia Main module, auto-initialising if needed."""
    if _jl is None:
        return init()
    return _jl



# ---------------------------------------------------------------------------
# At import time: if a dev path was previously saved, tell juliapkg about it
# now — before juliacall has a chance to start Julia and lock the environment.
# ---------------------------------------------------------------------------
_saved = _load_saved_config()
if _saved:
    _apply_dev_path(_saved)