#!/usr/bin/env python3
"""Resolve Multi-VSL (WACV 2025) class ids to Vietnamese glosses, front view only.

Inputs
  Multi-VSL_WACV_2025/data/1_1000_label.numbers      id (1..1000) -> Vietnamese word
  Multi-VSL_WACV_2025/data/label_1_1000/*.csv        filename -> class label (0..999)
  WACV-Data-SLR/Data.zip                             87,817 mp4 (labels join 84,764)

Verified alignment: `label + 1 == numbers id` (checked against the word range each
recording session encodes in its filename: 0.057% violations, vs 0.339% for the
1-based alternative).

Outputs (in Multi-VSL_WACV_2025/data/)
  glosses_1_1000.csv        label, id, word, url            -- the resolved lexicon
  front_{split}.csv         name, label, word               -- front (center) view only
  lexicon_coverage.json     overlap with the Full_TriVis gloss vocabulary

Note on views: Multi-VSL names its frontal view `center`; there is no `left`/`right`
ambiguity to resolve, so "front view only" == the `*_center_ord1.csv` splits.

Note on the matching unit: a Full_TriVis gloss is `|`-separated, and each item is ONE
sign that may span several Vietnamese words ("ông bà", "bãi cỏ"). So coverage is
computed per pipe-separated ITEM, not per whitespace token.
"""
import argparse
import csv
import json
import os
import re
import unicodedata
from collections import Counter

REPO = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
MV = os.path.join(REPO, 'Multi-VSL_WACV_2025', 'data')
LAB = os.path.join(MV, 'label_1_1000')


def norm(s):
    """Normalize a gloss for matching: NFC, lowercase, drop parentheticals/punctuation."""
    s = unicodedata.normalize('NFC', str(s)).lower().strip()
    s = re.sub(r'\([^)]*\)', ' ', s)          # 'mét (m)' -> 'mét', 'hay là (hoặc là)' -> 'hay là'
    s = re.sub(r'[^\w\s]', ' ', s, flags=re.UNICODE)
    return re.sub(r'\s+', ' ', s).strip()


GRAMMAR_NOTES = {'đại từ', 'giới từ', 'động từ', 'danh từ', 'tính từ', 'trạng từ',
                 '2 người', '3 người', 'số nhiều', 'số ít'}


def variants(word):
    """All surface forms a Multi-VSL entry can match.

    Entries are richly annotated, and each convention needs handling:
      'cùng / với (giới từ)'      -> 'cùng', 'với'      ('/' = alternatives)
      'một ít/ một chút'          -> 'một ít', 'một chút'
      'hay là (hoặc là)'          -> 'hay là', 'hoặc là' (paren = alternative)
      'họ (2 người) (đại từ)'     -> 'họ'                (paren = grammar note, dropped)
      'có … không?'               -> 'có không'          (ellipsis = slot)
    """
    w = unicodedata.normalize('NFC', str(word)).lower()
    out = set()

    # parenthetical content is either an alternative form or a grammatical note
    for alt in re.findall(r'\(([^)]*)\)', w):
        a = norm(alt)
        if a and a not in GRAMMAR_NOTES and not a.isdigit():
            out.add(a)

    base = re.sub(r'\([^)]*\)', ' ', w)          # strip all parentheticals
    for piece in re.split(r'[/;,]', base):       # '/' etc. separate alternatives
        n = norm(piece)
        if n:
            out.add(n)

    out.add(norm(base))
    return {v for v in out if v}


