"""Reset bracket stiffness optimization study by deleting database and optionally cleaning Nastran files."""
import optuna
from pathlib import Path
import argparse

study_dir = Path(__file__).parent
storage = f"sqlite:///{study_dir / '2_results' / 'study.db'}"
study_name = "bracket_stiffness_optimization_atomizerfield"

def clean_nastran_files():
    """Remove old Nastran solver output files."""
    model_dir = study_dir / "1_setup" / "model"
    nastran_extensions = ['*.op2', '*.f06', '*.log', '*.f04', '*.pch', '*.DBALL', '*.MASTER', '*.asg', '*.diag']
    temp_patterns = ['_temp*.txt', '*_temp_*']

    deleted = []
    for pattern in nastran_extensions + temp_patterns:
        for f in model_dir.glob(pattern):
            try:
                f.unlink()
                deleted.append(f.name)
            except Exception as e:
                print(f"[WARNING] Could not delete {f.name}: {e}")

    return deleted

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Reset bracket optimization study")
    parser.add_argument('--clean', action='store_true', help='Also clean Nastran output files')
    parser.add_argument('--clean-only', action='store_true', help='Only clean Nastran files, keep database')
    args = parser.parse_args()

    if not args.clean_only:
        try:
            optuna.delete_study(study_name=study_name, storage=storage)
            print(f"[OK] Deleted study: {study_name}")
        except KeyError:
            print(f"[INFO] Study '{study_name}' not found (database may not exist)")
        except Exception as e:
            print(f"[ERROR] Error: {e}")

    if args.clean or args.clean_only:
        deleted = clean_nastran_files()
        if deleted:
            print(f"[OK] Deleted {len(deleted)} Nastran files:")
            for f in deleted[:5]:
                print(f"     - {f}")
            if len(deleted) > 5:
                print(f"     ... and {len(deleted) - 5} more")
        else:
            print("[INFO] No Nastran files to clean")