#!/usr/bin/env python3
"""E0: select the real universe (seeded, nested samples) from FinData coverage and export 1-minute bars for 2024.
Outputs: experiments/results/e0_universe.json (symbols, dataset version, export ids, sha256), cache Parquet files
under $RAID_CACHE_DIR (never committed). Requires FINDATA_API_KEY_FILE and the compiled core on sys.path."""
import os, sys, json, time, yaml, hashlib
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python"))
import numpy as np, raid
cfg = yaml.safe_load(open(os.path.join(os.path.dirname(__file__), "..", "configs", "e0_universe.yaml"))); cfg["start"] = str(cfg["start"]); cfg["end"] = str(cfg["end"])
out_path = os.path.join(os.path.dirname(__file__), "..", "results", "e0_universe.json")
t0 = time.time()
cov = raid.findata.coverage(cfg["dataset_id"], cfg["min_sessions_total"])
active = [c for c in cov if c["d0"] <= cfg["start"] and c["d1"] >= cfg["end"] and (not cfg["exclude_delisted_suffix"] or not c["symbol"].endswith("-DELISTED"))]
active.sort(key=lambda c: c["symbol"])
rng = np.random.default_rng(cfg["seed"]); perm = rng.permutation(len(active))
nmax = max(cfg["sample_sizes"]); chosen = [active[i]["symbol"] for i in perm[:nmax]]
samples = {str(n): chosen[:n] for n in cfg["sample_sizes"]}
print(f"coverage symbols {len(cov)}, active through {cfg['year']}: {len(active)}, chosen {nmax}; first {chosen[:5]}")
exports = []; versions = set()
for i in range(0, nmax, cfg["symbols_per_export"]):
    syms = chosen[i:i + cfg["symbols_per_export"]]; t1 = time.time()
    e = raid.findata.export(syms, cfg["start"], cfg["end"], dataset_id=cfg["dataset_id"])
    versions.add(e["dataset_version"]); exports.append({"symbols": syms, **{k: e[k] for k in ("export_id", "path", "sha256", "dataset_version", "rows", "bytes", "from_cache")}})
    print(f"export {i // cfg['symbols_per_export'] + 1}/{(nmax + cfg['symbols_per_export'] - 1) // cfg['symbols_per_export']}: {e['rows']:,} rows {e['bytes'] / 1e6:.1f} MB cache={e['from_cache']} {time.time() - t1:.1f}s", flush=True)
json.dump({"config": cfg, "n_coverage": len(cov), "n_active": len(active), "samples": samples, "exports": exports, "dataset_versions": sorted(versions),
           "total_rows": sum(e["rows"] for e in exports), "total_bytes": sum(e["bytes"] for e in exports), "seconds": time.time() - t0,
           "universe_sha256": hashlib.sha256(",".join(chosen).encode()).hexdigest()}, open(out_path, "w"), indent=1)
print("wrote", out_path, "rows", sum(e["rows"] for e in exports), "seconds", round(time.time() - t0, 1))