#!/usr/bin/env python3
"""Build a 3-view (front/left/right) Full_TriVis CSV that `prepare_vsl_data.py` can pack.

`Full_TriVis/split_lab_front.csv` covers only the front view: 24,151 clips.
`Full_TriVis/Label/Dataset_full_lab.csv` is the master label file and holds all three
views of the same sentences -- 72,453 rows = 24,151 x 3, scene `lab` throughout -- which
matches the 72,453 lab `.npz` files under `Full_TriVis/skeleton/` exactly.

The split is INHERITED, not redrawn. `splits.json` used `split_key = Sign_sentence`, so
assigning each master row the split of its `Sign_sentence` keeps every view of a sentence
in the same split. That preserves the text-disjoint property (a sentence never spans
splits), and it keeps the new test set a strict superset of the old one, so front-only
results stay directly comparable -- the front subset of the new test split IS the old test
split.

npz path = "Full_TriVis/skeleton/" + Sentence_video_path with .mp4 -> .npz.
"""
import argparse
import collections
import csv
import os

REPO = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
COLS = ['split', 'npz_path', 'Sentence', 'Sign_sentence', 'Category', 'group', 'view',
        'scene', 'ID_sentence', 'Sentence_video_path']


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--front-csv', default=os.path.join(REPO, 'Full_TriVis',
                                                       'split_lab_front.csv'))
    ap.add_argument('--master-csv', default=os.path.join(REPO, 'Full_TriVis', 'Label',
                                                        'Dataset_full_lab.csv'))
    ap.add_argument('--skeleton-root', default='Full_TriVis/skeleton')
    ap.add_argument('--repo-root', default=REPO)
    ap.add_argument('--views', default='front,left,right')
    ap.add_argument('--out', default=os.path.join(REPO, 'Full_TriVis',
                                                 'split_lab_3view.csv'))
    ap.add_argument('--check-exists', action='store_true',
                    help='stat every npz and drop the missing ones (slow but safe)')
    args = ap.parse_args()

    want = set(args.views.split(','))

    with open(args.front_csv, newline='', encoding='utf-8') as f:
        front = list(csv.DictReader(f))
    split_of = {r['Sign_sentence']: r['split'] for r in front}
    print(f'{len(front)} front rows -> {len(split_of)} distinct Sign_sentence with a split')

    with open(args.master_csv, newline='', encoding='utf-8') as f:
        master = list(csv.DictReader(f))
    print(f'{len(master)} master rows, views {sorted({r["view"] for r in master})}')

    rows, missing_split, missing_file = [], 0, 0
    for r in master:
        if r['view'] not in want:
            continue
        sp = split_of.get(r['Sign_sentence'])
        if sp is None:
            missing_split += 1
            continue
        npz = os.path.join(args.skeleton_root,
                           r['Sentence_video_path'].replace('.mp4', '.npz'))
        if args.check_exists and not os.path.exists(os.path.join(args.repo_root, npz)):
            missing_file += 1
            continue
        rows.append({'split': sp, 'npz_path': npz, 'Sentence': r['Sentence'],
                     'Sign_sentence': r['Sign_sentence'], 'Category': r['Category'],
                     'group': r['group'], 'view': r['view'], 'scene': r['scene'],
                     'ID_sentence': r['ID_sentence'],
                     'Sentence_video_path': r['Sentence_video_path']})

    with open(args.out, 'w', newline='', encoding='utf-8') as f:
        w = csv.DictWriter(f, fieldnames=COLS)
        w.writeheader()
        w.writerows(rows)

    print(f'wrote {args.out}: {len(rows)} rows '
          f'(no split {missing_split}, missing npz {missing_file})')
    print('  split x view:')
    c = collections.Counter((r['split'], r['view']) for r in rows)
    for sp in ('train', 'val', 'test'):
        line = '  '.join(f'{v} {c[(sp, v)]}' for v in sorted(want))
        print(f'    {sp:<6} {line}   total {sum(c[(sp, v)] for v in want)}')

    # leak check, the same one splits.json reports
    by = collections.defaultdict(set)
    for r in rows:
        by[r['Sign_sentence']].add(r['split'])
    span = sum(1 for v in by.values() if len(v) > 1)
    print(f'  leakage_check: {"PASS" if span == 0 else "FAIL"}: {span} Sign_sentence '
          f'values span multiple splits')


if __name__ == '__main__':
    main()