scripts/run_chronos_long_training.py

"""
Long Chronos-2 LoRA fine-tuning with large context window, tuned for local
CPU/memory. Single training run (Chronos does not support resuming LoRA fit).

Usage (from project root):
  PYTHONPATH=. python scripts/run_chronos_long_training.py [--device cpu] [--num-steps 4000]

Uses: context_days=28, batch_size=16 by default. Set --num-steps for total steps.
"""

from __future__ import annotations

import os
import sys
import threading
from datetime import datetime, timezone
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

# Limit CPU threads to avoid oversubscription (set before importing torch)
if "OMP_NUM_THREADS" not in os.environ:
    try:
        import multiprocessing
        n = multiprocessing.cpu_count()
        os.environ["OMP_NUM_THREADS"] = str(min(n, 10))
    except Exception:
        pass

from config.settings import OUTPUTS_DIR
from src.chronos_forecaster import (
    ChronosForecaster,
    STEPS_PER_DAY,
)
from sklearn.metrics import mean_absolute_error
import numpy as np
import pandas as pd


def _quick_val_mae(forecaster: ChronosForecaster, df: pd.DataFrame, train_ratio: float, n_windows: int = 5) -> float:
    """Compute MAE on first n_windows test windows (daytime-only) for convergence check."""
    sparse = forecaster.load_sparse_data()
    daytime_ts = set(sparse["timestamp_utc"])
    split_idx = int(len(df) * train_ratio)
    test_starts = list(range(split_idx, len(df) - STEPS_PER_DAY, STEPS_PER_DAY))[:n_windows]
    actual_list, pred_list = [], []
    for start_idx in test_starts:
        f = forecaster.forecast_day(df, start_idx, STEPS_PER_DAY, covariate_mode="all")
        actual_slice = df.iloc[start_idx : start_idx + STEPS_PER_DAY]
        daytime_mask = actual_slice["timestamp_utc"].isin(daytime_ts).values[:len(f)]
        if daytime_mask.sum() < 5:
            continue
        actual_list.append(actual_slice["A"].values[:len(f)][daytime_mask])
        pred_list.append(np.clip(f["median"].values[daytime_mask], 0, None))
    if not actual_list:
        return float("nan")
    return float(mean_absolute_error(np.concatenate(actual_list), np.concatenate(pred_list)))


