#!/usr/bin/env python
"""Stacking-ensemble predictions.

A per-class linear stacking head over the ensemble's per-item probability outputs.
Given the bundled probability features (probs/), it produces the label vector and
checks it against the bundled prediction file.

    python predict.py --track t1
    python predict.py --track t2

"""
import argparse
import json
import os

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))


def predict(track):
    spec = json.load(open(os.path.join(HERE, f"{track}_stacker.json")))
    lab = spec["label_order"]                       
    cols = spec["npy_column_order"]                 #
    perm = [cols.index(c) for c in lab]
    n = spec["n_items"]

    feats = []
    for name in spec["sources"]:                    # fixed source order
        a = np.load(os.path.join(HERE, "probs", name)).astype(np.float64)[:, perm]
        a = a / a.sum(1, keepdims=True)
        feats.append(a)
    X = np.stack(feats).transpose(1, 0, 2).reshape(n, -1)   # n x 3J, matches training

    W = np.array(spec["coef"]); b = np.array(spec["intercept"])
    logits = X @ W.T + b
    pred = [lab[i] for i in logits.argmax(1)]

    ref_path = os.path.join(HERE, f"{track}_perclass_stack.txt")
    ref = open(ref_path).read().split()
    ok = pred == ref
    out = os.path.join(HERE, f"{track}_perclass_stack_out.txt")
    open(out, "w").write("\n".join(pred) + "\n")
    print(f"[{track}] {len(pred)} labels -> {out}  |  matches bundled predictions: {ok}")
    if not ok:
        raise SystemExit(f"[{track}] mismatch vs bundled predictions")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--track", choices=["t1", "t2"], required=True)
    predict(ap.parse_args().track)