def load_numbers(path):
    from numbers_parser import Document
    rows = Document(path).sheets[0].tables[0].rows(values_only=True)
    hdr = [str(c) for c in rows[0]]
    out = {}
    for r in rows[1:]:
        if r[0] is None:
            continue
        out[int(r[0])] = {'word': str(r[1]).strip(),
                          'url': str(r[2]).strip() if len(r) > 2 and r[2] else ''}
    return out, hdr


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--numbers', default=os.path.join(MV, '1_1000_label.numbers'))
    ap.add_argument('--offset', type=int, default=1, help='numbers id = csv label + offset')
    ap.add_argument('--trivis-index', default='./dataset/VSL_upper/train_index.json')
    ap.add_argument('--trivis-csv', default=os.path.join(REPO, 'Full_TriVis', 'split_lab_front.csv'))
    args = ap.parse_args()

    words, hdr = load_numbers(args.numbers)
    print(f'numbers table columns: {hdr}')
    print(f'ids {min(words)}..{max(words)}  ({len(words)} present)')
    missing = [i for i in range(1, 1001) if i not in words]
    if missing:
        print(f'!! ids with no word: {missing}  -> those classes stay unlabelled')

    # ---------------- resolved lexicon ----------------
    out_lex = os.path.join(MV, 'glosses_1_1000.csv')
    with open(out_lex, 'w', newline='', encoding='utf-8') as f:
        w = csv.writer(f)
        w.writerow(['label', 'numbers_id', 'word', 'url'])
        for lab in range(1000):
            e = words.get(lab + args.offset)
            w.writerow([lab, lab + args.offset, e['word'] if e else '', e['url'] if e else ''])
    print(f'wrote {out_lex}')

    label2word = {lab: words[lab + args.offset]['word']
                  for lab in range(1000) if lab + args.offset in words}

    # ---------------- front-view-only splits ----------------
    counts = {}
    for split in ('train', 'val', 'test'):
        src = os.path.join(LAB, f'{split}_1_1000_center_ord1.csv')
        rows = list(csv.DictReader(open(src)))
        dst = os.path.join(MV, f'front_{split}.csv')
        with open(dst, 'w', newline='', encoding='utf-8') as f:
            w = csv.writer(f)
            w.writerow(['name', 'label', 'word'])
            for r in rows:
                lab = int(r['label'])
                w.writerow([r['name'], lab, label2word.get(lab, '')])
        counts[split] = len(rows)
        print(f'wrote {dst}  ({len(rows)} front-view clips)')

    # ---------------- coverage against Full_TriVis ----------------
    # matching unit = pipe-separated gloss item (one sign), not whitespace token
    tri_items = Counter()
    with open(args.trivis_csv, newline='', encoding='utf-8') as f:
        for r in csv.DictReader(f):
            for it in str(r['Sign_sentence']).split('|'):
                n = norm(it)
                if n:
                    tri_items[n] += 1

    mv_index = {}
    for lab, wd in label2word.items():
        for v in variants(wd):
            mv_index.setdefault(v, lab)

    covered = {k: v for k, v in tri_items.items() if k in mv_index}
    tot_types, tot_tokens = len(tri_items), sum(tri_items.values())
    cov_types, cov_tokens = len(covered), sum(covered.values())

    print(f'\n--- Full_TriVis gloss coverage by the Multi-VSL 1000-word lexicon ---')
    print(f'Full_TriVis distinct gloss items (signs): {tot_types}')
    print(f'  covered by lexicon: {cov_types}  ({100*cov_types/tot_types:.1f}% of types)')
    print(f'Full_TriVis total gloss occurrences: {tot_tokens}')
    print(f'  covered: {cov_tokens}  ({100*cov_tokens/tot_tokens:.1f}% of tokens)')
    print(f'Multi-VSL words matched at least once: '
          f'{len(set(mv_index[k] for k in covered))}/1000')

    print('\ntop covered signs:', [w for w, _ in Counter(covered).most_common(12)])
    unc = Counter({k: v for k, v in tri_items.items() if k not in mv_index})
    print('top UNcovered signs:', [w for w, _ in unc.most_common(12)])

    res = {'offset': args.offset, 'front_view_counts': counts,
           'lexicon_words': len(label2word), 'unlabelled_ids': missing,
           'trivis_gloss_types': tot_types, 'covered_types': cov_types,
           'type_coverage_pct': round(100 * cov_types / tot_types, 2),
           'trivis_gloss_tokens': tot_tokens, 'covered_tokens': cov_tokens,
           'token_coverage_pct': round(100 * cov_tokens / tot_tokens, 2)}
    with open(os.path.join(MV, 'lexicon_coverage.json'), 'w') as f:
        json.dump(res, f, indent=2, ensure_ascii=False)
    print(f"\nwrote {os.path.join(MV, 'lexicon_coverage.json')}")


if __name__ == '__main__':
    main()