#!/usr/bin/env python3
"""
tkbo_bayes_gp_cli.py

Bayesian optimization / adaptive learning CLI using tkbo.

This program is intended as a PHYSBO/tkbo learning example.
It keeps a practical command-line interface while delegating optimizer,
surrogate, and acquisition behavior to tkbo.

Supported input
---------------
CSV / Excel table.

Column rules follow tkbo style:
    max:xxx    objective to maximize
    min:xxx    objective to minimize; internally transformed to -y
    =1.23:xxx  objective to approach target value; transformed to -(y-1.23)^2
    -xxx       excluded column
    numeric columns other than target are descriptors/features

Modes
-----
read     : show detected target/features and observed/candidate counts
ask      : fit current observed data and suggest next candidate rows
predict  : fit or load model and output posterior mean/std for all valid rows
tell     : write newly measured values back to a copy of the table
loop     : simple ask/tell loop with a Python expression simulator

Examples
--------
Read data:
    python tkbo_bayes_gp_cli.py --mode read --infile data.xlsx

sklearn GPR + EI:
    python tkbo_bayes_gp_cli.py --mode ask --infile data.xlsx --model sklearn_gpr --acquisition ei --n-points 3

PHYSBO native backend:
    python tkbo_bayes_gp_cli.py --mode ask --infile data.xlsx --model physbo --acquisition EI

PHYSBO posterior surrogate + tkbo acquisition:
    python tkbo_bayes_gp_cli.py --mode ask --infile data.xlsx --model physbo_gp --acquisition stein

Save fitted optimizer/model payload:
    python tkbo_bayes_gp_cli.py --mode ask --infile data.xlsx --model-file state.joblib

Reuse saved optimizer/model payload:
    python tkbo_bayes_gp_cli.py --mode predict --infile data.xlsx --load-model state.joblib

Update measured values:
    python tkbo_bayes_gp_cli.py --mode tell --infile data.xlsx --tell-values 10:3.14,11:2.71
"""

from __future__ import annotations

import argparse
import math
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional, Sequence

import joblib
import numpy as np
import pandas as pd

# tkbo imports.  This script intentionally uses tkbo rather than reimplementing
# the optimizer layer.
from tkbo.factory import create_optimizer
from tkbo.data import read_table, infer_columns, split_observed_candidates


@dataclass
class DataBundle:
    df: pd.DataFrame
    target: str
    target_label: str
    target_mode: str
    target_value: Optional[float]
    features: list[str]
    X: np.ndarray
    y_original: np.ndarray
    y_bo: np.ndarray
    observed_mask: np.ndarray
    valid_feature_mask: np.ndarray


# =============================================================================
# Data handling
# =============================================================================


def parse_target_column(name: str) -> tuple[str, str, Optional[float]]:
    s = str(name)
    m = re.match(r"^=([+-]?\d*\.?\d+(?:[eE][+-]?\d+)?):(.*)$", s)
    if m:
        return m.group(2), "value", float(m.group(1))

    m = re.match(r"^(max\d*|t\d*|o\d*):(.*)$", s, flags=re.IGNORECASE)
    if m:
        return m.group(2), "max", None

    m = re.match(r"^(min\d*):(.*)$", s, flags=re.IGNORECASE)
    if m:
        return m.group(2), "min", None

    return s, "max", None


def transform_target_for_bo(y: np.ndarray, mode: str, target_value: Optional[float]) -> np.ndarray:
    y = np.asarray(y, dtype=float).copy()
    if mode == "min":
        return -y
    if mode == "value":
        if target_value is None:
            raise ValueError("target_value is required for '=value:' target mode")
        return -((y - float(target_value)) ** 2)
    return y


def inverse_mean_for_display(y_bo_mean: np.ndarray, mode: str, target_value: Optional[float]) -> np.ndarray:
    """Best-effort inverse for display.

    For '=value:' mode, the inverse is not unique, so BO-scale mean is kept.
    """
    y_bo_mean = np.asarray(y_bo_mean, dtype=float)
    if mode == "min":
        return -y_bo_mean
    return y_bo_mean


