"""
Extract optimization history from Optuna database and create incremental JSON file.
"""

import sys
import json
from pathlib import Path
import optuna

def main():
    if len(sys.argv) < 2:
        print("Usage: python extract_history_from_db.py <path/to/study.db>")
        sys.exit(1)

    db_file = Path(sys.argv[1])
    if not db_file.exists():
        print(f"ERROR: Database not found: {db_file}")
        sys.exit(1)

    # Load Optuna study
    storage = f"sqlite:///{db_file}"
    study_name = db_file.parent.parent.name  # Extract from path

    try:
        study = optuna.load_study(study_name=study_name, storage=storage)
    except:
        # Try to get first study if name doesn't match
        studies = optuna.get_all_study_names(storage)
        if not studies:
            print("ERROR: No studies found in database")
            sys.exit(1)
        study = optuna.load_study(study_name=studies[0], storage=storage)

    print(f"Study: {study.study_name}")
    print(f"Trials: {len(study.trials)}")

    # Extract history
    history = []
    for trial in study.trials:
        if trial.state != optuna.trial.TrialState.COMPLETE:
            continue

        record = {
            'trial_number': trial.number,
            'design_variables': trial.params,
            'results': trial.user_attrs,  # May be empty if not stored
            'objective': trial.value
        }
        history.append(record)

    # Write to JSON
    output_file = db_file.parent / 'optimization_history_incremental.json'
    with open(output_file, 'w') as f:
        json.dump(history, f, indent=2)

    print(f"\nHistory exported to: {output_file}")
    print(f"  {len(history)} completed trials")

if __name__ == "__main__":
    main()