from autogluon.timeseries import TimeSeriesPredictor
from utils_vol import fetch_close_series, realized_vol, rv_to_autogluon_df

def train_bolt_small(ticker="AAPL", start="2015-01-01", interval="1d",
                     prediction_length=30, time_limit=900):
    """
    Trainiert Chronos-Bolt-Small mit AutoGluon auf CPU.
    time_limit in Sekunden (Standard: 15 min).
    """
    print(f"[AutoFT] Lade {ticker}...")
    close = fetch_close_series(ticker, start=start, interval=interval)
    rv = realized_vol(close)
    df = rv_to_autogluon_df(rv)

    predictor = TimeSeriesPredictor(
        path="/mnt/data/AutogluonChronosBoltSmall",
        prediction_length=prediction_length,
        eval_metric="WQL",
        verbosity=2,
    )

    predictor.fit(
        train_data=df,
        enable_ensemble=False,
        num_val_windows=1,
        hyperparameters={
            "Chronos": {
                "model_path": "autogluon/chronos-bolt-small",
                "fine_tune": True,
                "fine_tune_steps": 200,
                "fine_tune_lr": 1e-4,
                "context_length": 128,
                "quantile_levels": [0.1, 0.5, 0.9],
            }
        },
        time_limit=time_limit,
    )

    print("✅ Training abgeschlossen. Modellpfad:", predictor.path)
    return predictor