def load_data(
    path: str | Path,
    target: str | None = None,
    features: str | list[str] | None = None,
    sheet_name: str | int = 0,
) -> DataBundle:
    df = read_table(path, sheet_name=sheet_name)
    target, features = infer_columns(df, target=target, features=features)
    label, mode, target_value = parse_target_column(target)

    # Use tkbo's split function first, then add feature-NaN handling.
    X, y_original, observed_mask = split_observed_candidates(df, target, features)
    X = np.asarray(X, dtype=float)
    y_original = np.asarray(y_original, dtype=float).reshape(-1)
    valid_feature_mask = ~np.any(np.isnan(X), axis=1)
    y_bo = transform_target_for_bo(y_original, mode, target_value)

    return DataBundle(
        df=df,
        target=target,
        target_label=label,
        target_mode=mode,
        target_value=target_value,
        features=features,
        X=X,
        y_original=y_original,
        y_bo=y_bo,
        observed_mask=observed_mask,
        valid_feature_mask=valid_feature_mask,
    )


def write_table(df: pd.DataFrame, path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.suffix.lower() in {".xlsx", ".xlsm", ".xls"}:
        df.to_excel(path, index=False)
    else:
        df.to_csv(path, index=False)


def default_outfile(infile: str | Path, suffix: str = "-tkbo-predict.xlsx") -> str:
    p = Path(infile)
    return str(p.with_name(p.stem + suffix))


def default_model_file(infile: str | Path) -> str:
    p = Path(infile)
    return str(p.with_name(p.stem + "-tkbo-state.joblib"))


# =============================================================================
# tkbo optimizer construction
# =============================================================================


def build_tkbo_optimizer(args: argparse.Namespace):
    """Create tkbo optimizer from CLI arguments.

    model choices:
      sklearn_gpr : custom backend + sklearn GaussianProcessRegressor surrogate
      physbo      : PHYSBO native backend using PHYSBO EI/PI/TS
      physbo_gp   : PHYSBO posterior surrogate + tkbo acquisition
      random      : random search backend
      grid        : sequential grid backend
      custom      : explicit surrogate/acquisition
    """
    model = args.model.lower()
    acquisition = args.acquisition

    if model == "physbo":
        return create_optimizer(
            model="physbo",
            score_mode=acquisition.upper(),
            num_rand_basis=args.num_rand_basis,
            interval=args.interval,
            random_seed=args.random_seed,
            maximize=bool(args.maximize),
        )

    if model == "physbo_gp":
        return create_optimizer(
            model="physbo_gp",
            acquisition=acquisition.lower(),
            surrogate__num_rand_basis=args.num_rand_basis,
            surrogate__interval=args.interval,
            surrogate__score_mode=args.physbo_score.upper(),
            surrogate__random_seed=args.random_seed,
            acquisition__kappa=args.kappa,
            acquisition__xi=args.xi,
            acquisition__eps=args.stein_eps,
            acquisition__std_weight=args.stein_std_weight,
            acquisition__grad_weight=args.stein_grad_weight,
            maximize=bool(args.maximize),
        )

    if model == "sklearn_gpr":
        return create_optimizer(
            model="sklearn_gpr",
            acquisition=acquisition.lower(),
            surrogate__alpha=args.alpha,
            surrogate__normalize_y=bool(args.normalize_y),
            surrogate__n_restarts_optimizer=args.n_restarts_optimizer,
            surrogate__random_state=args.random_seed,
            surrogate__length_scale=args.length_scale,
            surrogate__noise_level=args.noise_level,
            acquisition__kappa=args.kappa,
            acquisition__xi=args.xi,
            acquisition__eps=args.stein_eps,
            acquisition__std_weight=args.stein_std_weight,
            acquisition__grad_weight=args.stein_grad_weight,
            maximize=bool(args.maximize),
        )

    if model in {"random", "grid"}:
        return create_optimizer(model=model, random_seed=args.random_seed)

    return create_optimizer(
        model=args.model,
        surrogate=args.surrogate,
        acquisition=acquisition.lower(),
        random_seed=args.random_seed,
        maximize=bool(args.maximize),
    )


def initialize_optimizer(args: argparse.Namespace, bundle: DataBundle):
    valid_rows = bundle.valid_feature_mask
    X_valid = bundle.X[valid_rows]
    y_valid = bundle.y_bo[valid_rows]
    observed_valid = bundle.observed_mask[valid_rows]

    if not args.load_model and int(observed_valid.sum()) == 0 and args.model not in {"random", "grid"}:
        raise ValueError("No observed rows are available. Fill at least one target value, or use --model random/grid.")

    opt = build_tkbo_optimizer(args)
    opt.initialize(X_valid, y=y_valid, observed_mask=observed_valid)
    return opt, np.where(valid_rows)[0]


def load_optimizer_payload(path: str | Path):
    payload = joblib.load(path)
    if isinstance(payload, dict) and "optimizer" in payload:
        return payload
    return {"optimizer": payload}


def save_optimizer_payload(path: str | Path, optimizer: Any, bundle: DataBundle, args: argparse.Namespace) -> None:
    payload = {
        "optimizer": optimizer,
        "target": bundle.target,
        "target_label": bundle.target_label,
        "target_mode": bundle.target_mode,
        "target_value": bundle.target_value,
        "features": bundle.features,
        "model": args.model,
        "surrogate": args.surrogate,
        "acquisition": args.acquisition,
    }
    joblib.dump(payload, path)


# =============================================================================
# Output helpers
# =============================================================================


def safe_predict_optimizer(opt: Any, X_valid: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    try:
        mean, std = opt.predict(X_valid, return_std=True)
        return np.asarray(mean, dtype=float).reshape(-1), np.asarray(std, dtype=float).reshape(-1)
    except Exception:
        n = len(X_valid)
        return np.full(n, np.nan), np.full(n, np.nan)


def safe_acquisition_optimizer(opt: Any, X_valid: np.ndarray) -> np.ndarray:
    try:
        score = opt.acquisition(X_valid)
        return np.asarray(score, dtype=float).reshape(-1)
    except Exception:
        return np.full(len(X_valid), np.nan)


def build_output_frame(
    bundle: DataBundle,
    valid_row_indices: np.ndarray,
    mean_valid: np.ndarray,
    std_valid: np.ndarray,
    score_valid: Optional[np.ndarray] = None,
    selected_original_indices: Optional[np.ndarray] = None,
) -> pd.DataFrame:
    out = bundle.df.copy()
    out["__observed__"] = bundle.observed_mask
    out["__valid_features__"] = bundle.valid_feature_mask
    out["__pred_mean_bo__"] = np.nan
    out["__pred_std_bo__"] = np.nan
    out["__pred_mean_display__"] = np.nan
    out.loc[out.index[valid_row_indices], "__pred_mean_bo__"] = mean_valid
    out.loc[out.index[valid_row_indices], "__pred_std_bo__"] = std_valid
    out.loc[out.index[valid_row_indices], "__pred_mean_display__"] = inverse_mean_for_display(
        mean_valid, bundle.target_mode, bundle.target_value
    )

    if score_valid is not None:
        out["__acquisition_score__"] = np.nan
        out.loc[out.index[valid_row_indices], "__acquisition_score__"] = score_valid
        out["__suggest_rank__"] = np.nan
        finite_valid = np.where(np.isfinite(score_valid))[0]
        order_valid = finite_valid[np.argsort(score_valid[finite_valid])[::-1]]
        for rank, valid_pos in enumerate(order_valid, 1):
            original_idx = valid_row_indices[valid_pos]
            out.loc[out.index[original_idx], "__suggest_rank__"] = rank

    if selected_original_indices is not None:
        out["__selected__"] = False
        for idx in selected_original_indices:
            out.loc[out.index[int(idx)], "__selected__"] = True

    return out


# =============================================================================
# CLI modes
# =============================================================================


def mode_read(args: argparse.Namespace) -> None:
    bundle = load_data(args.infile, target=args.target, features=args.features, sheet_name=args.sheet_name)
    observed = bundle.observed_mask & bundle.valid_feature_mask
    candidates = (~bundle.observed_mask) & bundle.valid_feature_mask
    invalid = ~bundle.valid_feature_mask

    print(f"Input: {args.infile}")
    print(f"Rows: {len(bundle.df)}")
    print(f"Target: {bundle.target}")
    print(f"Target label: {bundle.target_label}")
    print(f"Target mode: {bundle.target_mode}")
    print(f"Target value: {bundle.target_value}")
    print(f"Features: {bundle.features}")
    print(f"Observed rows with valid features: {int(observed.sum())}")
    print(f"Candidate rows with blank target and valid features: {int(candidates.sum())}")
    print(f"Rows skipped due to NaN in features: {int(invalid.sum())}")
    print(bundle.df.head(args.head).to_string(index=False))


def mode_ask_or_predict(args: argparse.Namespace) -> None:
    bundle = load_data(args.infile, target=args.target, features=args.features, sheet_name=args.sheet_name)

    if args.load_model:
        payload = load_optimizer_payload(args.load_model)
        opt = payload["optimizer"]
        if args.features is None and "features" in payload:
            # Reload with saved features/target for safer predict reuse.
            bundle = load_data(
                args.infile,
                target=args.target or payload.get("target"),
                features=payload.get("features"),
                sheet_name=args.sheet_name,
            )
            valid_row_indices = np.where(bundle.valid_feature_mask)[0]
        else:
            valid_row_indices = np.where(bundle.valid_feature_mask)[0]
    else:
        opt, valid_row_indices = initialize_optimizer(args, bundle)

    X_valid = bundle.X[valid_row_indices]
    mean_valid, std_valid = safe_predict_optimizer(opt, X_valid)
    score_valid = None
    selected_original_indices = np.array([], dtype=int)

    if args.mode == "ask":
        result = opt.ask(n_points=args.n_points)
        selected_valid_positions = np.asarray(result.indices, dtype=int).reshape(-1)
        selected_original_indices = valid_row_indices[selected_valid_positions]
        score_valid = safe_acquisition_optimizer(opt, X_valid)

        print("\nSuggested candidates:")
        for rank, valid_pos in enumerate(selected_valid_positions, 1):
            original_idx = int(valid_row_indices[valid_pos])
            excel_row = original_idx + 2
            x_dict = {col: bundle.df.iloc[original_idx][col] for col in bundle.features}
            score = np.nan
            if result.scores is not None and rank - 1 < len(result.scores):
                score = float(result.scores[rank - 1])
            elif score_valid is not None and valid_pos < len(score_valid):
                score = float(score_valid[valid_pos])
            print(
                f"  #{rank}: index={original_idx}, tkbo_index={valid_pos}, Excel row={excel_row}, "
                f"mean_bo={mean_valid[valid_pos]:.6g}, std={std_valid[valid_pos]:.6g}, "
                f"score={score:.6g}, X={x_dict}"
            )

    if args.save:
        outfile = args.outfile or default_outfile(args.infile)
        out = build_output_frame(
            bundle,
            valid_row_indices=valid_row_indices,
            mean_valid=mean_valid,
            std_valid=std_valid,
            score_valid=score_valid,
            selected_original_indices=selected_original_indices,
        )
        write_table(out, outfile)
        print(f"Saved output: {outfile}")

    if args.save_model and not args.load_model:
        model_file = args.model_file or default_model_file(args.infile)
        save_optimizer_payload(model_file, opt, bundle, args)
        print(f"Saved tkbo optimizer state: {model_file}")


def mode_tell(args: argparse.Namespace) -> None:
    if not args.tell_values:
        raise ValueError("--mode tell requires --tell-values, e.g. --tell-values 10:3.14,11:2.71")

    bundle = load_data(args.infile, target=args.target, features=args.features, sheet_name=args.sheet_name)
    df = bundle.df.copy()
    for pair in args.tell_values.split(","):
        idx_s, val_s = pair.split(":", 1)
        idx = int(idx_s.strip())
        val = float(val_s.strip())
        if idx < 0 or idx >= len(df):
            raise IndexError(f"tell index out of range: {idx}")
        df.loc[df.index[idx], bundle.target] = val

    outfile = args.outfile or default_outfile(args.infile, suffix="-tkbo-tell.xlsx")
    write_table(df, outfile)
    print(f"Saved updated table: {outfile}")


def mode_loop(args: argparse.Namespace) -> None:
    """Simple simulator loop.

    The simulator is a Python expression evaluated with feature names available
    as variables. Example:
        --simulator "sin(x1) + 0.1*x2"

    This is for learning only. For real experiments, use ask -> measure -> tell.
    """
    if not args.simulator:
        raise ValueError("--mode loop requires --simulator, e.g. --simulator 'sin(x)'")

    bundle = load_data(args.infile, target=args.target, features=args.features, sheet_name=args.sheet_name)
    opt, valid_row_indices = initialize_optimizer(args, bundle)
    y_current = bundle.y_bo.copy()
    observed_current = bundle.observed_mask.copy()

    safe_math = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
    safe_math.update({"np": np})

    history: list[dict[str, Any]] = []
    for step in range(args.loop_steps):
        result = opt.ask(n_points=args.n_points)
        if len(result.indices) == 0:
            break
        for valid_pos in np.asarray(result.indices, dtype=int).reshape(-1):
            original_idx = int(valid_row_indices[valid_pos])
            local_vars = {col: float(bundle.df.iloc[original_idx][col]) for col in bundle.features}
            y_new_original = float(eval(args.simulator, {"__builtins__": {}}, {**safe_math, **local_vars}))
            y_new_bo = float(transform_target_for_bo(np.array([y_new_original]), bundle.target_mode, bundle.target_value)[0])
            opt.tell(int(valid_pos), y_new_bo)
            y_current[original_idx] = y_new_bo
            observed_current[original_idx] = True
            history.append(
                {
                    "step": step + 1,
                    "index": original_idx,
                    "tkbo_index": int(valid_pos),
                    "excel_row": original_idx + 2,
                    "y_original": y_new_original,
                    "y_bo": y_new_bo,
                }
            )
            print(f"step={step + 1} index={original_idx} Excel row={original_idx + 2} y={y_new_original:.6g}")

    if args.save:
        out = bundle.df.copy()
        for h in history:
            out.loc[out.index[h["index"]], bundle.target] = h["y_original"]
        outfile = args.outfile or default_outfile(args.infile, suffix="-tkbo-loop.xlsx")
        write_table(out, outfile)
        print(f"Saved loop result: {outfile}")


# =============================================================================
# argparse
# =============================================================================


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Bayesian optimization CLI using tkbo")

    p.add_argument("--mode", default="ask", choices=["read", "ask", "predict", "tell", "loop"])
    p.add_argument("--infile", required=True)
    p.add_argument("--outfile", default="")
    p.add_argument("--sheet-name", default=0)
    p.add_argument("--target", default=None)
    p.add_argument("--features", default=None, help="Comma-separated feature columns")

    p.add_argument("--model", default="sklearn_gpr", help="sklearn_gpr, physbo, physbo_gp, custom, random, grid")
    p.add_argument("--surrogate", default="sklearn_gpr")
    p.add_argument("--acquisition", default="ei", help="ei, pi, ucb, lcb, entropy, stein; EI/PI/TS for physbo native")
    p.add_argument("--physbo-score", default="EI", help="Score used while PHYSBO surrogate builds posterior")
    p.add_argument("--n-points", type=int, default=1)
    p.add_argument("--maximize", type=int, default=1, choices=[0, 1])

    p.add_argument("--standardize", type=int, default=1, choices=[0, 1], help="Kept for interface compatibility; sklearn_gpr tkbo surrogate may handle scaling internally depending on implementation")
    p.add_argument("--normalize-y", type=int, default=1, choices=[0, 1])
    p.add_argument("--random-seed", type=int, default=None)
    p.add_argument("--num-rand-basis", type=int, default=200)
    p.add_argument("--interval", type=int, default=0)

    p.add_argument("--alpha", type=float, default=1.0e-10)
    p.add_argument("--length-scale", type=float, default=1.0)
    p.add_argument("--noise-level", type=float, default=None)
    p.add_argument("--n-restarts-optimizer", type=int, default=10)

    p.add_argument("--xi", type=float, default=0.0)
    p.add_argument("--kappa", type=float, default=2.0)
    p.add_argument("--stein-eps", type=float, default=1.0e-5)
    p.add_argument("--stein-std-weight", type=float, default=1.0)
    p.add_argument("--stein-grad-weight", type=float, default=1.0)

    p.add_argument("--save", type=int, default=1, choices=[0, 1])
    p.add_argument("--save-model", type=int, default=1, choices=[0, 1])
    p.add_argument("--model-file", default="")
    p.add_argument("--load-model", default="")
    p.add_argument("--tell-values", default="")
    p.add_argument("--loop-steps", type=int, default=5)
    p.add_argument("--simulator", default="")
    p.add_argument("--head", type=int, default=5)
    p.add_argument("--debug", type=int, default=0, choices=[0, 1])
    return p


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if isinstance(args.sheet_name, str) and args.sheet_name.isdigit():
        args.sheet_name = int(args.sheet_name)

    try:
        if args.mode == "read":
            mode_read(args)
        elif args.mode in {"ask", "predict"}:
            mode_ask_or_predict(args)
        elif args.mode == "tell":
            mode_tell(args)
        elif args.mode == "loop":
            mode_loop(args)
        else:
            raise ValueError(f"Unknown mode: {args.mode}")
    except Exception as exc:
        print(f"ERROR: {exc}")
        if args.debug:
            raise
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())