import optuna
import time
import json
import os
import argparse
def main():

parser = argparse.ArgumentParser(description='Optuna performance comparison')

parser.add_argument('-t','--n_trials', type=int, default=1000, help='Number of trials')

parser.add_argument('-b','--branch_name', type=str, default='master', help='Branch name for output')

parser.add_argument('-o','--n_objectives', type=int, default=4, help='Number of objective return values (1-4)')

parser.add_argument('-m','--multivariate', action='store_true', help='Multivariate')
args = parser.parse_args()

# Validate n_objectives
if args.n_objectives < 1 or args.n_objectives > 9:
    raise ValueError("n_objectives must be between 1 and 4")

def objective(trial):
    w = trial.suggest_float("w", -5, 5)
    x = trial.suggest_float("x", -5, 5)
    y = trial.suggest_float("y", -5, 5)
    z = trial.suggest_float("z", -5, 5)
    
    objectives = [
        w ** 2 + x ** 2 + y ** 2 + z ** 2, 
        w + x + y + z, 
        (x-2)**2 + (y-2)**2 + (z-2)**2, 
        (w-1)**2 + (x-1)**2 + (y-1)**2 + (z-1)**2,
        w + x + y,
        x + y + z,
        w + y + z,
        w + x + z,
        w * x * y,
    ][:args.n_objectives]
    
    return tuple(objectives)

sampler = optuna.samplers.TPESampler(
    multivariate= args.multivariate,
    seed=42,
)
times = {}

study = optuna.create_study(
    directions=["minimize"]*args.n_objectives,
    sampler=sampler,
)

start_time = time.time()

for i in range(args.n_trials):
    trial = study.ask()
    values = objective(trial)
    study.tell(trial, values)
    print(f"Trial#{trial.number}: {values=} with params={trial.params}")
    diff = time.time() - start_time
    times[i] = diff

dir_name = f"{args.branch_name}_{'multivariate' if args.multivariate else 'independent'}_{args.n_objectives}"
os.makedirs(dir_name, exist_ok=True)

with open(f"./{dir_name}/time.json", "w") as f:
    json.dump(times, f, indent=4)

with open(f"./{dir_name}/best_trials.json", "w") as f:
    best_trials = [trial.values for trial in study.best_trials]
    json.dump(best_trials, f, indent=4)

if name == "main":

main()