#!/usr/bin/env python3
"""
Score the canine-demo human evaluation.

Usage:
    python3 human_eval/score_human_eval.py merged.csv

`merged.csv` holds both annotators' rows (one row per clip per annotator):

    clip_id,annotator,environment,activity,intention,interpretation,notes
    clip01,A,2,1,2,4,cat visible; movement brief
    clip01,B,2,2,2,4,
    ...

Attribute scores (environment/activity/intention) are 0/1/2:
    2 = Correct, 1 = Plausible, 0 = Incorrect
Interpretation is 1-5 (1 = nonsensical, 5 = accurate & informative).

Reports, per attribute and pooled:
    Accuracy       - mean fraction of ratings == 2 (avg. over annotators)
    Acceptability  - mean fraction of ratings >= 1 (avg. over annotators)
    Agreement      - % of clips where both annotators gave the same score
    Fleiss' kappa  - chance-corrected agreement (2 raters, 3 categories)
plus the interpretation mean +/- std, and a LaTeX block for the paper.
"""
import sys
import csv
import math
from collections import defaultdict

ATTRS = ['environment', 'activity', 'intention']


def load(path):
    rows = []
    with open(path, newline='') as f:
        for r in csv.DictReader(f):
            r = {k.strip().lower(): (v.strip() if isinstance(v, str) else v)
                 for k, v in r.items() if k}
            rows.append(r)
    return rows


def to_int(x):
    try:
        return int(float(x))
    except (TypeError, ValueError):
        return None


def fleiss_kappa(items, categories):
    """items: list of dicts {category: count} for each subject (row of the
    N x k table). All subjects must have the same number of ratings."""
    n = None
    table = []
    for it in items:
        row = [it.get(c, 0) for c in categories]
        s = sum(row)
        if s == 0:
            continue
        if n is None:
            n = s
        if s != n:            # unequal ratings per subject -> undefined
            return None
        table.append(row)
    if not table or n < 2:
        return None
    N = len(table)
    # P_j : proportion of all assignments to category j
    col = [sum(table[i][j] for i in range(N)) for j in range(len(categories))]
    total = N * n
    Pj = [c / total for c in col]
    # P_i : extent of agreement for subject i
    Pi = [(sum(x * x for x in table[i]) - n) / (n * (n - 1)) for i in range(N)]
    Pbar = sum(Pi) / N
    Pe = sum(p * p for p in Pj)
    if abs(1 - Pe) < 1e-12:
        return None
    return (Pbar - Pe) / (1 - Pe)


def summarize(scores_by_ann, agree_pairs, kappa_items):
    # scores_by_ann: list of all attribute scores (pooled over annotators)
    acc = sum(1 for s in scores_by_ann if s == 2) / len(scores_by_ann)
    accept = sum(1 for s in scores_by_ann if s >= 1) / len(scores_by_ann)
    agree = (sum(1 for a, b in agree_pairs if a == b) / len(agree_pairs)
             if agree_pairs else float('nan'))
    kappa = fleiss_kappa(kappa_items, [0, 1, 2])
    return acc, accept, agree, kappa


def main():
    if len(sys.argv) != 2:
        print(__doc__)
        sys.exit(1)
    rows = load(sys.argv[1])

    # index[clip][attr] = {annotator: score}
    index = defaultdict(lambda: defaultdict(dict))
    interp = defaultdict(dict)      # index[clip][annotator] = interp score
    for r in rows:
        clip = r.get('clip_id') or r.get('clip')
        ann = (r.get('annotator') or r.get('rater') or '').upper()
        if not clip or not ann:
            continue
        for a in ATTRS:
            v = to_int(r.get(a))
            if v is not None:
                index[clip][a][ann] = v
        iv = to_int(r.get('interpretation'))
        if iv is not None:
            interp[clip][ann] = iv

    clips = sorted(index)
    print(f'\nClips scored: {len(clips)}')
    annset = sorted({ann for c in index for a in ATTRS for ann in index[c][a]})
    print(f'Annotators:   {", ".join(annset)}\n')

    # ---- per-attribute ----
    per_attr = {}
    pooled_scores, pooled_pairs, pooled_items = [], [], []
    for a in ATTRS:
        scores, pairs, items = [], [], []
        for c in clips:
            d = index[c][a]
            vals = list(d.values())
            scores += vals
            cnt = defaultdict(int)
            for v in vals:
                cnt[v] += 1
            items.append(cnt)
            if len(vals) == 2:
                pairs.append(tuple(vals))
        per_attr[a] = summarize(scores, pairs, items)
        pooled_scores += scores
        pooled_pairs += pairs
        pooled_items += items

    overall = summarize(pooled_scores, pooled_pairs, pooled_items)

    # ---- interpretation ----
    ivals = [v for c in interp for v in interp[c].values()]
    imean = sum(ivals) / len(ivals) if ivals else float('nan')
    istd = (math.sqrt(sum((x - imean) ** 2 for x in ivals) / len(ivals))
            if ivals else float('nan'))

    def fk(x):
        return 'n/a' if x is None else f'{x:.3f}'

    print(f'{"Attribute":<14}{"Accuracy":>10}{"Accept.":>10}'
          f'{"Agree%":>9}{"Fleiss k":>10}')
    print('-' * 53)
    for a in ATTRS:
        acc, acp, ag, k = per_attr[a]
        print(f'{a:<14}{acc*100:>9.1f}%{acp*100:>9.1f}%{ag*100:>8.1f}%{fk(k):>10}')
    acc, acp, ag, k = overall
    print('-' * 53)
    print(f'{"OVERALL":<14}{acc*100:>9.1f}%{acp*100:>9.1f}%{ag*100:>8.1f}%{fk(k):>10}')
    print(f'\nInterpretation (1-5):  mean = {imean:.2f} +/- {istd:.2f}  '
          f'(n={len(ivals)})')

    # ---- LaTeX ----
    print('\n% ---------- paste into evaluation.tex ----------')
    print('% Table 1 numbers (attribute rows + overall)')
    for a in ATTRS:
        acc, acp, ag, k = per_attr[a]
        print(f'% {a.capitalize():<14}& {acc*100:.1f}\\% & {acp*100:.1f}\\% '
              f'& {ag*100:.1f}\\% & {fk(k)} \\\\')
    acc, acp, ag, k = overall
    print(f'\\textbf{{Overall}} & {acc*100:.1f}\\% & {acp*100:.1f}\\% '
          f'& {ag*100:.1f}\\% & {fk(k)} \\\\')
    print(f'% Interpretation quality: ${imean:.2f}\\pm{istd:.2f}$ '
          f'(1--5 scale, N={len(clips)}, two annotators)')
    print('% ------------------------------------------------\n')


if __name__ == '__main__':
    main()
