"""Optional real-data source: Salesforce/GiftEvalPretrain (OFF by default).

GiftEvalPretrain's card says Apache-2.0 but "for research purposes only"; the user has
not chosen between synthetic-only and synthetic+GiftEvalPretrain, so train.py keeps this
behind `--data synthetic+gifteval` (default `--data synthetic`).

Contamination filter: every GiftEvalPretrain sub-dataset whose normalised name matches
(or is an explicit alias of) a dataset in the GIFT-Eval test pool, fev-bench tasks or the
TIME benchmark is dropped. The resolved include/exclude lists are written next to each
checkpoint (data_manifest.json) for the model card's contamination audit.

Data layout on the Hub: one folder per sub-dataset holding HF `save_to_disk` arrow
shards (`data-0000k-of-0000n.arrow`) with columns item_id, start, freq, target
(1-D or 2-D list). We download shards lazily, memory-map them, and crop random windows
on a background thread.
"""

from __future__ import annotations

import queue
import math
import random
import re
import threading

import numpy as np
import torch

REPO = "Salesforce/GiftEvalPretrain"

GIFT_EVAL_TEST = [
    "m4_yearly", "m4_quarterly", "m4_monthly", "m4_weekly", "m4_daily", "m4_hourly", "electricity", "solar",
    "hospital", "covid_deaths", "us_births", "saugeenday", "temperature_rain_with_missing",
    "kdd_cup_2018_with_missing", "car_parts_with_missing", "restaurant", "hierarchical_sales", "LOOP_SEATTLE",
    "SZ_TAXI", "M_DENSE", "ett1", "ett2", "jena_weather", "bitbrains_fast_storage", "bitbrains_rnd",
    "bizitobs_application", "bizitobs_service", "bizitobs_l2c",
]
FEV_BENCH = [
    "ETT", "LOOP_SEATTLE", "M_DENSE", "SZ_TAXI", "australian_tourism", "bizitobs_l2c", "boomlet", "ecdc_ili",
    "entsoe", "epf_be", "epf_de", "epf_fr", "epf_np", "epf_pjm", "ercot", "favorita_stores", "favorita_transactions",
    "fred_md_2025", "fred_qd_2025", "gvar", "hermes", "hierarchical_sales", "hospital", "hospital_admissions",
    "jena_weather", "kdd_cup_2022", "m5", "proenfo_gfc12", "proenfo_gfc14", "proenfo_gfc17", "redset", "restaurant",
    "rohlik_orders", "rohlik_sales", "rossmann", "solar", "solar_with_weather", "uci_air_quality", "uk_covid_nation",
    "uk_covid_utla", "us_consumption", "walmart", "world_co2_emissions", "world_life_expectancy", "world_tourism",
]
TIME_BENCH = [
    "Australia_Solar", "Auto_Production_SF", "CPHL", "Coastal_T_S", "Commodity_Import", "Commodity_Production",
    "Crypto", "ECDC_COVID", "EWELD_Load", "Finland_Traffic", "Global_Influenza", "Global_Price", "Housing_Inventory",
    "JOLTS", "Job_Claims", "MetroPT-3", "NE_China_Wind", "Oil_Price", "Online_Retail_2_UCI", "OpenElectricity_NEM",
    "Port_Activity", "SG_Carpark", "SG_PM25", "SG_Weather", "Smart_Manufacturing", "Supply_Chain_Customer",
    "Supply_Chain_Location", "US_Labor", "US_Term_Structure", "Uncertainty_1M", "Vehicle_Sales", "Vehicle_Supply",
    "WUI_Global", "Water_Quality_Darwin", "azure2019", "current_velocity", "epf_electricity_price",
]
# GiftEvalPretrain names that are the same source as a benchmark dataset under another name.
ALIASES_EXCLUDE = {
    "favorita_sales": "fev favorita_stores", "favorita_transactions": "fev favorita_transactions",
    "m5": "fev m5", "kdd2022": "fev kdd_cup_2022", "gfc12_load": "fev proenfo_gfc12",
    "gfc14_load": "fev proenfo_gfc14", "gfc17_load": "fev proenfo_gfc17", "fred_md": "fev fred_md_2025",
    "solar_power": "GIFT-Eval solar (conservative)", "ecdc_ili": "fev ecdc_ili",
}


