# -*- coding: utf-8 -*-
"""Ensemble AdaptNER predictions by weighted majority vote per entity-type column.

Each input is a Konooz submission file: token + 21 space-separated B/I/O columns
(one channel per entity type), sentences separated by blank lines. All inputs must
share the same tokens in the same order. For every token and every one of the 21
columns, we take a weighted vote over the models' tags ({O, B-TYPE, I-TYPE}) and
emit the highest-weight tag (ties break toward the earlier/higher-weight model).

Usage
-----
    python ensemble_predictions.py \
        --model FILE1.txt:0.74 --model FILE2.txt:0.72 --model FILE3.txt:0.71 \
        --out namaa_pred_ensemble.txt --zip submission_ensemble.zip

Weights are optional (default 1.0 each -> plain majority vote). Pass them after a
colon; the natural choice is each model's Konooz-dev micro-F1.
"""
from __future__ import annotations

import argparse
import io
import os
import zipfile
from collections import defaultdict

N_COLS = 21


def read_rows(path):
    """Return list of rows; each row is None (blank line) or [token, *21tags]."""
    rows = []
    for line in io.open(path, encoding="utf-8"):
        s = line.rstrip("\n")
        if not s.strip():
            rows.append(None)
        else:
            rows.append(s.split())
    return rows


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model", action="append", required=True,
                    help="path[:weight] of a prediction file; repeatable")
    ap.add_argument("--out", default="namaa_pred_ensemble.txt")
    ap.add_argument("--zip", dest="zip_path", default="")
    ap.add_argument("--mode", choices=["vote", "union"], default="vote",
                    help="vote = weighted majority per column; union = any non-O wins (max recall)")
    args = ap.parse_args()

    models, weights = [], []
    for spec in args.model:
        if ":" in spec and spec.rsplit(":", 1)[1].replace(".", "", 1).isdigit():
            path, w = spec.rsplit(":", 1); w = float(w)
        else:
            path, w = spec, 1.0
        models.append(read_rows(path)); weights.append(w)
        print(f"[in] {os.path.basename(path)}  weight={w}")

    n = len(models[0])
    for i, m in enumerate(models):
        if len(m) != n:
            raise ValueError(f"file {i} has {len(m)} lines, expected {n} (inputs must align)")

    out_rows = []
    for r in range(n):
        if models[0][r] is None:
            out_rows.append(None); continue
        token = models[0][r][0]
        cols = []
        for c in range(N_COLS):
            score = defaultdict(float)
            for m, w in zip(models, weights):
                tag = m[r][1 + c]
                if args.mode == "union":
                    # any B/I outweighs O
                    score[tag] += (w + 10.0 if tag != "O" else w)
                else:
                    score[tag] += w
            # pick max-weight tag; ties -> earlier model order via stable max on (weight, -firstseen)
            best = max(score.items(), key=lambda kv: kv[1])[0]
            cols.append(best)
        out_rows.append([token, *cols])

    with io.open(args.out, "w", encoding="utf-8", newline="\n") as f:
        for row in out_rows:
            f.write("" if row is None else " ".join(row))
            f.write("\n")
    n_ent = sum(1 for row in out_rows if row for t in row[1:] if t.startswith("B-"))
    print(f"[out] {args.out}  (entities: {n_ent})")

    if args.zip_path:
        with zipfile.ZipFile(args.zip_path, "w", zipfile.ZIP_DEFLATED) as z:
            z.write(args.out, arcname=os.path.basename(args.out))
        print(f"[zip] {args.zip_path}")


if __name__ == "__main__":
    main()