# -*- coding: utf-8 -*-
"""Low-level Julia execution utilities for running generated Julia code.

This module provides functions for executing Julia code generated by TVBO
templates (DifferentialEquations.jl, BifurcationKit.jl, NetworkDynamics.jl)
and extracting results back into Python/NumPy.

Uses juliacall (``tvbo.adapters.julia``) for the Python↔Julia bridge,
keeping full Julia objects available for interactive inspection.
"""

from __future__ import annotations

import numpy as np

from tvbo.adapters.julia import eval_with_auto_install, get_julia


def ensure_packages(*packages: str) -> None:
    """Ensure Julia packages are installed, installing if missing.

    Parameters
    ----------
    packages : str
        Package names to ensure are available.
    """
    for pkg in packages:
        # Just try to load — eval_with_auto_install handles missing packages
        eval_with_auto_install(f"using {pkg}")


def run_julia_code(code: str, compiled_modules: bool = True):
    """Execute a block of Julia code in the Main module.

    Parameters
    ----------
    code : str
        Julia source code to evaluate.
    compiled_modules : bool
        Passed to ``get_julia()``. Default True to use precompiled packages.

    Returns
    -------
    object
        The return value of the last Julia expression.
    """
    _jl, _Main = get_julia(compiled_modules=compiled_modules)
    return eval_with_auto_install(code)


def extract_ode_solution() -> tuple[np.ndarray, np.ndarray, object]:
    """Extract time and state arrays from a Julia ODE/SDE solution.

    Assumes a variable ``sol`` exists in Julia Main after running
    a DifferentialEquations.jl / NetworkDynamics.jl solve.

    Returns
    -------
    t : np.ndarray
        Time points, shape ``(n_t,)``.
    u : np.ndarray
        State array, shape ``(n_states, n_t)``.
    sol : object
        The raw Julia solution proxy for interactive use.
    """
    sol = eval_with_auto_install("sol")
    t = np.array(eval_with_auto_install("Array(sol.t)"), dtype=float)
    u = np.array(eval_with_auto_install("Array(sol)"), dtype=float)
    return t, u, sol


def extract_bifurcation_result():
    """Extract a bifurcation result from Julia Main.

    Assumes ``bifurcation_result`` exists after running a
    BifurcationKit.jl script.

    Returns
    -------
    object
        The Julia BifurcationResult proxy.
    """
    return eval_with_auto_install("bifurcation_result")


def solution_to_dataarray(
    t: np.ndarray,
    u: np.ndarray,
    sv_names: list[str],
    n_nodes: int,
):
    """Build an ``xr.DataArray`` from a flat Julia ODE solution.

    Reshapes ``u`` from flat ``(n_states, n_t)`` into dims
    ``(time, variable, node)`` with named coordinates.

    Parameters
    ----------
    t : np.ndarray
        Time array, shape ``(n_t,)``.
    u : np.ndarray
        Raw solution, shape ``(n_nodes * n_sv, n_t)``.
    sv_names : list[str]
        State variable names (length determines n_sv).
    n_nodes : int
        Number of nodes.

    Returns
    -------
    xr.DataArray
        Dims ``(time, variable, node)`` with coords.
    """
    import xarray as xr

    n_sv = len(sv_names)
    data = u.T.reshape(len(t), n_nodes, n_sv).transpose(0, 2, 1)
    return xr.DataArray(
        data=data,
        dims=['time', 'variable', 'node'],
        coords={
            'time': np.asarray(t),
            'variable': list(sv_names),
            'node': [str(i) for i in range(n_nodes)],
        },
    )
