#!/usr/bin/env python3
"""
kronos_reference_capture.py
===========================
First build artifact for the "Kronos-small inference in native MQL5" project.

What it does
------------
1. Loads Kronos-small + Kronos-Tokenizer-base and PINS their exact resolved
   architecture to JSON (so the tokenizer config gets captured automatically).
2. Freezes ONE fully reproducible input window and saves it as the contract the
   MQL5 side will read verbatim -- MQL5 never re-synthesizes the input, it reads
   these bytes, so the data generation here does not need to be portable.
3. Replicates KronosPredictor's preprocessing EXACTLY (per-window z-score with
   eps=1e-5, population std, clip +/-5) and saves every intermediate.
4. Captures a DETERMINISTIC verification ladder -- no sampling anywhere -- at each
   model boundary, so each MQL5 component can be checked against ground truth as
   it is built:

       x_norm  ->  (s1_ids, s2_ids)  ->  s1_logits  ->  s2_logits  ->  reconstruction

5. (Optional) one SEEDED end-to-end forecast as a coarse full-pipeline check.
   This one is stochastic; it is documented as such and is not a bit-exact target.

Every array is written twice:
  - <name>.npy                 for Python-side inspection
  - <name>.bin + manifest.json little-endian, row-major, for MQL5 FileReadArray

Usage
-----
Run from inside the cloned Kronos repo (so `from model import ...` resolves), or
point --kronos_repo at it:

    python kronos_reference_capture.py --kronos_repo /path/to/Kronos --out ./kronos_refs

Notes
-----
- CPU + float32 are forced for reproducibility. Do not run this on GPU.
- np.std defaults to population std (ddof=0); KronosPredictor relies on that, so
  MQL5 must also use population std (divide by N, not N-1).
"""

import argparse
import json
import os
import sys

import numpy as np


def to_py(v):
    """Make a config value JSON-serializable."""
    try:
        import torch
        if isinstance(v, torch.Tensor):
            return v.tolist()
    except Exception:
        pass
    if isinstance(v, (np.integer,)):
        return int(v)
    if isinstance(v, (np.floating,)):
        return float(v)
    if isinstance(v, (np.ndarray,)):
        return v.tolist()
    return v


