from __future__ import annotations

import argparse
import traceback
from pathlib import Path
import numpy as np
import pandas as pd

from .factory import create_optimizer
from .data import read_table, infer_columns, split_observed_candidates


def build_parser():
    p = argparse.ArgumentParser(description="tkbo adaptive optimization CLI")
    p.add_argument("--mode", default="ask", choices=["read", "ask", "loop"])
    p.add_argument("--infile", required=True)
    p.add_argument("--outfile", default="")
    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="physbo, sklearn_gpr, physbo_gp, custom, random, grid")
    p.add_argument("--surrogate", default="sklearn_gpr")
    p.add_argument("--acquisition", default="ei", help="ei, pi, ucb, lcb, stein, entropy")
    p.add_argument("--n-points", type=int, default=1)
    p.add_argument("--maximize", 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("--show", type=int, default=0, choices=[0, 1])
    p.add_argument("--save", type=int, default=1, choices=[0, 1])
    return p


def run(args):
    df = read_table(args.infile)
    target, features = infer_columns(df, target=args.target, features=args.features)

    print("Input:", args.infile)
    print("target:", target)
    print("features:", features)

    if args.mode == "read":
        print(df.head())
        print(df.dtypes)
        print("ndata:", len(df), "nfeatures:", len(features))
        print("observed target:", df[target].notna().sum())
        return

    X, y, observed_mask = split_observed_candidates(df, target, features)

    params = {}
    if args.model == "physbo":
        params.update(
            score_mode=args.acquisition.upper(),
            num_rand_basis=args.num_rand_basis,
            interval=args.interval,
            random_seed=args.random_seed,
            maximize=bool(args.maximize),
        )
        opt = create_optimizer(model="physbo", **params)
    elif args.model == "physbo_gp":
        opt = create_optimizer(
            model="physbo_gp",
            acquisition=args.acquisition,
            surrogate__num_rand_basis=args.num_rand_basis,
            surrogate__interval=args.interval,
            surrogate__score_mode=args.acquisition.upper(),
            surrogate__random_seed=args.random_seed,
            maximize=bool(args.maximize),
        )
    elif args.model == "sklearn_gpr":
        opt = create_optimizer(
            model="sklearn_gpr",
            acquisition=args.acquisition,
            surrogate__random_state=args.random_seed,
            maximize=bool(args.maximize),
        )
    else:
        opt = create_optimizer(
            model=args.model,
            surrogate=args.surrogate,
            acquisition=args.acquisition,
            random_seed=args.random_seed,
            maximize=bool(args.maximize),
        )

    opt.initialize(X, y=y, observed_mask=observed_mask)
    result = opt.ask(args.n_points)

    print("\nSuggested candidates:")
    for rank, idx in enumerate(result.indices, 1):
        print(f"  #{rank}: index={idx}, Excel line={idx+2}, X={X[idx]}")
        if result.scores is not None:
            print(f"       score={result.scores[rank-1]:.6g}")

    mean, std = opt.predict(X, return_std=True)
    out = df.copy()
    out["tkbo_pred_mean"] = mean
    out["tkbo_pred_std"] = std
    out["tkbo_observed"] = observed_mask
    out["tkbo_suggested"] = False
    out.loc[result.indices, "tkbo_suggested"] = True

    if args.save:
        outfile = args.outfile
        if not outfile:
            infile = Path(args.infile)
            outfile = str(infile.with_name(infile.stem + "_tkbo_ask.xlsx"))
        out.to_excel(outfile, index=False)
        print("\nSaved:", outfile)


def main():
    parser = build_parser()
    args = parser.parse_args()
    try:
        run(args)
    except Exception:
        traceback.print_exc()
        raise


if __name__ == "__main__":
    main()