"""P1: point-in-time stratified universe -> data/manifests/universe_main.json (+ pilot subset)."""
from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from chronos_x.data.client import FinDataClient  # noqa: E402
from chronos_x.data.universe import UniverseConfig, build_candidates, stratified_sample, manifest, load_sic_map  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-target", type=int, default=300)
    ap.add_argument("--sec-tickers", default="/tmp/wp41_ext/sec_company_tickers.json")
    ap.add_argument("--sec-submissions", default="/tmp/wp41_ext/sec_submissions")
    ap.add_argument("--n-pilot", type=int, default=20)
    ap.add_argument("--n-feasibility", type=int, default=50)
    ap.add_argument("--min-coverage", type=float, default=0.0)
    ap.add_argument("--coverage-file", default=None)
    ap.add_argument("--out", default="universe_main.json")
    a = ap.parse_args()
    cfg = UniverseConfig(n_target=a.n_target, min_coverage=a.min_coverage, coverage_file=a.coverage_file)
    c = FinDataClient(log_path=ROOT / "results" / "logs" / "provenance_universe.jsonl")
    sic = load_sic_map(Path(a.sec_tickers), Path(a.sec_submissions))
    print(f"SIC map: {len(sic)} tickers")
    cand = build_candidates(c, cfg, sic)
    print(f"candidates: {cand.shape[0]} (stocks {int((cand.asset_class=='stock').sum())}, etf {int((cand.asset_class=='etf').sum())}, delisted {int(cand.delisted.sum())})")
    sample = stratified_sample(cand, cfg)
    ds_version = None
    try:
        ds_version = c.dataset("market_bars.stock.1m.raw").get("version")
    except Exception:
        pass
    # pilot and feasibility subsets drawn from the pooled (non-held-out) instruments, stratified by asset class
    rng = random.Random(cfg.seed + 7)
    pooled = sample[sample["role"] == "pooled"]
    pilot = sorted(rng.sample(sorted(pooled[pooled.asset_class == "stock"]["symbol"]), int(a.n_pilot * 0.75)) +
                   rng.sample(sorted(pooled[pooled.asset_class == "etf"]["symbol"]), a.n_pilot - int(a.n_pilot * 0.75)))
    rest = [s for s in sorted(pooled["symbol"]) if s not in pilot]
    feas = sorted(pilot + rng.sample(rest, a.n_feasibility - len(pilot)))
    man = manifest(sample, cand, cfg, ds_version, {"pilot_symbols": pilot, "feasibility_symbols": feas})
    out = ROOT / "data" / "manifests"
    out.mkdir(parents=True, exist_ok=True)
    (out / a.out).write_text(json.dumps(man, indent=1))
    cand.drop(columns=["name"]).to_csv(out / a.out.replace("universe", "candidates").replace(".json", ".csv"), index=False)
    summary = sample.groupby(["asset_class", "adv_tercile", "vol_tercile"]).size().unstack(fill_value=0)
    print(summary)
    print(sample.groupby(["asset_class", "group"]).size())
    print(f"selected {len(sample)}, held out {int((sample.role=='heldout_instrument').sum())}, delisted {int(sample.delisted.sum())}")
    print("pilot:", pilot)


if __name__ == "__main__":
    main()