# Export visual groups as a CSV which helps humanities researchers
# interpret and use the results.
#
# Author : Abhishek Dutta <adutta@robots.ox.ac.uk>
# Date   : 2022-10-03

import argparse
import sqlite3

def main():
    parser = argparse.ArgumentParser(prog="export-vgroups-as-csv",
                                     description="Export visual groups of VISE as a CSV file")

    parser.add_argument("--metadata-db",
                        required=True,
                        type=str,
                        help="sqlite database containing metadata (generated by VISE)")
    parser.add_argument("--groups",
                       required=True,
                       type=str,
                       help="sqlite database containing visual groups (generated by VISE)")
    parser.add_argument("--out-csv",
                        required=True,
                        type=str,
                        help="write groups to this file as a CSV")

    args = parser.parse_args()
    metadata_connection = sqlite3.connect(args.metadata_db)
    metadata_cur = metadata_connection.cursor()

    group_connection = sqlite3.connect(args.groups)
    group_cur = group_connection.cursor()

    fid_filename_map = {}
    for row in metadata_cur.execute("SELECT file_id, filename FROM file_metadata"):
        fid_filename_map[ row[0] ] = row[1]

    with open(args.out_csv, 'w') as f:
        f.write('set_id,set_size,set_files\n')
        for row in group_cur.execute("SELECT set_id, set_size, member_id_list FROM vgroup"):
            fid_list = row[2].split(',')
            filename_list = []
            for fid_str in fid_list:
                fid = int(fid_str)
                filename_list.append( '"' + fid_filename_map[fid] + '"' )
            f.write("%d,%d,%s\n" % (row[0], row[1], ','.join(filename_list)) )
        print('Saved visual groups to %s' % (args.out_csv))

if __name__ == '__main__':
    main()
