"""Stage B: cache teacher quantile forecasts (TimesFM-2.5-200M + Chronos-2) for distillation.

Run in the teacher venv (~/tvenv, system torch):
  ~/tvenv/bin/python -m modular_ts.teacher_gen --source synth --n 200000 --out teach/synth
  ~/tvenv/bin/python -m modular_ts.teacher_gen --source gep   --n 100000 --out teach/gep --gep_selection gep_selection.json

Each shard teach/<src>/shard_XXXX.npz holds
  ctx  (n, C) float32   context, NaN = unobserved (left-padded for short contexts)
  fut  (n, F) float32   ground-truth continuation (synthetic only; NaN for GEP = teacher-only input)
  q_tfm, q_cr2 (n, H, 9) float32  teacher quantiles 0.1..0.9 for the first H future steps
Teacher weights: google/timesfm-2.5-200m-pytorch (Apache-2.0), amazon/chronos-2 (Apache-2.0).
GiftEvalPretrain inputs use the same GIFT-Eval/fev/TIME leakage filter as run1 (realdata.filter_datasets).
NEVER uses Timestral weights.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import time

import numpy as np
import torch

Q9 = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]


def interp_nans(x: np.ndarray) -> np.ndarray:
    m = np.isfinite(x)
    if m.all():
        return x
    idx = np.arange(len(x))
    return np.interp(idx, idx[m], x[m]).astype(np.float32)


class Teachers:
    def __init__(self, C, H, dev="cuda", tfm_bs=256, cr_bs=512):
        import timesfm
        from chronos import Chronos2Pipeline
        torch.set_float32_matmul_precision("high")
        self.H = H
        self.cr = Chronos2Pipeline.from_pretrained("amazon/chronos-2", device_map=dev)
        self.cr_bs = cr_bs
        self.tfm = timesfm.TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch")
        self.tfm.compile(timesfm.ForecastConfig(max_context=C, max_horizon=H, normalize_inputs=True, per_core_batch_size=tfm_bs,
                                                use_continuous_quantile_head=True, force_flip_invariance=True,
                                                infer_is_positive=True, fix_quantile_crossing=True))

    def __call__(self, ctx: np.ndarray):
        """ctx (n, C) with NaN -> q_tfm, q_cr2 (n, H, 9)."""
        seqs = []
        for row in ctx:
            f = np.isfinite(row)
            first = int(np.argmax(f)) if f.any() else len(row) - 1
            seqs.append(row[first:])
        with torch.no_grad():
            qs, _ = self.cr.predict_quantiles([torch.from_numpy(s) for s in seqs], prediction_length=self.H,
                                              quantile_levels=Q9, batch_size=self.cr_bs)
        q_cr2 = torch.stack([q[0] for q in qs]).float().cpu().numpy()
        _, qt = self.tfm.forecast(horizon=self.H, inputs=[interp_nans(s) for s in seqs])
        q_tfm = np.asarray(qt)[:, :self.H, 1:10].astype(np.float32)
        return q_tfm, q_cr2


def synth_windows(gen, n, C, F, rng):
    vals, roles = gen.group_batch(n, 1)
    v = vals[:, 0].float().cpu().numpy()
    T = C + F
    # context length: 50% full C, else log-uniform in [48, C]; 15% of series get iid missing values
    clen = np.where(rng.random(n) < 0.5, C, np.exp(rng.uniform(math.log(48), math.log(C), n))).astype(int)
    obs = np.ones((n, T), bool)
    obs[np.arange(T)[None] < (C - clen)[:, None]] = False
    miss = rng.random(n) < 0.15
    pm = rng.uniform(0, 0.1, n)[:, None]
    obs &= ~(miss[:, None] & (rng.random((n, T)) < pm))
    v = np.where(obs, v, np.nan).astype(np.float32)
    return v[:, :C], v[:, C:]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", choices=["synth", "gep"], required=True)
    ap.add_argument("--n", type=int, default=100000)
    ap.add_argument("--shard", type=int, default=8192)
    ap.add_argument("--C", type=int, default=2048)
    ap.add_argument("--F", type=int, default=1024)
    ap.add_argument("--H", type=int, default=512)
    ap.add_argument("--out", required=True)
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--gep_selection", default="gep_selection.json")
    ap.add_argument("--max_minutes", type=float, default=0)
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    dev = torch.device("cuda")
    rng = np.random.default_rng(args.seed)
    T = args.C + args.F
    if args.source == "synth":
        from modular_ts.synth import SynthGen
        gen = SynthGen(T, dev, seed=args.seed)
    else:
        from modular_ts.realdata import FastGiftEvalPretrain
        gep = FastGiftEvalPretrain(T, args.shard * 2 * (T // 32), dev, json.load(open(args.gep_selection)), seed=args.seed,
                                   mv_frac=0.0, short_ctx_p=0.3)
        json.dump(gep.manifest(), open(os.path.join(args.out, "gep_manifest.json"), "w"), indent=1)
    tea = Teachers(args.C, args.H)
    t0, done, k = time.time(), 0, 0
    while done < args.n:
        if args.max_minutes and (time.time() - t0) / 60 > args.max_minutes:
            print("time limit", flush=True)
            break
        p = os.path.join(args.out, f"shard_{k:04d}.npz")
        if os.path.exists(p):
            done += args.shard; k += 1
            continue
        if args.source == "synth":
            gen = SynthGen(T, dev, seed=args.seed * 1000 + k)  # per-shard seed -> resumable without duplicates
            rng = np.random.default_rng(args.seed * 1000 + k)
            ctx, fut = synth_windows(gen, args.shard, args.C, args.F, rng)
        else:
            gep.g.manual_seed(args.seed * 1000 + k); gep.rng.seed(args.seed * 1000 + k)
            b = gep.next_batch()
            v = torch.where(b["observed"][:, 0], b["values"][:, 0], torch.full_like(b["values"][:, 0], float("nan"))).cpu().numpy()
            ok = np.isfinite(v[:, :args.C]).sum(1) >= 48
            v = v[ok][:args.shard]
            ctx, fut = v[:, :args.C], np.full((len(v), args.F), np.nan, np.float32)
        ts = time.time()
        q_tfm, q_cr2 = tea(ctx)
        good = np.isfinite(q_tfm).all((1, 2)) & np.isfinite(q_cr2).all((1, 2))
        np.savez(p, ctx=ctx[good], fut=fut[good], q_tfm=q_tfm[good].astype(np.float32), q_cr2=q_cr2[good].astype(np.float32))
        done += int(good.sum()); k += 1
        print(json.dumps({"shard": k, "n": int(good.sum()), "dropped": int((~good).sum()), "teach_s": round(time.time() - ts, 1),
                          "done": done, "min": round((time.time() - t0) / 60, 2)}), flush=True)
    print("TEACH_DONE", done, flush=True)


if __name__ == "__main__":
    main()