"""Run a Kronos-small forecast end-to-end on the K-line data bundled with the repo.

Loads NeoQuasar/Kronos-Tokenizer-base + NeoQuasar/Kronos-small from Hugging Face,
predicts 120 five-minute bars from a 400-bar context, then compares the forecast
against the held-out ground truth and saves a plot + CSV.
"""
import random
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")  # headless: save to file instead of opening a window
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

REPO_ROOT = Path(__file__).resolve().parent / "Kronos"
sys.path.insert(0, str(REPO_ROOT))

from model import Kronos, KronosTokenizer, KronosPredictor

DATA_PATH = REPO_ROOT / "tests" / "data" / "regression_input.csv"
OUT_DIR = Path(__file__).resolve().parent / "output"
LOOKBACK = 400
PRED_LEN = 120
SEED = 123


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def main() -> None:
    set_seed(SEED)
    OUT_DIR.mkdir(exist_ok=True)

    print("Loading tokenizer and model from Hugging Face Hub...")
    tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")
    model = Kronos.from_pretrained("NeoQuasar/Kronos-small")
    tokenizer.eval()
    model.eval()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model loaded: Kronos-small ({n_params / 1e6:.1f}M params)")

    predictor = KronosPredictor(model, tokenizer, device="cpu", max_context=512)

    df = pd.read_csv(DATA_PATH, parse_dates=["timestamps"])
    print(f"Data: {DATA_PATH.name}, {len(df)} rows, "
          f"{df['timestamps'].iloc[0]} .. {df['timestamps'].iloc[-1]}")

    x_df = df.loc[:LOOKBACK - 1, ["open", "high", "low", "close", "volume", "amount"]]
    x_timestamp = df.loc[:LOOKBACK - 1, "timestamps"]
    y_timestamp = df.loc[LOOKBACK:LOOKBACK + PRED_LEN - 1, "timestamps"]

    print(f"Forecasting {PRED_LEN} bars from a {LOOKBACK}-bar context (CPU)...")
    pred_df = predictor.predict(
        df=x_df,
        x_timestamp=x_timestamp,
        y_timestamp=y_timestamp,
        pred_len=PRED_LEN,
        T=1.0,
        top_p=0.9,
        sample_count=1,
        verbose=True,
    )

    print("\nForecasted Data Head:")
    print(pred_df.head())

    # Compare against held-out ground truth
    truth_df = df.loc[LOOKBACK:LOOKBACK + PRED_LEN - 1].set_index("timestamps")
    price_cols = ["open", "high", "low", "close"]
    mae = np.mean(np.abs(pred_df[price_cols].values - truth_df[price_cols].values))
    mape = np.mean(
        np.abs(pred_df[price_cols].values - truth_df[price_cols].values)
        / truth_df[price_cols].values
    ) * 100
    print(f"\nPrice MAE vs ground truth:  {mae:.4f}")
    print(f"Price MAPE vs ground truth: {mape:.2f}%")

    pred_csv = OUT_DIR / "kronos_small_forecast.csv"
    pred_df.to_csv(pred_csv, index_label="timestamps")

    # Plot: history + forecast vs ground truth
    hist = df.loc[:LOOKBACK - 1].set_index("timestamps")
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True)

    ax1.plot(hist.index, hist["close"], color="gray", linewidth=1, label="History")
    ax1.plot(truth_df.index, truth_df["close"], color="blue", linewidth=1.5, label="Ground Truth")
    ax1.plot(pred_df.index, pred_df["close"], color="red", linewidth=1.5, label="Kronos-small Forecast")
    ax1.set_ylabel("Close Price")
    ax1.legend(loc="best")
    ax1.grid(True, alpha=0.4)
    ax1.set_title(f"Kronos-small: {PRED_LEN}-step forecast ({LOOKBACK}-bar context)")

    ax2.plot(hist.index, hist["volume"], color="gray", linewidth=1, label="History")
    ax2.plot(truth_df.index, truth_df["volume"], color="blue", linewidth=1.5, label="Ground Truth")
    ax2.plot(pred_df.index, pred_df["volume"], color="red", linewidth=1.5, label="Kronos-small Forecast")
    ax2.set_ylabel("Volume")
    ax2.legend(loc="best")
    ax2.grid(True, alpha=0.4)

    plt.tight_layout()
    plot_path = OUT_DIR / "kronos_small_forecast.png"
    plt.savefig(plot_path, dpi=150)

    print(f"\nSaved forecast CSV to {pred_csv}")
    print(f"Saved plot to {plot_path}")


if __name__ == "__main__":
    main()