OptunaWrapper.py

Go to the documentation of this file.

import sqlite3
import json
import inspect
import textwrap
import linecache
import ast
import base64
import pickle
import os
import sysconfig
from contextlib import contextmanager
from typing import Any, Optional, Iterable
from urllib.parse import unquote
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
 
from .LHSIterator import LHSIterator
 
 
# ============================================================
# Helpers
# ============================================================
 
def variable_to_pickle64b(value: Any) -> str:
    data = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
    return base64.b64encode(data).decode("ascii")
 
 
def pickle64b_to_variable(payload: str) -> Any:
    data = base64.b64decode(payload.encode("ascii"))
    return pickle.loads(data)
 
 
 
def _capture_external_caller_frame():
    """
    Return the first 'user' frame (not this lib, not stdlib).
    Works better with contextmanager wrappers (contextlib).
    """
    lib_file = os.path.abspath(__file__) if "__file__" in globals() else None
 
    stdlib_dir = os.path.abspath(sysconfig.get_paths()["stdlib"])
    purelib_dir = os.path.abspath(sysconfig.get_paths().get("purelib", ""))
    platlib_dir = os.path.abspath(sysconfig.get_paths().get("platlib", ""))
 
    def is_internal_frame(fr) -> bool:
        filename = fr.f_code.co_filename
 
        # Cas notebook custom: "Cell 12" -> ce n'est pas un vrai path, on le garde
        if not os.path.isabs(filename):
            # ex: "Cell 3", "<ipython-input-...>", "<stdin>"
            return False
 
        filename = os.path.abspath(filename)
 
        # 1) ce fichier de lib
        if lib_file is not None and filename == lib_file:
            return True
 
        # 2) stdlib (contextlib, inspect, etc.)
        if filename.startswith(stdlib_dir + os.sep):
            return True
 
        # 3) site-packages (optionnel, mais souvent utile)
        if purelib_dir and filename.startswith(purelib_dir + os.sep):
            return True
        if platlib_dir and filename.startswith(platlib_dir + os.sep):
            return True
 
        return False
 
    fr = inspect.currentframe()
    try:
        fr = fr.f_back
        while fr is not None:
            # debug éventuel:
            # print("STACK:", fr.f_code.co_filename, fr.f_code.co_name)
 
            if not is_internal_frame(fr):
                return fr
 
            fr = fr.f_back
 
        return None
    finally:
        del fr
 
 
def _capture_external_caller_namespace():
    fr = _capture_external_caller_frame()
    if fr is None:
        return {}, {}
    return fr.f_globals, fr.f_locals
 
 
def _get_function_source(fn) -> str:
    """
    Robust-ish getsource for normal .py + custom notebooks if linecache is populated.
    """
    try:
        return textwrap.dedent(inspect.getsource(fn))
    except Exception:
        code = getattr(fn, "__code__", None)
        if code is None:
            raise
 
        filename = code.co_filename
        lines = linecache.getlines(filename)
        if not lines:
            raise
 
        src = "".join(lines)
        tree = ast.parse(src)
 
        # Try exact function match by name + line number
        target_name = fn.__name__
        target_lineno = getattr(code, "co_firstlineno", None)
 
        for node in tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == target_name:
                if target_lineno is None or node.lineno == target_lineno:
                    if hasattr(node, "end_lineno") and node.end_lineno is not None:
                        return textwrap.dedent("".join(lines[node.lineno - 1: node.end_lineno]))
                    return textwrap.dedent("".join(lines[node.lineno - 1:]))
 
        raise
 
 
# ============================================================
# Optuna Archive
# ============================================================
 