def dump_config(obj, curated_keys):
    """Best-effort capture of a HF-mixin module's resolved config."""
    cfg = {}
    # 1) try the stored hub-mixin config dict
    for attr in ("config", "_hub_mixin_config", "_config"):
        c = getattr(obj, attr, None)
        if isinstance(c, dict):
            cfg.update({k: to_py(v) for k, v in c.items()})
    # 2) scrape curated attributes directly off the instantiated module
    for k in curated_keys:
        if hasattr(obj, k):
            cfg[k] = to_py(getattr(obj, k))
    return cfg


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--kronos_repo", default=".",
                    help="Path to cloned Kronos repo (the dir containing model/)")
    ap.add_argument("--out", default="./kronos_refs")
    ap.add_argument("--tokenizer_id", default="NeoQuasar/Kronos-Tokenizer-base")
    ap.add_argument("--model_id", default="NeoQuasar/Kronos-small")
    ap.add_argument("--lookback", type=int, default=256, help="context length, <=512")
    ap.add_argument("--pred_len", type=int, default=16)
    ap.add_argument("--clip", type=float, default=5.0)
    ap.add_argument("--eps", type=float, default=1e-5)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--skip_forecast", action="store_true",
                    help="Skip the seeded end-to-end forecast (stochastic) step")
    args = ap.parse_args()

    sys.path.insert(0, os.path.abspath(args.kronos_repo))
    import torch
    import pandas as pd
    from model import Kronos, KronosTokenizer

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    torch.set_grad_enabled(False)
    device = "cpu"  # reproducible references only

    os.makedirs(args.out, exist_ok=True)
    manifest = {"arrays": {}, "scalars": {}, "notes": []}

    def save(name, arr, dtype):
        a = np.ascontiguousarray(np.asarray(arr).astype(dtype))
        np.save(os.path.join(args.out, name + ".npy"), a)
        a.tofile(os.path.join(args.out, name + ".bin"))
        manifest["arrays"][name] = {
            "file": name + ".bin",
            "dtype": np.dtype(dtype).name,
            "shape": list(a.shape),
            "order": "little-endian, row-major (C)",
        }
        print(f"  saved {name:24s} shape={tuple(a.shape)} dtype={np.dtype(dtype).name}")

    # ------------------------------------------------------------------ load
    print("Loading models (CPU, float32) ...")
    tok = KronosTokenizer.from_pretrained(args.tokenizer_id).to(device).float().eval()
    mdl = Kronos.from_pretrained(args.model_id).to(device).float().eval()

    tok_cfg = dump_config(tok, [
        "d_in", "d_model", "n_heads", "ff_dim", "enc_layers", "dec_layers",
        "s1_bits", "s2_bits", "codebook_dim",
    ])
    mdl_cfg = dump_config(mdl, [
        "d_model", "n_heads", "ff_dim", "n_layers", "s1_bits", "s2_bits",
        "s1_vocab_size", "s2_vocab_size", "learn_te",
    ])
    with open(os.path.join(args.out, "config_tokenizer.json"), "w") as f:
        json.dump(tok_cfg, f, indent=2)
    with open(os.path.join(args.out, "config_predictor.json"), "w") as f:
        json.dump(mdl_cfg, f, indent=2)
    print("Pinned tokenizer config:", json.dumps(tok_cfg))
    print("Pinned predictor config:", json.dumps(mdl_cfg))

    # ------------------------------------------------- freeze an input window
    L, P = args.lookback, args.pred_len
    price_cols = ["open", "high", "low", "close"]
    feat_cols = price_cols + ["volume", "amount"]
    stamp_cols = ["minute", "hour", "weekday", "day", "month"]

    idx = pd.date_range("2023-01-02 00:00:00", periods=L + P, freq="h")
    rw = np.cumsum(np.random.randn(L + P).astype(np.float64)) * 0.5 + 100.0
    close = rw
    open_ = np.concatenate([[close[0]], close[:-1]])
    high = np.maximum(open_, close) + np.abs(np.random.randn(L + P)) * 0.2
    low = np.minimum(open_, close) - np.abs(np.random.randn(L + P)) * 0.2
    volume = np.abs(np.random.randn(L + P)) * 1000.0 + 500.0
    amount = volume * (open_ + high + low + close) / 4.0
    df_all = pd.DataFrame(
        {"open": open_, "high": high, "low": low, "close": close,
         "volume": volume, "amount": amount}, index=idx)

    hist_df = df_all.iloc[:L].copy()
    x_ts = pd.Series(df_all.index[:L])
    y_ts = pd.Series(df_all.index[L:L + P])

    def stamps(ts):
        out = np.zeros((len(ts), 5), dtype=np.float32)
        out[:, 0] = ts.dt.minute
        out[:, 1] = ts.dt.hour
        out[:, 2] = ts.dt.weekday
        out[:, 3] = ts.dt.day
        out[:, 4] = ts.dt.month
        return out

    x_raw = hist_df[feat_cols].values.astype(np.float32)        # (L, 6)
    x_stamp = stamps(x_ts)                                       # (L, 5)
    y_stamp = stamps(y_ts)                                       # (P, 5)

    # --------------------------------------------- preprocessing (exact copy)
    x_mean = x_raw.mean(axis=0)                                  # (6,)
    x_std = x_raw.std(axis=0)                                    # (6,) population std, ddof=0
    x_norm = (x_raw - x_mean) / (x_std + args.eps)
    x_norm = np.clip(x_norm, -args.clip, args.clip).astype(np.float32)

    print("\nStage 0  preprocessing")
    save("x_raw", x_raw, np.float32)
    save("x_norm", x_norm, np.float32)
    save("x_mean", x_mean, np.float32)
    save("x_std", x_std, np.float32)
    save("x_stamp", x_stamp, np.float32)
    save("y_stamp", y_stamp, np.float32)

    xt = torch.from_numpy(x_norm)[None].to(device)              # (1, L, 6)
    xs = torch.from_numpy(x_stamp)[None].to(device)             # (1, L, 5)

    # ----------------------------------------------- Stage 1: tokenizer.encode
    print("\nStage 1  tokenizer.encode (half=True -> [s1_ids, s2_ids])")
    z_idx = tok.encode(xt, half=True)
    s1_ids = z_idx[0] if isinstance(z_idx, (list, tuple)) else z_idx
    s2_ids = z_idx[1] if isinstance(z_idx, (list, tuple)) else None
    save("s1_ids", s1_ids.cpu().numpy(), np.int32)
    if s2_ids is not None:
        save("s2_ids", s2_ids.cpu().numpy(), np.int32)

    # --------------------------------- Stage 2: predictor decode_s1 / decode_s2
    print("\nStage 2  predictor logits (deterministic, teacher-forced)")
    try:
        s1_logits, context = mdl.decode_s1(s1_ids, s2_ids, xs)
        save("s1_logits_full", s1_logits.cpu().numpy(), np.float32)
        save("s1_logits_last", s1_logits[0, -1].cpu().numpy(), np.float32)
        # deterministic condition for s2: argmax of the last-step s1 logits
        s1_pick = s1_logits[:, -1:].argmax(dim=-1)              # (1, 1)
        save("s1_pick_last", s1_pick.cpu().numpy(), np.int32)
        try:
            s2_logits = mdl.decode_s2(context, s1_pick)
            save("s2_logits_last", s2_logits.reshape(-1).cpu().numpy(), np.float32)
        except Exception as e:
            manifest["notes"].append(f"decode_s2 capture failed: {e!r} "
                                     f"(check arg convention vs auto_regressive_inference)")
            print("  [warn] decode_s2 failed:", repr(e))
    except Exception as e:
        manifest["notes"].append(f"decode_s1 capture failed: {e!r}")
        print("  [warn] decode_s1 failed:", repr(e))

    # ------------------------------------ Stage 3: tokenizer.decode round-trip
    print("\nStage 3  tokenizer.decode (reconstruction of the context window)")
    try:
        recon = tok.decode([s1_ids, s2_ids], half=True)        # (1, L, 6) normalized
        recon_np = recon[0].cpu().numpy()
        save("recon_norm", recon_np, np.float32)
        save("recon_denorm", recon_np * (x_std + args.eps) + x_mean, np.float32)
    except Exception as e:
        manifest["notes"].append(f"tokenizer.decode capture failed: {e!r}")
        print("  [warn] tokenizer.decode failed:", repr(e))

    # ------------------------------- Stage 4: seeded end-to-end forecast (opt.)
    if not args.skip_forecast:
        print("\nStage 4  seeded end-to-end forecast (STOCHASTIC, seed-pinned)")
        try:
            from model import KronosPredictor
            torch.manual_seed(args.seed)
            predictor = KronosPredictor(mdl, tok, device=device,
                                        max_context=512, clip=args.clip)
            pred_df = predictor.predict(
                df=hist_df, x_timestamp=x_ts, y_timestamp=y_ts,
                pred_len=P, T=1.0, top_k=0, top_p=0.9,
                sample_count=1, verbose=False)
            save("forecast_seeded", pred_df[feat_cols].values, np.float32)
            manifest["notes"].append(
                "forecast_seeded uses torch.manual_seed(seed), sample_count=1, "
                "T=1.0, top_p=0.9 -- reproducible in PyTorch but NOT a bit-exact "
                "MQL5 target; use the logits-level refs for exact verification.")
        except Exception as e:
            manifest["notes"].append(f"seeded forecast failed: {e!r}")
            print("  [warn] forecast failed:", repr(e))

    # ----------------------------------------------------------- write manifest
    manifest["scalars"] = {
        "lookback": L, "pred_len": P, "clip": args.clip, "eps": args.eps,
        "seed": args.seed, "feature_order": feat_cols, "stamp_order": stamp_cols,
        "std": "population (ddof=0)",
        "tokenizer_config": tok_cfg, "predictor_config": mdl_cfg,
    }
    with open(os.path.join(args.out, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)

    print(f"\nDone. {len(manifest['arrays'])} arrays written to {args.out}")
    if manifest["notes"]:
        print("Notes / warnings:")
        for n in manifest["notes"]:
            print("  -", n)


if __name__ == "__main__":
    main()