def _norm(s: str) -> str:
    s = s.lower()
    s = re.sub(r"_with_missing$", "", s)
    s = re.sub(r"_(\d+[a-z]+|daily|hourly|weekly|monthly|yearly|quarterly)$", "", s)
    return re.sub(r"[^a-z0-9]", "", s)


def list_pretrain_datasets() -> list[str]:
    from huggingface_hub import HfApi
    tree = HfApi().list_repo_tree(REPO, repo_type="dataset")
    return sorted(t.path for t in tree if "." not in t.path)


def filter_datasets(names: list[str]) -> tuple[list[str], dict[str, str]]:
    bench = {}
    for pool, lst in (("GIFT-Eval", GIFT_EVAL_TEST), ("fev-bench", FEV_BENCH), ("TIME", TIME_BENCH)):
        for b in lst:
            bench[_norm(b)] = f"{pool} {b}"
    keep, dropped = [], {}
    for n in names:
        k = _norm(n)
        hit = ALIASES_EXCLUDE.get(n) or bench.get(k)
        if hit is None:
            hit = next((v for bk, v in bench.items() if len(bk) >= 5 and (bk in k or k in bk)), None)
        if hit:
            dropped[n] = hit
        else:
            keep.append(n)
    return keep, dropped


class GiftEvalPretrainStream:
    """Random fixed-length windows from GiftEvalPretrain, batched as (B, V, T) like SynthStream.

    Univariate series give V=1 groups; multivariate series give up to `max_v` variates.
    `datasets` is optional; used only to memory-map arrow shards.
    """

    def __init__(self, T: int, tokens_per_batch: int, device="cpu", seed=0, datasets_subset: list[str] | None = None,
                 max_shards_per_dataset: int = 2, max_v: int = 16, patch=32, min_ctx=4, cache_dir=None, prefetch=4):
        from datasets import Dataset as HFDataset
        from huggingface_hub import HfApi, hf_hub_download

        names = datasets_subset or list_pretrain_datasets()
        self.keep, self.dropped = filter_datasets(names)
        self.T, self.N, self.tok, self.patch = T, T // patch, tokens_per_batch, patch
        self.dev, self.min_ctx, self.max_v = torch.device(device), min_ctx, max_v
        self.rng = random.Random(seed)
        api = HfApi()
        self.tables = []
        for name in self.keep:
            files = sorted(f for f in api.list_repo_files(REPO, repo_type="dataset") if f.startswith(name + "/") and f.endswith(".arrow"))
            for f in files[:max_shards_per_dataset]:
                p = hf_hub_download(REPO, f, repo_type="dataset", cache_dir=cache_dir)
                self.tables.append((name, HFDataset.from_file(p)))
        assert self.tables, "no GiftEvalPretrain shards left after filtering"
        self.q: queue.Queue = queue.Queue(maxsize=prefetch)
        threading.Thread(target=self._worker, daemon=True).start()

    def manifest(self) -> dict:
        return {"source": REPO, "included": self.keep, "excluded": self.dropped}

    def _window(self):
        name, tab = self.rng.choice(self.tables)
        row = tab[self.rng.randrange(len(tab))]
        x = np.asarray(row["target"], dtype=np.float32)
        x = x[None] if x.ndim == 1 else x
        if x.shape[0] > self.max_v:
            x = x[self.rng.sample(range(x.shape[0]), self.max_v)]
        L = x.shape[1]
        if L >= self.T:
            st = self.rng.randrange(L - self.T + 1)
            w = x[:, st : st + self.T]
        else:  # left-pad with NaN (unobserved)
            w = np.concatenate([np.full((x.shape[0], self.T - L), np.nan, np.float32), x], 1)
        return w

    def _make(self):
        # bucket by V: draw one window, then fill the batch with windows of the same V
        first = self._window()
        V = first.shape[0]
        Bg = max(1, self.tok // (V * self.N))
        ws = [first]
        tries = 0
        while len(ws) < Bg and tries < 20 * Bg:
            w = self._window()
            tries += 1
            if w.shape[0] == V:
                ws.append(w)
        arr = np.stack(ws)
        obs = np.isfinite(arr)
        vals = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
        return vals, obs

    def _worker(self):
        while True:
            self.q.put(self._make())

    def next_batch(self, V=None) -> dict:
        vals, obs = self.q.get()
        Bg, V, _ = vals.shape
        s = torch.randint(self.min_ctx, self.N - 1, (Bg,))
        return {
            "values": torch.from_numpy(vals).to(self.dev, non_blocking=True),
            "observed": torch.from_numpy(obs).to(self.dev, non_blocking=True),
            "roles": torch.zeros(Bg, V, dtype=torch.long, device=self.dev),
            "s": s.to(self.dev),
        }


class MixedStream:
    """Draw from `real` with probability `real_frac`, else from `synth`."""

    def __init__(self, synth, real, real_frac: float, seed=0):
        self.synth, self.real, self.p = synth, real, real_frac
        self.rng = random.Random(seed)

    def next_batch(self, V=None):
        return self.real.next_batch() if self.rng.random() < self.p else self.synth.next_batch(V)


# ---------------------------------------------------------------------------------------------
# Fast GPU-resident stream (used for run1). The thread/row-based stream above cannot keep up
# with ~12+ steps/s x 1024 windows; this one loads a size-capped selection of shards once into
# one flat float32 tensor on the GPU and crops whole batches with a single vectorised gather.
# ---------------------------------------------------------------------------------------------

def _arrow_channels(path):
    """Return (values float32 1-D, ch_off int64, ch_len int64, row_first_ch, row_nvar)."""
    import pyarrow as pa
    try:
        tab = pa.ipc.open_stream(path).read_all()
    except Exception:
        tab = pa.ipc.open_file(path).read_all()
    col = tab.column("target").combine_chunks()
    t = col.type
    if pa.types.is_list(t) and not (pa.types.is_list(t.value_type) or pa.types.is_fixed_size_list(t.value_type)):
        inner, nvar = col, np.ones(len(col), np.int64)
    elif pa.types.is_fixed_size_list(t):
        inner, nvar = col.flatten(), np.full(len(col), t.list_size, np.int64)
    else:  # list<list<float>>
        inner = col.flatten()
        off = np.asarray(col.offsets)
        nvar = (off[1:] - off[:-1]).astype(np.int64)
    off = np.asarray(inner.offsets).astype(np.int64)
    vals = inner.values.to_numpy(zero_copy_only=False).astype(np.float32, copy=False)
    base = off[0]
    vals = vals[base:off[-1]]
    ch_off = off[:-1] - base
    ch_len = off[1:] - off[:-1]
    row_first = np.concatenate([[0], np.cumsum(nvar)[:-1]]).astype(np.int64)
    return vals, ch_off, ch_len, row_first, nvar


class FastGiftEvalPretrain:
    def __init__(self, T, tokens_per_batch, device, selection: dict, seed=0, patch=32, cache_dir=None,
                 mv_frac=0.3, max_v=16, short_ctx_p=0.3, ds_temp=0.5, workers=16):
        from concurrent.futures import ThreadPoolExecutor
        from huggingface_hub import hf_hub_download
        names = sorted(selection)
        self.keep, self.dropped = filter_datasets(names)
        assert set(self.keep) == set(names), f"selection contains filtered datasets: {set(names) - set(self.keep)}"
        self.T, self.P, self.N, self.tok = T, patch, T // patch, tokens_per_batch
        self.dev, self.max_v, self.mv_frac, self.short_ctx_p = torch.device(device), max_v, mv_frac, short_ctx_p
        self.g = torch.Generator(device=self.dev); self.g.manual_seed(seed + 777)
        jobs = [(n, f) for n in names for f in selection[n]]
        with ThreadPoolExecutor(workers) as ex:
            paths = list(ex.map(lambda nf: hf_hub_download(REPO, nf[1], repo_type="dataset", cache_dir=cache_dir), jobs))
        print(f"[gep] downloaded {len(paths)} shards", flush=True)
        with ThreadPoolExecutor(8) as ex:
            parsed = list(ex.map(_arrow_channels, paths))
        vals_all, ch_off_all, ch_len_all, ch_ds, mv_rows = [], [], [], [], []  # mv_rows: (ds, first_ch_global, nvar)
        base, nch = 0, 0
        ds_idx = {n: i for i, n in enumerate(names)}
        self.ds_points = np.zeros(len(names))
        for (n, _), (v, co, cl, rf, nv) in zip(jobs, parsed):
            keep = cl >= 8
            d = ds_idx[n]
            vals_all.append(v)
            ch_off_all.append(co + base); ch_len_all.append(cl); ch_ds.append(np.full(len(cl), d, np.int64))
            m = (nv >= 2)
            if m.any():
                mv_rows.append(np.stack([np.full(m.sum(), d), rf[m] + nch, nv[m]], 1))
            self.ds_points[d] += cl[keep].sum()
            base += len(v); nch += len(cl)
        flat = np.concatenate(vals_all)
        ch_off, ch_len, ch_ds = map(np.concatenate, (ch_off_all, ch_len_all, ch_ds))
        good = ch_len >= 8
        print(f"[gep] {len(names)} datasets, {nch} channels, {len(flat)/1e9:.2f}B points", flush=True)
        self.flat = torch.from_numpy(flat).to(self.dev)
        del flat, vals_all
        # univariate channel pool grouped by dataset
        order = np.argsort(ch_ds[good], kind="stable")
        gi = np.nonzero(good)[0][order]
        self.u_off = torch.from_numpy(ch_off[gi]).to(self.dev)
        self.u_len = torch.from_numpy(ch_len[gi]).to(self.dev)
        cnt = np.bincount(ch_ds[gi], minlength=len(names))
        self.u_start = torch.from_numpy(np.concatenate([[0], np.cumsum(cnt)[:-1]])).to(self.dev)
        self.u_cnt = torch.from_numpy(cnt).to(self.dev)
        w = np.where(cnt > 0, self.ds_points ** ds_temp, 0.0)
        self.u_w = torch.tensor(w / w.sum(), dtype=torch.float, device=self.dev)
        self.weights = {n: float(x) for n, x in zip(names, w / w.sum())}
        # multivariate rows
        self.ch_off_t = torch.from_numpy(ch_off).to(self.dev)
        self.ch_len_t = torch.from_numpy(ch_len).to(self.dev)
        self.mv = None
        if mv_rows:
            mvr = np.concatenate(mv_rows)
            mvr = mvr[ch_len[mvr[:, 1]] >= 8]
            self.mv = mvr
            self.mv_by_v = {}
            for V in (2, 4, 8, 16):
                sub = mvr[mvr[:, 2] >= V]
                if len(sub) == 0:
                    continue
                dcnt = np.bincount(sub[:, 0], minlength=len(names))
                o = np.argsort(sub[:, 0], kind="stable"); sub = sub[o]
                wv = np.where(dcnt > 0, self.ds_points ** ds_temp, 0.0)
                self.mv_by_v[V] = dict(first=torch.from_numpy(sub[:, 1]).to(self.dev), nvar=torch.from_numpy(sub[:, 2]).to(self.dev),
                                       start=torch.from_numpy(np.concatenate([[0], np.cumsum(dcnt)[:-1]])).to(self.dev),
                                       cnt=torch.from_numpy(dcnt).to(self.dev), w=torch.tensor(wv / wv.sum(), dtype=torch.float, device=self.dev))
        self.rng = random.Random(seed)
        self.ar = torch.arange(T, device=self.dev)

    def manifest(self):
        try:
            excl = filter_datasets(list_pretrain_datasets())[1]
        except Exception as e:  # offline
            excl = f"unavailable: {e!r}"
        return {"source": REPO, "selection": "size-capped shard selection (gep_selection.json)", "included": self.keep,
                "excluded_by_contamination_filter": excl, "sampling_weights": self.weights, "points_per_dataset": {n: int(x) for n, x in zip(sorted(self.weights), self.ds_points)}}

    def _rand(self, n):
        return torch.rand(n, generator=self.g, device=self.dev)

    def _pick(self, w, start, cnt, n):
        d = torch.multinomial(w, n, replacement=True, generator=self.g)
        return start[d] + (self._rand(n) * cnt[d]).long().clamp_max(cnt[d] - 1)

    def _crop(self, off, L):
        """off, L: (R,) -> values (R,T), observed (R,T). Same start for rows of one group handled by caller."""
        T = self.T
        st = torch.where(L >= T, (self._rand(L.shape[0]) * (L - T + 1)).long(), L - T)
        return st

    def _gather(self, off, L, st, keep_len):
        pos = st[:, None] + self.ar[None]
        valid = (pos >= 0) & (pos < L[:, None]) & (self.ar[None] >= (self.T - keep_len)[:, None])
        idx = (off[:, None] + pos.clamp(min=0)).clamp_max(self.flat.numel() - 1)
        v = self.flat[idx]
        obs = valid & torch.isfinite(v)
        return torch.where(obs, v, torch.zeros_like(v)), obs

    def _s(self, obs_bv):
        # obs_bv (B,V,T): first observed patch per group -> min-context split after it
        B = obs_bv.shape[0]
        any_p = obs_bv.view(B, -1, self.N, self.P).any(-1).any(1)  # (B,N)
        fs = torch.where(any_p.any(1), any_p.float().argmax(1), torch.zeros(B, dtype=torch.long, device=self.dev))
        span = (self.N - 2 - fs).clamp_min(1)
        s = fs + 1 + (self._rand(B) * span).long()
        return s.clamp(1, self.N - 2)

    def _keep_len(self, n):
        k = torch.full((n,), self.T, device=self.dev, dtype=torch.long)
        short = self._rand(n) < self.short_ctx_p
        kl = torch.exp(math.log(64) + self._rand(n) * (math.log(self.T) - math.log(64))).long()
        return torch.where(short, kl, k)

    @torch.no_grad()
    def next_batch(self, V=None):
        if self.mv is not None and self.rng.random() < self.mv_frac:
            V = self.rng.choice(sorted(self.mv_by_v))
            m = self.mv_by_v[V]
            Bg = max(1, self.tok // (V * self.N))
            r = self._pick(m["w"], m["start"], m["cnt"], Bg)
            first, nvar = m["first"][r], m["nvar"][r]
            j0 = (self._rand(Bg) * (nvar - V + 1)).long()
            ch = first[:, None] + j0[:, None] + torch.arange(V, device=self.dev)[None]  # (Bg,V)
            off, L = self.ch_off_t[ch], self.ch_len_t[ch]
            Lmin = L.min(1).values
            st = self._crop(None, Lmin)
            kl = self._keep_len(Bg)
            v, o = self._gather(off.reshape(-1), L.reshape(-1), st.repeat_interleave(V), kl.repeat_interleave(V))
            v, o = v.view(Bg, V, self.T), o.view(Bg, V, self.T)
        else:
            V = 1
            Bg = max(1, self.tok // self.N)
            c = self._pick(self.u_w, self.u_start, self.u_cnt, Bg)
            off, L = self.u_off[c], self.u_len[c]
            st = self._crop(off, L)
            v, o = self._gather(off, L, st, self._keep_len(Bg))
            v, o = v[:, None], o[:, None]
        return {"values": v, "observed": o, "roles": torch.zeros(Bg, V, dtype=torch.long, device=self.dev), "s": self._s(o)}


if __name__ == "__main__":
    names = list_pretrain_datasets()
    keep, dropped = filter_datasets(names)
    print(f"{len(names)} sub-datasets; keep {len(keep)}; drop {len(dropped)}")
    for k, v in dropped.items():
        print(f"  DROP {k:40s} <- {v}")