#!/usr/bin/env python3
"""Fan the existing BARTpho gloss predictions out to all three views.

`output_vsl/text2gloss/pred_full_trivis.json` is keyed by front-view clip name. The
left/right clips of the same recording carry the SAME sentence and the SAME reference
gloss, only a different camera, so no model run is needed -- the prediction is a function
of the sentence alone. Rebuilding the map sentence -> prediction and re-emitting it per
clip of the 3-view CSV keeps the predictions bit-identical to the front-only evaluation,
which is what makes the two comparable.
"""
import argparse
import collections
import csv
import json
import os

REPO = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')


def signs(s):
    parts = [t.strip() for t in str(s).split('|')] if '|' in str(s) else str(s).split()
    return [p for p in parts if p]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--pred', default='output_vsl/text2gloss/pred_full_trivis.json')
    ap.add_argument('--csv', default=os.path.join(REPO, 'Full_TriVis',
                                                  'split_lab_3view.csv'))
    ap.add_argument('--out', default='output_vsl/text2gloss/pred_3view.json')
    args = ap.parse_args()

    with open(args.pred, encoding='utf-8') as f:
        pred = json.load(f)
    by_sent = {}
    for v in pred.values():
        by_sent.setdefault(v['sentence'], (v['pred_gloss'], v['pred_signs']))
    print(f'{len(pred)} front clips -> {len(by_sent)} distinct sentences with a prediction')

    out, miss = {}, 0
    with open(args.csv, newline='', encoding='utf-8') as f:
        rows = list(csv.DictReader(f))
    for r in rows:
        name = os.path.splitext(os.path.basename(r['npz_path']))[0]
        hit = by_sent.get(r['Sentence'].strip())
        if hit is None:
            miss += 1
            continue
        pg, ps = hit
        out[name] = {'split': r['split'], 'view': r['view'],
                     'sentence': r['Sentence'].strip(),
                     'pred_gloss': pg, 'pred_signs': ps,
                     'ref_gloss': ' '.join(signs(r['Sign_sentence'])),
                     'ref_signs': signs(r['Sign_sentence'])}

    os.makedirs(os.path.dirname(args.out) or '.', exist_ok=True)
    with open(args.out, 'w', encoding='utf-8') as f:
        json.dump(out, f, ensure_ascii=False)
    print(f'wrote {args.out}: {len(out)} clips ({miss} sentences had no prediction)')
    c = collections.Counter((v['split'], v['view']) for v in out.values())
    for sp in ('train', 'val', 'test'):
        print(f"  {sp:<6} " + '  '.join(f'{v} {c[(sp, v)]}'
                                        for v in ('front', 'left', 'right')))


if __name__ == '__main__':
    main()