def main() -> None:
    import argparse
    parser = argparse.ArgumentParser(description="Chronos-2 long LoRA training (single run)")
    parser.add_argument("--device", default="cpu", help="torch device (cpu or mps)")
    parser.add_argument("--context-days", type=int, default=28, help="context window in days")
    parser.add_argument("--batch-size", type=int, default=16, help="batch size (safe for 32GB RAM)")
    parser.add_argument("--num-steps", type=int, default=4000, help="total training steps")
    parser.add_argument("--learning-rate", type=float, default=1e-5, help="learning rate")
    parser.add_argument("--progress-minutes", type=int, default=10, help="print timestamp and progress every N minutes")
    parser.add_argument("--output-dir", type=str, default=None, help="output dir for checkpoints")
    args = parser.parse_args()

    output_dir = args.output_dir or str(OUTPUTS_DIR / "chronos_finetuned_long")
    OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)

    print("Loading data...")
    forecaster = ChronosForecaster(device=args.device, context_days=args.context_days)
    df = forecaster.load_data()
    print(f"  Grid: {len(df)} rows, context={args.context_days}d ({forecaster.context_steps} steps)")
    train_ratio = 0.75
    split_idx = int(len(df) * train_ratio)

    print("\nBaseline (zero-shot) validation MAE (5 windows)...")
    baseline_mae = _quick_val_mae(forecaster, df, train_ratio, n_windows=5)
    print(f"  {baseline_mae:.4f}")

    print(f"\nLoRA fine-tuning: {args.num_steps} steps, batch_size={args.batch_size}, lr={args.learning_rate}...")
    stop_event = threading.Event()
    interval_sec = max(1, args.progress_minutes * 60)

    def _progress_reporter():
        while True:
            if stop_event.wait(interval_sec):
                break
            ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
            print(f"[{ts}] Chronos LoRA training still in progress ({args.num_steps} steps total)...", flush=True)

    progress_thread = threading.Thread(target=_progress_reporter, daemon=True)
    progress_thread.start()
    try:
        forecaster.finetune(
            df,
            train_ratio=train_ratio,
            covariate_mode="all",
            num_steps=args.num_steps,
            learning_rate=args.learning_rate,
            batch_size=args.batch_size,
            output_dir=output_dir,
        )
    finally:
        stop_event.set()
        progress_thread.join(timeout=interval_sec + 5)

    print("\nValidation MAE after training (5 windows)...")
    val_mae = _quick_val_mae(forecaster, df, train_ratio, n_windows=5)
    print(f"  {val_mae:.4f} (baseline {baseline_mae:.4f})")

    # Full benchmark with finetuned model (append lora row to CSV)
    print("\nRunning full walk-forward benchmark (finetuned model, mode=all)...")
    sparse = forecaster.load_sparse_data()
    daytime_ts = set(sparse["timestamp_utc"])
    test_starts = list(range(split_idx, len(df) - STEPS_PER_DAY, STEPS_PER_DAY))
    all_actual, all_pred = [], []
    for start_idx in test_starts:
        f = forecaster.forecast_day(df, start_idx, STEPS_PER_DAY, covariate_mode="all")
        actual_slice = df.iloc[start_idx : start_idx + STEPS_PER_DAY]
        daytime_mask = actual_slice["timestamp_utc"].isin(daytime_ts).values[:len(f)]
        if daytime_mask.sum() < 5:
            continue
        all_actual.append(actual_slice["A"].values[:len(f)][daytime_mask])
        all_pred.append(np.clip(f["median"].values[daytime_mask], 0, None))
    lora_mae = None
    if all_actual:
        from sklearn.metrics import mean_squared_error, r2_score
        a = np.concatenate(all_actual)
        p = np.concatenate(all_pred)
        lora_mae = float(mean_absolute_error(a, p))
        lora_rmse = float(np.sqrt(mean_squared_error(a, p)))
        lora_r2 = float(r2_score(a, p))
        print(f"  LoRA / all: MAE={lora_mae:.4f}  RMSE={lora_rmse:.4f}  R²={lora_r2:.4f}  ({len(all_actual)} windows, {len(a)} steps)")

        # Load existing benchmark CSV, append lora row, save
        bench_path = OUTPUTS_DIR / "chronos_benchmark.csv"
        if bench_path.exists():
            existing = pd.read_csv(bench_path)
            lora_row = pd.DataFrame([{
                "mode": "lora / all",
                "MAE": lora_mae,
                "RMSE": lora_rmse,
                "R2": lora_r2,
                "n_windows": len(all_actual),
                "n_steps": len(a),
            }])
            combined = pd.concat([existing, lora_row], ignore_index=True)
            combined.to_csv(bench_path, index=False)
            print(f"  Appended lora row → {bench_path}")
        else:
            pd.DataFrame([{
                "mode": "lora / all", "MAE": lora_mae, "RMSE": lora_rmse, "R2": lora_r2,
                "n_windows": len(all_actual), "n_steps": len(a),
            }]).to_csv(bench_path, index=False)

    # Sample forecast plot
    print("\nGenerating sample forecast plot...")
    forecaster.plot_sample_forecast(df)

    # Summary and next steps
    print("\n" + "=" * 60)
    print("TRAINING COMPLETE — Next steps")
    print("=" * 60)
    print(f"  Checkpoints: {output_dir}")
    print(f"  Benchmark:  {OUTPUTS_DIR / 'chronos_benchmark.csv'} (lora / all row appended)")
    print(f"  Plot:       {OUTPUTS_DIR / 'chronos_forecast_sample.png'} (from finetuned model)")
    print("  • Refresh the Streamlit app: Prediction Accuracy tab will show LoRA / all in the table.")
    print("  • Sample forecast image is from the finetuned model.")
    if lora_mae is not None:
        zs_mae = 3.91  # typical zero-shot 'all' on this eval
        delta = (zs_mae - lora_mae) / zs_mae * 100
        print(f"  • LoRA MAE {lora_mae:.2f} vs zero-shot ~{zs_mae:.2f} ({delta:+.0f}% change).")
    print("=" * 60)
    print("Done.")


if __name__ == "__main__":
    main()