"""
Walk-forward ML combination experiment.

    PYTHONPATH=. uv run --with lightgbm python -m research.lab.run_ml --market us
"""

import argparse
import os
import pickle
import time
import warnings


from research.lab import ml
from research.lab.factor_scan import load, phase_avg
from research.lab.factors import FACTORS
from research.lab.report import fmt

warnings.filterwarnings("ignore")
OUT = os.path.join("data", "research", "results")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--market", default="us")
    ap.add_argument("--models", nargs="*", default=["lgbm", "ridge", "ic_weight"])
    ap.add_argument("--horizons", nargs="*", type=int, default=[5, 20])
    ap.add_argument("--first-test-year", type=int, default=2010)
    args = ap.parse_args()
    mk = "us" if args.market == "us" else "cn"
    p, periods = load(args.market)
    names = [n for n, f in FACTORS.items() if mk in f.markets]
    t = time.time()
    feats = ml.feature_ranks(p, names)
    print(f"{len(names)} features in {time.time() - t:.0f}s", flush=True)
    evalp = {"design": (f"{args.first_test_year}-01-01", periods["design"][1]), "holdout": periods["holdout"]}
    os.makedirs(OUT, exist_ok=True)
    for model in args.models:
        for h in args.horizons:
            t = time.time()
            score = ml.walk_forward(p, feats, model=model, horizon=h, first_test_year=args.first_test_year)
            with open(os.path.join(OUT, f"ml_{args.market}_{model}_h{h}.pkl"), "wb") as fh:
                pickle.dump(score, fh)
            for top_n in (10, 20):
                parts = []
                for per, (a, b) in evalp.items():
                    parts.append(f"{per}: {fmt(phase_avg(p, score, a, b, top_n, 5))}")
                print(f"{model:9s} h{h:<3d} top{top_n:<3d}", " | ".join(parts), f"[{time.time() - t:.0f}s]", flush=True)


if __name__ == "__main__":
    main()