# train_autogluon.py
from autogluon.timeseries import TimeSeriesPredictor, TimeSeriesDataFrame
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,              # Sekunden (15 Min). Bei Bedarf anpassen.
):
    """
    Trainiert Chronos-Bolt-Small auf CPU via AutoGluon mit CPU-freundlichen Limits.
    Explizite Business-Day-Frequenz ('B') verhindert Frequency-Fehler.
    """
    print(f"[AutoFT] Lade {ticker} ...")
    close = fetch_close_series(ticker, start=start, interval=interval)
    rv = realized_vol(close)

    # tidy DataFrame: columns = item_id, timestamp, target
    df = rv_to_autogluon_df(rv)

    # TimeSeriesDataFrame mit expliziter Frequenz erzeugen
    tsdf = TimeSeriesDataFrame.from_data_frame(
        df,
        id_column="item_id",
        timestamp_column="timestamp",
        # KEIN target_column-Argument in AG 1.4.0 – 'target' wird implizit erkannt
        freq="B",
    )
    # auf reguläres Business-Day-Gitter bringen (Lücken = NaN)
    tsdf = tsdf.convert_frequency("B")

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

    predictor.fit(
        train_data=tsdf,
        enable_ensemble=False,
        num_val_windows=1,
        hyperparameters={
            "Chronos": {
                "model_path": "autogluon/chronos-bolt-small",
                "fine_tune": True,
                "fine_tune_steps": 200,     # klein halten für CPU
                "fine_tune_lr": 1e-4,
                "context_length": 128,      # klein halten für CPU
                "quantile_levels": [0.1, 0.5, 0.9],
            }
        },
        time_limit=time_limit,               # harter Cap, damit HF nicht timeoutet
    )

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