class OptunaArchive:
    """
    Lightweight helper to standardize what is stored in Optuna study/trial user_attrs.
 
    Main idea:
    - Typed helpers for trial outputs (timeseries)
    - Explicit function registration (suggest/run/eval)
    - Optional context capture via `with archiver.register(): ...`
      to archive helper functions + runtime objects automatically
      into study_user_attributes['backupContext'].
    """
 
    KEY_ENTRY_SUGGEST = "entrypoint:suggest"
    KEY_ENTRY_RUN = "entrypoint:run"
    KEY_ENTRY_EVAL = "entrypoint:eval"
    KEY_CONTEXT = "backupContext"
 
    # contexte partagé unique par DB (pas de study_map)
    CONTEXT_ID = "__default__"
 
    def __init__(self, study=None):
        self.study = None
        self.backupContext: dict[str, dict[str, Any]] = {}
 
        # cache des entrypoints tant qu'on n'a pas de study
        self._entrypoints_cache: dict[str, str] = {}
 
        # connexion sqlite custom (ouverte après set_study)
        self._db_path: Optional[str] = None
        self._con_custom: Optional[sqlite3.Connection] = None
 
        if study is not None:
            self.set_study(study)
 
    # ---------------------------
    # Study attach / DB init / flush
    # ---------------------------
    def set_study(self, study):
        self.study = study
        self._db_path = self._extract_sqlite_path_from_study(study)
 
        # ouvre connexion vers la même DB SQLite qu'Optuna
        if self._con_custom is not None:
            try:
                self._con_custom.close()
            except Exception:
                pass
        self._con_custom = sqlite3.connect(self._db_path)
 
        self._ensure_isicell_tables()
        self._flush_context_if_possible()
        self._flush_entrypoints_if_possible()
        return self
 
    def _require_custom_db(self):
        if self._con_custom is None:
            raise RuntimeError(
                "No custom DB connection yet. Call set_study(study) first."
            )
 
 
    def _require_study(self):
        if self.study is None:
            raise RuntimeError("OptunaArchive.study is None. Attach a study first (OptunaArchive(study) or set_study()).")
 
 
    def _ensure_isicell_tables(self):
        self._require_custom_db()
        cur = self._con_custom.cursor()
 
        cur.execute("""
            CREATE TABLE IF NOT EXISTS isicell_context (
                context_id TEXT NOT NULL,
                key       TEXT NOT NULL,
                mode      TEXT NOT NULL,
                payload   TEXT NOT NULL,
                PRIMARY KEY (context_id, key)
            );
        """)
 
        cur.execute("""
            CREATE TABLE IF NOT EXISTS isicell_context_entrypoints (
                context_id TEXT NOT NULL,
                kind       TEXT NOT NULL,   -- suggest | run | eval
                name       TEXT NOT NULL,
                PRIMARY KEY (context_id, kind)
            );
        """)
 
        self._con_custom.commit()
 
    def _flush_context_if_possible(self):
        """
        Flush in-memory backupContext to custom table.
        Upsert behavior (safe if called multiple times).
        """
        if self._con_custom is None or not self.backupContext:
            return
 
        self._ensure_isicell_tables()
        cur = self._con_custom.cursor()
 
        for name, entry in self.backupContext.items():
            mode = entry.get("mode")
            payload = entry.get("payload")
 
            # payload must be text in SQL table
            if mode == "json":
                payload_txt = json.dumps(payload)
            else:
                # mode source / pickle64b already string-like
                payload_txt = str(payload)
 
            cur.execute("""
                INSERT INTO isicell_context (context_id, key, mode, payload)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(context_id, key) DO UPDATE SET
                    mode=excluded.mode,
                    payload=excluded.payload;
            """, (self.CONTEXT_ID, name, mode, payload_txt))
 
        self._con_custom.commit()
 
    def _flush_entrypoints_if_possible(self):
        """
        Flush cached entrypoints to custom table.
        """
        if self._con_custom is None or not self._entrypoints_cache:
            return
 
        self._ensure_isicell_tables()
        cur = self._con_custom.cursor()
 
        for kind, name in self._entrypoints_cache.items():
            cur.execute("""
                INSERT INTO isicell_context_entrypoints (context_id, kind, name)
                VALUES (?, ?, ?)
                ON CONFLICT(context_id, kind) DO UPDATE SET
                    name=excluded.name;
            """, (self.CONTEXT_ID, kind, name))
 
        self._con_custom.commit()
 
    # ---------------------------
    # Serialization helpers
    # ---------------------------
    @staticmethod
    def _to_jsonable(x: Any) -> Any:
        if isinstance(x, np.ndarray):
            return x.tolist()
        if isinstance(x, (np.integer, np.floating)):
            return x.item()
        if isinstance(x, pd.Series):
            return {
                "__kind__": "series",
                "name": x.name,
                "value": x.to_dict()
            }
        if isinstance(x, pd.DataFrame):
            return {
                "__kind__": "dataframe",
                "orient": "split",
                "value": x.to_dict("split")
            }
        if isinstance(x, dict):
            return {k: OptunaArchive._to_jsonable(v) for k, v in x.items()}
        if isinstance(x, (list, tuple)):
            return [OptunaArchive._to_jsonable(v) for v in x]
        # Let plain python scalars/strings/bool/None pass
        return x
 
    def _encode_context_value(self, value: Any) -> dict[str, Any]:
        """
        Context encoding policy:
        - functions/classes/callables -> source code (mode='source') if possible
        - JSON-serializable (after _to_jsonable) -> mode='json'
        - fallback -> pickle64b
        """
        if callable(value):
            src = _get_function_source(value)
            return {"mode": "source", "payload": src}
 
        try:
            j = self._to_jsonable(value)
            json.dumps(j)  # validate serializable
            return {"mode": "json", "payload": j}
        except Exception:
            return {"mode": "pickle64b", "payload": variable_to_pickle64b(value)}
 
    # ---------------------------
    # Context manager capture
    # ---------------------------
    @contextmanager
    def register(self):
        """
        Capture newly created symbols in caller scope and store them in backupContext.
 
        Example
        -------
        archiver = OptunaArchive(study)
        with archiver.register():
            data = {...}
            def dataCatcherOnStep(...): ...
            def run_simu(...): ...
            def optim(trial): ...
        # all new symbols now archived in study_user_attributes['backupContext']
        """
        g0, l0 = _capture_external_caller_namespace()
        before_g = set(g0.keys())
        before_l = set(l0.keys())
        try:
            yield self
        finally:
            g1, l1 = _capture_external_caller_namespace()
 
            new_names = [k for k in l1.keys() if k not in before_l]
            new_names += [k for k in g1.keys() if k not in before_g and k not in new_names]
 
            for name in new_names:
                if name.startswith("__"):
                    continue
                value = l1[name] if name in l1 else g1[name]
                try:
                    self.backupContext[name] = self._encode_context_value(value)
                except Exception as e:
                    raise RuntimeError(f"Failed to archive symbol '{name}'.\nerror: {e}") from e
 
            # si study déjà attachée, flush immédiat; sinon ça reste en cache
            self._flush_context_if_possible()
 
    # ---------------------------
    # Manual context registration (optional)
    # ---------------------------
    def context_add(self, name: str, value: Any):
        self.backupContext[name] = self._encode_context_value(value)
        self._flush_context_if_possible()
 
    def context_add_many(self, mapping: dict[str, Any]):
        for k, v in mapping.items():
            self.context_add(k, v)
 
    def _extract_sqlite_path_from_study(self, study) -> str:
        """
        Extract sqlite file path from an Optuna study backed by RDBStorage(sqlite),
        even if wrapped by _CachedStorage.
        """
        storage = getattr(study, "_storage", None)
        if storage is None:
            raise RuntimeError("Study has no _storage attribute.")
 
        # Unwrap common Optuna wrappers (_CachedStorage, etc.)
        visited = set()
        while storage is not None and id(storage) not in visited:
            visited.add(id(storage))
 
            engine = getattr(storage, "engine", None)
            if engine is not None:
                url = str(engine.url)
                if not url.startswith("sqlite:///"):
                    raise RuntimeError(f"Only sqlite storage is supported. Got storage URL: {url}")
 
                path = url[len("sqlite:///"):]  # keep leading slash if absolute path
                path = unquote(path)
 
                if not os.path.isabs(path):
                    path = os.path.abspath(path)
                return path
 
            # Try common wrapper attributes used by Optuna versions
            if hasattr(storage, "_backend"):
                storage = storage._backend
                continue
            if hasattr(storage, "_storage"):
                storage = storage._storage
                continue
 
            break
 
        raise RuntimeError(
            "Could not access underlying RDBStorage engine from study._storage "
            "(possibly an unsupported Optuna storage wrapper/version)."
        )
 
    def _set_entrypoint(self, key: str, name: str):
        """
        API inchangée:
          - vérifie que le symbole existe dans backupContext et est source-backed
          - stocke en cache + flush si possible
        """
        if name not in self.backupContext:
            raise KeyError(
                f"'{name}' not found in backupContext. "
                "Define it inside `with archiver.register(): ...` first."
            )
 
        entry = self.backupContext[name]
        if entry.get("mode") != "source":
            raise TypeError(
                f"Entrypoint '{name}' must be a source-backed symbol, "
                f"got mode='{entry.get('mode')}'."
            )
 
        # map old keys -> compact kinds in custom table
        if key == self.KEY_ENTRY_SUGGEST:
            kind = "suggest"
        elif key == self.KEY_ENTRY_RUN:
            kind = "run"
        elif key == self.KEY_ENTRY_EVAL:
            kind = "eval"
        else:
            raise KeyError(f"Unknown entrypoint key: {key}")
 
        self._entrypoints_cache[kind] = name
        self._flush_entrypoints_if_possible()
 
    def set_suggest_entrypoint(self, name: str):
        self._set_entrypoint(self.KEY_ENTRY_SUGGEST, name)
 
    def set_run_entrypoint(self, name: str):
        self._set_entrypoint(self.KEY_ENTRY_RUN, name)
 
    def set_eval_entrypoint(self, name: str):
        self._set_entrypoint(self.KEY_ENTRY_EVAL, name)
 
    # ---------------------------
    # Generic study/trial data
    # ---------------------------
    def study_add_data(self, key: str, value: Any):
        self._require_study()
        self.study.set_user_attr(key, self._to_jsonable(value))
 
    @staticmethod
    def trial_add_data(trial, key: str, value: Any):
        trial.set_user_attr(key, OptunaArchive._to_jsonable(value))
 
    # ---------------------------
    # Axes helpers (for timeseries / vectors)
    # ---------------------------
    def study_add_axis(self, name: str, values: Any):
        self._require_study()
        self.study.set_user_attr(f"axis:{name}", self._to_jsonable(values))
 
    # ---------------------------
    # Typed trial helpers
    # ---------------------------
    def trial_add_timeseries(
        self,
        trial,
        key: str,
        values: Any,
        axes: list[str],
        columns: Optional[list[str]] = None,
        value_col: str = "value",
        layout: str = "aligned",
    ):
        """
        Store trial-level flattened values and a study-level spec.
 
        layout='product' : cartesian product of axes
        layout='aligned' : zipped/aligned axes (same length as values)
        """
        self._require_study()
 
        if not isinstance(axes, (list, tuple)) or len(axes) < 1:
            raise ValueError("axes must be a non-empty list of axis names.")
        if columns is None:
            columns = list(axes)
        if len(columns) != len(axes):
            raise ValueError("columns and axes must have the same length.")
        if layout not in ("product", "aligned"):
            raise ValueError("layout must be 'product' or 'aligned'.")
 
        for ax in axes:
            if f"axis:{ax}" not in self.study.user_attrs:
                raise KeyError(f"Missing study axis 'axis:{ax}'. Call study_add_axis('{ax}', ...) first.")
 
        spec_key = f"timeseries_spec:{key}"
        spec = {
            "axes": list(axes),
            "columns": list(columns),
            "value_col": value_col,
            "layout": layout,
        }
 
        existing = self.study.user_attrs.get(spec_key, None)
        if existing is None:
            self.study.set_user_attr(spec_key, spec)
        elif existing != spec:
            raise ValueError(f"Incompatible timeseries spec for key '{key}'. Existing={existing}, New={spec}")
 
        # Always store trial payload
        trial.set_user_attr(key, self._to_jsonable(values))
 
    # ---------------------------
    # Validation
    # ---------------------------
    def check_minimal(self, require_run_function: bool = True):
        missing = []
 
        if "suggest" not in self._entrypoints_cache:
            missing.append("suggest entrypoint")
        if require_run_function and "run" not in self._entrypoints_cache:
            missing.append("run entrypoint")
 
        if missing:
            raise RuntimeError(
                f"Missing required archive entrypoints: {missing}. "
                "Use set_suggest_entrypoint/set_run_entrypoint before optimize()."
            )
 
        # Si study déjà attachée, on vérifie aussi que tout est bien flush en DB custom
        if self._con_custom is not None:
            cur = self._con_custom.cursor()
            rows = cur.execute("""
                SELECT kind, name
                FROM isicell_context_entrypoints
                WHERE context_id = ?;
            """, (self.CONTEXT_ID,)).fetchall()
            kinds_in_db = {k for k, _ in rows}
 
            if "suggest" not in kinds_in_db:
                raise RuntimeError("Missing 'suggest' entrypoint in isicell_context_entrypoints.")
            if require_run_function and "run" not in kinds_in_db:
                raise RuntimeError("Missing 'run' entrypoint in isicell_context_entrypoints.")
 
 
