"""Local reimplementation of the shared-task metric, so our offline scores
line up with the leaderboard.

    Favg2 = (F1_Favor + F1_Against) / 2      -- the actual metric, no None
    Favg3 = (F1_Favor + F1_Against + F1_None) / 3
"""
import argparse
import sys

import pandas as pd
from sklearn.metrics import accuracy_score, f1_score

VALID_LABELS = ["Against", "Favor", "None"]
LABEL2ID = {"Against": 0, "Favor": 1, "None": 2}


def compute_metrics(gold_labels, pred_labels):
    y_true = [LABEL2ID[x] for x in gold_labels]
    y_pred = [LABEL2ID[x] for x in pred_labels]

    f_against = f1_score(y_true, y_pred, labels=[0], average="macro")
    f_favor = f1_score(y_true, y_pred, labels=[1], average="macro")
    f_none = f1_score(y_true, y_pred, labels=[2], average="macro")

    return {
        "Favg2": (f_favor + f_against) / 2.0,
        "Favg3": (f_favor + f_against + f_none) / 3.0,
        "F_Favor": f_favor,
        "F_Against": f_against,
        "F_None": f_none,
        "Accuracy": accuracy_score(y_true, y_pred),
    }


def score(gold_df, preds, verbose=True):
    if len(preds) != len(gold_df):
        raise ValueError(
            f"Length mismatch: gold={len(gold_df)}, preds={len(preds)}"
        )

    invalid = sorted(set(preds) - set(VALID_LABELS))
    if invalid:
        raise ValueError(f"Invalid prediction labels: {invalid}")

    gold_df = gold_df.copy()
    gold_df["stance"] = gold_df["stance"].astype(str).str.strip()
    gold_df["target"] = gold_df["target"].astype(str).str.strip()
    gold_df["pred"] = preds

    per_target = {}
    for target in sorted(gold_df["target"].unique()):
        sub = gold_df[gold_df["target"] == target]
        per_target[target] = compute_metrics(
            sub["stance"].tolist(), sub["pred"].tolist()
        )

    overall = compute_metrics(
        gold_df["stance"].tolist(), gold_df["pred"].tolist()
    )

    if verbose:
        for t, m in per_target.items():
            print(
                f"  [{t}] Favg2={m['Favg2']:.4f} Favg3={m['Favg3']:.4f} "
                f"(Fav={m['F_Favor']:.4f} Agn={m['F_Against']:.4f} "
                f"None={m['F_None']:.4f}) Acc={m['Accuracy']:.4f}"
            )
        print(
            f"  OVERALL Favg2={overall['Favg2']:.4f} "
            f"Favg3={overall['Favg3']:.4f} Acc={overall['Accuracy']:.4f}"
        )

    return {"overall": overall, "per_target": per_target}


def load_gold(csv_path):
    df = pd.read_csv(csv_path, keep_default_na=False, encoding="utf-8-sig")
    df.columns = df.columns.astype(str).str.strip()
    for col in ("target", "stance"):
        if col not in df.columns:
            raise ValueError(f"Gold file must contain a '{col}' column.")
    return df


def read_preds_txt(path):
    with open(path, "r", encoding="utf-8") as f:
        preds = [line.strip() for line in f if line.strip() != ""]
    if preds and preds[0].lower() == "stance":
        preds = preds[1:]
    return preds


def validate_submission(preds, n_expected):
    if len(preds) != n_expected:
        return False, f"count {len(preds)} != expected {n_expected}"
    invalid = sorted(set(preds) - set(VALID_LABELS))
    if invalid:
        return False, f"invalid labels: {invalid}"
    return True, f"OK: {len(preds)} rows, labels in {VALID_LABELS}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold", required=True)
    ap.add_argument("--pred", required=True)
    args = ap.parse_args()

    gold_df = load_gold(args.gold)
    preds = read_preds_txt(args.pred)
    ok, msg = validate_submission(preds, len(gold_df))
    print(f"[validate] {msg}", file=sys.stderr)
    if not ok:
        sys.exit(1)
    score(gold_df, preds)


if __name__ == "__main__":
    main()