# ============================================================
# Optuna DB reader / replay
# ============================================================
 
class OptunaDB:
    """
    Read-only Optuna SQLite helper (fast SQL) + archived function loading (exec) + replay-ready helpers.
 
    Context and entrypoints are read from custom tables:
      - isicell_context
      - isicell_context_entrypoints
 
    Study-specific metadata (axes, timeseries specs, etc.) remains in study_user_attributes.
    """
 
    CONTEXT_ID = "__default__"
 
    def __init__(self, db_path: str, readonly: bool = True):
        self.db_path = db_path
        if readonly:
            self.con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        else:
            self.con = sqlite3.connect(db_path)
 
        self.suggest_fn = None
        self.run_fn = None
        self.eval_fn = None
 
        # charge le contexte global si tables présentes
        if self._has_isicell_tables():
            self.load_context()
        elif self.studies():
            # fallback legacy possible (si ancienne DB)
            # self.load_context_legacy(self.studies()[0])
            pass
 
    def close(self):
        self.con.close()
 
    # ============================================================
    # Basic metadata
    # ============================================================
    def studies(self) -> list[str]:
        rows = self.con.execute("SELECT study_name FROM studies ORDER BY study_name;").fetchall()
        return [r[0] for r in rows]
 
    def trial_to_study(self, trial_id: int) -> str:
        row = self.con.execute(
            """
            SELECT studies.study_name
            FROM trials
            INNER JOIN studies ON studies.study_id = trials.study_id
            WHERE trials.trial_id = ?;
            """,
            (int(trial_id),),
        ).fetchone()
        if row is None:
            raise KeyError(f"trial_id {trial_id} not found")
        return row[0]
 
    def study_attrs(self, study_name: str) -> dict[str, Any]:
        rows = self.con.execute(
            """
            SELECT key, value_json
            FROM study_user_attributes
            INNER JOIN studies ON studies.study_id = study_user_attributes.study_id
            WHERE studies.study_name = ?;
            """,
            (study_name,),
        ).fetchall()
        return {k: json.loads(v) for k, v in rows}
 
    def trial_attrs(self, study_name: str, key: Optional[str] = None) -> pd.DataFrame:
        if key is None:
            df = pd.read_sql_query(
                """
                SELECT studies.study_name, trial_user_attributes.trial_id, trial_user_attributes.key, trial_user_attributes.value_json
                FROM trial_user_attributes
                INNER JOIN trials ON trial_user_attributes.trial_id = trials.trial_id
                INNER JOIN studies ON studies.study_id = trials.study_id
                WHERE trials.state="COMPLETE" AND studies.study_name = ?;
                """,
                self.con,
                params=(study_name,),
            )
        else:
            df = pd.read_sql_query(
                """
                SELECT studies.study_name, trial_user_attributes.trial_id, trial_user_attributes.key, trial_user_attributes.value_json
                FROM trial_user_attributes
                INNER JOIN trials ON trial_user_attributes.trial_id = trials.trial_id
                INNER JOIN studies ON studies.study_id = trials.study_id
                WHERE trials.state="COMPLETE" AND studies.study_name = ? AND trial_user_attributes.key = ?;
                """,
                self.con,
                params=(study_name, key),
            )
 
        if df.empty:
            return pd.DataFrame(columns=["study_name", "trial_id", "key", "value"])
 
        df = df.rename(columns={"value_json": "value"})
        df["value"] = df["value"].map(json.loads)
        return df
 
    # ============================================================
    # Fitness / params (fast SQL)
    # ============================================================
    def fitness_long(self, study_name: Optional[str] = None) -> pd.DataFrame:
        if study_name is None:
            return pd.read_sql_query(
                """
                SELECT studies.study_name, trial_values.trial_id, trial_values.objective, trial_values.value
                FROM trial_values
                INNER JOIN trials  ON trial_values.trial_id = trials.trial_id
                INNER JOIN studies ON studies.study_id = trials.study_id
                WHERE trials.state="COMPLETE";
                """,
                self.con,
            )
 
        return pd.read_sql_query(
            """
            SELECT studies.study_name, trial_values.trial_id, trial_values.objective, trial_values.value
            FROM trial_values
            INNER JOIN trials  ON trial_values.trial_id = trials.trial_id
            INNER JOIN studies ON studies.study_id = trials.study_id
            WHERE trials.state="COMPLETE" AND studies.study_name = ?;
            """,
            self.con,
            params=(study_name,),
        )
 
    def fitness_wide(self, study_name: Optional[str] = None) -> pd.DataFrame:
        res = self.fitness_long(study_name=study_name)
        if res.empty:
            return res
 
        wide = res.pivot(index=["study_name", "trial_id"], columns="objective", values="value").reset_index()
        wide["trial_num"] = wide.groupby("study_name")["trial_id"].transform(lambda s: s - s.min())
 
        if 0 in wide.columns:
            wide = wide.sort_values(["study_name", "trial_num"]).reset_index(drop=True)
            wide["best"] = wide.groupby("study_name")[0].cummin()
 
        return wide
 
    def params(self, study_name: Optional[str] = None) -> pd.DataFrame:
        if study_name is None:
            res = pd.read_sql_query(
                """
                SELECT studies.study_name, trial_params.trial_id, trial_params.param_name, trial_params.param_value
                FROM trial_params
                INNER JOIN trials ON trial_params.trial_id = trials.trial_id
                INNER JOIN studies ON studies.study_id = trials.study_id
                WHERE trials.state="COMPLETE";
                """,
                self.con,
            )
        else:
            res = pd.read_sql_query(
                """
                SELECT studies.study_name, trial_params.trial_id, trial_params.param_name, trial_params.param_value
                FROM trial_params
                INNER JOIN trials ON trial_params.trial_id = trials.trial_id
                INNER JOIN studies ON studies.study_id = trials.study_id
                WHERE trials.state="COMPLETE" AND studies.study_name = ?;
                """,
                self.con,
                params=(study_name,),
            )
 
        if res.empty:
            return pd.DataFrame()
 
        return res.pivot(index=["study_name", "trial_id"], columns="param_name", values="param_value")
 
    # ============================================================
    # Timeseries helper
    # ============================================================
    def _normalize_ids(self, x, cast=int, name="ids"):
        """
        Normalize None | scalar | iterable -> None | list[cast(x)].
        Strings are treated as scalars (not iterables).
        """
        if x is None:
            return None
 
        # scalar str/int/etc.
        if isinstance(x, (str, bytes)):
            return [cast(x)]
        if not isinstance(x, Iterable):
            return [cast(x)]
 
        # iterable
        vals = [cast(v) for v in x]
        return vals
 
 
    def _sql_in_clause(self, values):
        """
        Returns ('(?,?,?)', [..]) for SQL IN clauses.
        values must be a non-empty list.
        """
        if values is None or len(values) == 0:
            raise ValueError("values must be a non-empty list")
        return "(" + ",".join(["?"] * len(values)) + ")", values
 
    def get_trial_timeseries(
        self,
        key: str,
        study_name=None,   # None | str | Iterable[str]
        trial_ids=None,    # None | int | Iterable[int]
    ) -> pd.DataFrame:
        study_names = self._normalize_ids(study_name, cast=str, name="study_name")
        trial_ids = self._normalize_ids(trial_ids, cast=int, name="trial_ids")
 
        if study_names is None:
            if trial_ids is None:
                # cas "tout"
                study_names = self.studies()
            else:
                # déduire les studies à partir des trial_ids
                if len(trial_ids) == 0:
                    return pd.DataFrame(columns=["study_name", "trial_id", "value"])
                in_clause, in_vals = self._sql_in_clause(trial_ids)
                rows = self.con.execute(
                    f"""
                    SELECT DISTINCT studies.study_name
                    FROM trials
                    INNER JOIN studies ON studies.study_id = trials.study_id
                    WHERE trials.trial_id IN {in_clause}
                    ORDER BY studies.study_name;
                    """,
                    in_vals,
                ).fetchall()
                study_names = [r[0] for r in rows]
 
        if not study_names:
            return pd.DataFrame(columns=["study_name", "trial_id", "value"])
 
        # ------------------------------------------------------------
        # Lire et valider les specs (une par study)
        # ------------------------------------------------------------
        specs_by_study = {}
        base_spec = None
        base_cols = None
        base_value_col = None
 
        for sn in study_names:
            attrs = self.study_attrs(sn)
            spec_key = f"timeseries_spec:{key}"
            if spec_key not in attrs:
                raise KeyError(f"Missing study_user_attributes['{spec_key}'] for study '{sn}'")
 
            spec = attrs[spec_key]
            axis_names = spec["axes"]
            col_names = spec["columns"]
            value_col = spec.get("value_col", "value")
            layout = spec.get("layout", "product")
 
            axis_values = []
            for ax in axis_names:
                k = f"axis:{ax}"
                if k not in attrs:
                    raise KeyError(f"Missing study_user_attributes['{k}'] in study '{sn}'")
                axis_values.append(attrs[k])
 
            # Spec complète (incluant axes values) pour compatibilité inter-studies
            spec_full = {
                "axes": axis_names,
                "columns": col_names,
                "value_col": value_col,
                "layout": layout,
                "axis_values": axis_values,
            }
 
            if base_spec is None:
                base_spec = spec_full
                base_cols = list(col_names)
                base_value_col = value_col
            else:
                # vérifier compatibilité stricte
                if spec_full != base_spec:
                    raise ValueError(
                        f"Timeseries spec mismatch for key '{key}' between studies. "
                        f"Study '{sn}' has a different spec/axes."
                    )
 
            specs_by_study[sn] = spec_full
 
        # ------------------------------------------------------------
        # Charger les payloads trial_user_attributes (toutes studies demandées)
        # ------------------------------------------------------------
        sql = """
            SELECT studies.study_name, trial_user_attributes.trial_id, trial_user_attributes.value_json
            FROM trial_user_attributes
            INNER JOIN trials ON trial_user_attributes.trial_id = trials.trial_id
            INNER JOIN studies ON studies.study_id = trials.study_id
            WHERE trials.state="COMPLETE"
            AND trial_user_attributes.key = ?
        """
        params = [key]
 
        if study_names is not None:
            in_clause, in_vals = self._sql_in_clause(study_names)
            sql += f" AND studies.study_name IN {in_clause}"
            params.extend(in_vals)
 
        if trial_ids is not None:
            if len(trial_ids) == 0:
                return pd.DataFrame(columns=["study_name", "trial_id", *base_cols, base_value_col])
            in_clause, in_vals = self._sql_in_clause(trial_ids)
            sql += f" AND trial_user_attributes.trial_id IN {in_clause}"
            params.extend(in_vals)
 
        raw = pd.read_sql_query(sql, self.con, params=params)
 
        if raw.empty:
            return pd.DataFrame(columns=["study_name", "trial_id", *base_cols, base_value_col])
 
        # ------------------------------------------------------------
        # Expand per-study (pour réutiliser axes/spec de chaque study)
        # ------------------------------------------------------------
        out_parts = []
 
        for sn, raw_sn in raw.groupby("study_name", sort=False):
            spec = specs_by_study[sn]
            axis_names = spec["axes"]
            col_names = spec["columns"]
            value_col = spec["value_col"]
            layout = spec["layout"]
            axis_values = spec["axis_values"]
 
            arr = np.asarray([json.loads(v) for v in raw_sn["value_json"]], dtype=float)
 
            if layout == "aligned":
                n = arr.shape[1]
                for ax, vals in zip(axis_names, axis_values):
                    if len(vals) != n:
                        raise ValueError(
                            f"Aligned layout mismatch for axis '{ax}' in study '{sn}': {len(vals)} != {n}"
                        )
 
                rows = []
                for i, (_, tid) in enumerate(raw_sn[["study_name", "trial_id"]].to_numpy()):
                    d = pd.DataFrame({"study_name": sn, "trial_id": tid, value_col: arr[i]})
                    for cname, vals in zip(col_names, axis_values):
                        d[cname] = vals
                    rows.append(d)
                out_parts.append(pd.concat(rows, ignore_index=True))
                continue
 
            if layout == "product":
                if len(axis_values) == 1:
                    cols = pd.Index(axis_values[0], name=col_names[0])
                else:
                    cols = pd.MultiIndex.from_product(axis_values, names=col_names)
 
                if arr.shape[1] != len(cols):
                    raise ValueError(
                        f"Timeseries size mismatch for key '{key}' in study '{sn}': "
                        f"payload length={arr.shape[1]} expected={len(cols)}"
                    )
 
                wide = pd.DataFrame(
                    arr,
                    index=pd.MultiIndex.from_frame(raw_sn[["study_name", "trial_id"]]),
                    columns=cols,
                )
                out_parts.append(
                    wide.reset_index().melt(
                        id_vars=["study_name", "trial_id"],
                        var_name=col_names if len(col_names) > 1 else col_names[0],
                        value_name=value_col,
                    )
                )
                continue
 
            raise ValueError(f"Unknown timeseries layout '{layout}' in study '{sn}'")
 
        if not out_parts:
            return pd.DataFrame(columns=["study_name", "trial_id", *base_cols, base_value_col])
 
        return pd.concat(out_parts, ignore_index=True)
 
    def get_trial_data(
        self,
        key: str,
        study_name=None,   # None | str | Iterable[str]
        trial_ids=None,    # None | int | Iterable[int]
        raw=False,
        sep='.'
    ) -> pd.DataFrame:
        if not raw:
            df = self.get_trial_data(key=key, study_name=study_name, trial_ids=trial_ids,raw=True)
            if df.empty:
                return df
 
            # All dict payloads -> normalize
            is_dict = df["value"].map(lambda x: isinstance(x, dict))
            if is_dict.all():
                expanded = pd.json_normalize(df["value"], sep=sep)
                expanded.index = df.index
                out = pd.concat([df[["study_name", "trial_id"]], expanded], axis=1)
                return out
 
            # Mixed payloads -> best effort:
            # normalize dict rows, keep others in "value"
            if is_dict.any():
                out = df[["study_name", "trial_id"]].copy()
                out["value"] = df["value"]
 
                dict_rows = df[is_dict]
                expanded = pd.json_normalize(dict_rows["value"], sep=sep)
                expanded.index = dict_rows.index
 
                out = pd.concat([out, expanded], axis=1)
                return out
 
            # No dict payloads
            return df
 
        study_names = self._normalize_ids(study_name, cast=str, name="study_name")
        trial_ids = self._normalize_ids(trial_ids, cast=int, name="trial_ids")
 
        sql = """
            SELECT studies.study_name, trial_user_attributes.trial_id, trial_user_attributes.value_json
            FROM trial_user_attributes
            INNER JOIN trials ON trial_user_attributes.trial_id = trials.trial_id
            INNER JOIN studies ON studies.study_id = trials.study_id
            WHERE trials.state="COMPLETE"
            AND trial_user_attributes.key = ?
        """
        params = [key]
 
        if study_names is not None:
            if len(study_names) == 0:
                return pd.DataFrame(columns=["study_name", "trial_id", "value"])
            in_clause, in_vals = self._sql_in_clause(study_names)
            sql += f" AND studies.study_name IN {in_clause}"
            params.extend(in_vals)
 
        if trial_ids is not None:
            if len(trial_ids) == 0:
                return pd.DataFrame(columns=["study_name", "trial_id", "value"])
            in_clause, in_vals = self._sql_in_clause(trial_ids)
            sql += f" AND trial_user_attributes.trial_id IN {in_clause}"
            params.extend(in_vals)
 
        df = pd.read_sql_query(sql, self.con, params=params)
 
        if df.empty:
            return pd.DataFrame(columns=["study_name", "trial_id", "value"])
 
        df = df.rename(columns={"value_json": "value"})
        df["value"] = df["value"].map(json.loads)
        return df
    # ============================================================
    # isicell custom context tables
    # ============================================================
    def _has_isicell_tables(self) -> bool:
        rows = self.con.execute(
            """
            SELECT name
            FROM sqlite_master
            WHERE type='table' AND name IN ('isicell_context', 'isicell_context_entrypoints');
            """
        ).fetchall()
        names = {r[0] for r in rows}
        return "isicell_context" in names and "isicell_context_entrypoints" in names
 
    def _read_isicell_context_rows(self) -> list[tuple[str, str, str]]:
        """
        Returns rows: (key, mode, payload) for CONTEXT_ID.
        """
        if not self._has_isicell_tables():
            return []
 
        return self.con.execute(
            """
            SELECT key, mode, payload
            FROM isicell_context
            WHERE context_id = ?
            ORDER BY key;
            """,
            (self.CONTEXT_ID,),
        ).fetchall()
 
    def _read_isicell_entrypoints(self) -> dict[str, str]:
        """
        Returns {"suggest": "optim", "run": "run_simu", ...}
        """
        if not self._has_isicell_tables():
            return {}
 
        rows = self.con.execute(
            """
            SELECT kind, name
            FROM isicell_context_entrypoints
            WHERE context_id = ?;
            """,
            (self.CONTEXT_ID,),
        ).fetchall()
        return {k: n for k, n in rows}
 
    # ============================================================
    # Context decode / load
    # ============================================================
    @staticmethod
    def _decode_json_runtime_value(v: Any) -> Any:
        if isinstance(v, dict) and v.get("__kind__") == "dataframe" and v.get("orient") == "split":
            vv = v["value"]
            return pd.DataFrame(data=vv["data"], index=vv["index"], columns=vv["columns"])
        if isinstance(v, dict) and v.get("__kind__") == "series":
            s = pd.Series(v["value"])
            s.name = v.get("name")
            return s
        return v
 
    def context_value(self, name: str) -> Any:
        """
        Return a data object from the shared isicell context (json/pickle only).
        """
        rows = self._read_isicell_context_rows()
        if not rows:
            raise RuntimeError("No isicell_context table (or empty context) in this DB.")
 
        for key, mode, payload_txt in rows:
            if key != name:
                continue
 
            if mode == "json":
                return self._decode_json_runtime_value(json.loads(payload_txt))
            if mode == "pickle64b":
                return pickle64b_to_variable(payload_txt)
 
            raise TypeError(f"'{name}' is source-backed (mode='source'), not a data object")
 
        raise KeyError(f"'{name}' not found in isicell_context")
 
    def load_context(self) -> dict[str, Any]:
        """
        Build exec namespace from shared isicell_context + caller namespace.
        Source-backed symbols are exec'ed into the same namespace.
        Entrypoints are resolved from isicell_context_entrypoints.
        """
        g, _ = _capture_external_caller_namespace()
 
        rows = self._read_isicell_context_rows()
        if not rows:
            return
 
        # 1) data first
        for key, mode, payload_txt in rows:
            if mode == "json":
                g[key] = self._decode_json_runtime_value(json.loads(payload_txt))
            elif mode == "pickle64b":
                g[key] = pickle64b_to_variable(payload_txt)
 
        # 2) source next
        for key, mode, payload_txt in rows:
            if mode == "source":
                exec(payload_txt, g, g)
 
        # 3) entrypoints
        eps = self._read_isicell_entrypoints()
        if "suggest" in eps and eps["suggest"] in g:
            self.suggest_fn = g[eps["suggest"]]
        if "run" in eps and eps["run"] in g:
            self.run_fn = g[eps["run"]]
        if "eval" in eps and eps["eval"] in g:
            self.eval_fn = g[eps["eval"]]
 
    # ============================================================
    # Trial params -> structured params
    # ============================================================
    @staticmethod
    def _coerce_param_value(x: Any) -> Any:
        if x is None or isinstance(x, (int, float, bool)):
            return x
 
        s = str(x)
        if s == "True":
            return True
        if s == "False":
            return False
 
        try:
            if any(ch in s for ch in [".", "e", "E"]):
                return float(s)
            return int(s)
        except Exception:
            pass
 
        try:
            return json.loads(s)
        except Exception:
            return s
 
    def get_trial_params_flat(self, trial_id: int) -> dict[str, Any]:
        df = pd.read_sql_query(
            """
            SELECT trial_params.param_name, trial_params.param_value
            FROM trial_params
            INNER JOIN trials ON trial_params.trial_id = trials.trial_id
            WHERE trials.state="COMPLETE"
              AND trial_params.trial_id = ?;
            """,
            self.con,
            params=(int(trial_id),),
        )
        return {r["param_name"]: self._coerce_param_value(r["param_value"]) for _, r in df.iterrows()}
 
    def build_params(self, trial_id: int):
        """
        Rebuild structured params from flat optuna params using archived suggest_fn.
        Handles categorical params stored as category indices.
        """
        if self.suggest_fn is None:
            self.load_context()
 
        if self.suggest_fn is None:
            raise RuntimeError("No suggest entrypoint found in isicell_context_entrypoints.")
 
        flat = self.get_trial_params_flat(trial_id=trial_id)
 
        bounds = LHSIterator.getBoundaries(self.suggest_fn)
        for p, info in bounds.items():
            if info.get("type") == "categorical" and p in flat:
                try:
                    flat[p] = info["choices"][int(flat[p])]
                except Exception:
                    # leave as-is if already decoded
                    pass
 
        fake_trial = LHSIterator.Trial({})
        fake_trial.params = flat
        fake_trial.modeSuggest = True
        return self.suggest_fn(fake_trial)
 
    def plot_param_distributions(
        self,
        fitness_threshold: float = 1500,
        objective: int = 0,
        study_name=None,          # None | str | Iterable[str]
        trial_ids=None,           # None | int | Iterable[int]
        figsize_width: float = 11,
        row_height: float = 0.42,
        min_fig_height: float = 3.0,
        label_fontsize: int = 8,
        min_cat_label_width: float = 0.08,
        max_cat_label_len: int = 14,
    ):
        """
        Plot normalized distributions of selected trials' parameters:
        - numeric params: violin plots on x in [0, 1]
        - categorical params: stacked horizontal bars of frequencies on x in [0, 1]
 
        Selection logic:
        - keep completed trials with fitness(objective) < fitness_threshold
        - optional filter by study_name and/or trial_ids
 
        Notes
        -----
        - Requires self.suggest_fn (loaded from isicell context) to infer param bounds/types.
        - Assumes categorical params in trial_params are stored as category indices (Optuna/LHSIterator path).
        - Uses seaborn + matplotlib (imported lazily).
        """
        import matplotlib.pyplot as plt
        import seaborn as sns
 
        # Ensure suggest_fn is available
        if self.suggest_fn is None:
            self.load_context()
        if self.suggest_fn is None:
            raise RuntimeError("No suggest entrypoint loaded. Cannot infer parameter bounds/types.")
 
        # Normalize filters
        study_names = self._normalize_ids(study_name, cast=str, name="study_name")
        trial_ids_norm = self._normalize_ids(trial_ids, cast=int, name="trial_ids")
 
        # Load fitness + params
        all_fitness = self.fitness_wide()  # columns: study_name, trial_id, obj cols...
        all_params = self.params()         # MultiIndex (study_name, trial_id)
 
        if all_fitness.empty or all_params.empty:
            raise RuntimeError("No completed trials found in DB.")
 
        # Objective column existence
        if objective not in all_fitness.columns:
            raise KeyError(f"Objective column {objective!r} not found in fitness_wide(). Available: {list(all_fitness.columns)}")
 
        # Filter fitness rows by threshold
        fit = all_fitness.loc[all_fitness[objective] < fitness_threshold, ["study_name", "trial_id", objective]].copy()
 
        if study_names is not None:
            fit = fit[fit["study_name"].isin(study_names)]
 
        if trial_ids_norm is not None:
            fit = fit[fit["trial_id"].isin(trial_ids_norm)]
 
        if fit.empty:
            raise RuntimeError("No trial matches the selected filters / fitness threshold.")
 
        # Selected trials index
        selected_pairs = set(map(tuple, fit[["study_name", "trial_id"]].to_numpy()))
 
        # Filter params on same (study_name, trial_id)
        idx_study = all_params.index.get_level_values("study_name")
        idx_trial = all_params.index.get_level_values("trial_id")
        mask = [(sn, tid) in selected_pairs for sn, tid in zip(idx_study, idx_trial)]
        tmp_params = all_params.loc[mask].copy()
 
        if tmp_params.empty:
            raise RuntimeError("No parameters found for selected trials.")
 
        # Bounds from archived suggest function
        bounds = LHSIterator.getBoundaries(self.suggest_fn)
 
        # Helper numeric conversion (kept from your logic)
        tmp_num_all = tmp_params.copy()
        for c in tmp_num_all.columns:
            tmp_num_all[c] = pd.to_numeric(tmp_num_all[c], errors="coerce")
 
        # Split params by type from suggest boundaries
        num_params = [k for k, v in bounds.items() if v.get("type") in ("int", "float") and k in tmp_params.columns]
        cat_params = [k for k, v in bounds.items() if v.get("type") == "categorical" and k in tmp_params.columns]
 
        # ---------- Numeric normalized values ----------
        num_rows = []
        for p in num_params:
            info = bounds[p]
            s = pd.to_numeric(tmp_params[p], errors="coerce")
 
            low = float(info["low"])
            high = float(info["high"])
 
            if high == low:
                continue  # borne fixe -> rien à tracer
 
            if info.get("log", False):
                # log-normalization if applicable
                s = s.where(s > 0)
                vals = (np.log(s) - np.log(low)) / (np.log(high) - np.log(low))
            else:
                vals = (s - low) / (high - low)
 
            vals = vals.clip(0, 1)
 
            label = f"{p} [{low:.3f},{high:.3f}]"
            for v in vals.dropna().to_numpy():
                num_rows.append({"param": label, "value": float(v), "kind": "numeric"})
 
        num_df = pd.DataFrame(num_rows)
 
        # ---------- Categorical frequencies ----------
        cat_rows = []
        cat_segment_labels = {}  # param_label -> list[(x_center, label, freq)]
 
        for p in cat_params:
            info = bounds[p]
            choices = info["choices"]
 
            # string labels for display (stable)
            def _choice_repr(x):
                try:
                    return repr(choices[int(x)])
                except Exception:
                    # if already decoded / weird value, keep readable fallback
                    return repr(x)
 
            s = tmp_params[p].map(_choice_repr)
            choice_labels = [repr(c) for c in choices]
 
            vc = s.value_counts(normalize=True)
            freqs = [float(vc.get(cl, 0.0)) for cl in choice_labels]
 
            # Build segments on x in [0,1]
            x0 = 0.0
            param_label = f"{p} ({len(choice_labels)} choices)"
            cat_segment_labels[param_label] = []
 
            for cl, f in zip(choice_labels, freqs):
                if f <= 0:
                    continue
                cat_rows.append({
                    "param": param_label,
                    "x0": x0,
                    "width": f,
                    "choice": cl,
                    "freq": f,
                    "kind": "categorical",
                })
                cat_segment_labels[param_label].append((x0 + f / 2, cl, f))
                x0 += f
 
        cat_df = pd.DataFrame(cat_rows)
 
        # ---------- Unified parameter order ----------
        param_order = []
        if not num_df.empty:
            param_order += list(dict.fromkeys(num_df["param"].tolist()))
        if not cat_df.empty:
            param_order += [p for p in dict.fromkeys(cat_df["param"].tolist()) if p not in param_order]
 
        if len(param_order) == 0:
            raise RuntimeError("No plottable parameter found (check bounds / selected trials).")
 
        # y positions (top to bottom)
        ypos = {p: i for i, p in enumerate(param_order)}
 
        # ---------- Plot ----------
        fig_h = max(min_fig_height, row_height * len(param_order) + 1.5)
        fig = plt.figure(figsize=(figsize_width, fig_h))
        ax = plt.gca()
 
        # Numeric violins (same axis [0,1])
        if not num_df.empty:
            num_df_plot = num_df.copy()
            num_df_plot["y"] = num_df_plot["param"].map(ypos)
 
            sns.violinplot(
                data=num_df_plot,
                x="value",
                y="param",
                order=param_order,
                orient="h",
                inner="quartile",
                cut=0,
                linewidth=1,
                ax=ax,
            )
 
        # Categorical stacked bars on same axes
        if not cat_df.empty:
            for _, r in cat_df.iterrows():
                y = ypos[r["param"]]
                ax.barh(
                    y=y,
                    width=r["width"],
                    left=r["x0"],
                    height=0.65,
                    alpha=0.9,
                )
 
            # Optional text labels inside segments (only if enough width)
            for p, segs in cat_segment_labels.items():
                y = ypos[p]
                for xc, cl, f in segs:
                    if f >= min_cat_label_width:
                        txt = cl
                        if len(txt) > max_cat_label_len:
                            txt = txt[: max_cat_label_len - 3] + "..."
                        ax.text(xc, y, txt, ha="center", va="center", fontsize=label_fontsize)
 
        # Unified axis styling
        ax.set_xlim(0, 1)
        ax.set_xlabel("Normalized value / category frequency")
        ax.set_ylabel("")
 
        # Light separators to improve readability
        for i in range(len(param_order)):
            ax.axhline(i + 0.5, linewidth=0.3, alpha=0.3)
 
        plt.tight_layout()
        plt.show()