# -------------------------------------------------------------------------
# Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual
# property and proprietary rights in and to this software, related
# documentation and any modifications thereto.  Any use, reproduction,
# disclosure or distribution of this software and related documentation
# without an express license agreement from NVIDIA CORPORATION is strictly
# prohibited.
#
# Written by Jiarui Xu
# -------------------------------------------------------------------------

import os
import os.path as osp
import argparse

import pandas as pd
import sqlite3
import pandas as pd
import os.path as osp
from urllib.parse import unquote
import re
from datadings.tools import locate_files
from yfcc100m.vars import FILES
from yfcc100m.convert_metadata import download_db
from pandarallel import pandarallel

pandarallel.initialize(progress_bar=True)


def key2path(key):
    img_path = osp.join(key[0:3], key[3:6], key + '.jpg')
    return img_path


def clean_caption(line):
    line = unquote(str(line))
    line = remove_html_tags(line)
    return line.replace('\n', ' ').replace('+', ' ')


def remove_html_tags(text):
    """Remove html tags from a string"""
    clean = re.compile('<.*?>')
    return re.sub(clean, '', text)

def parse_args():
    parser = argparse.ArgumentParser(description='Create YFCC subset sql db and tsv')
    parser.add_argument('--input-dir', help='input sql db file directory')
    parser.add_argument('--output-dir', help='output tsv directory')
    parser.add_argument(
        '--subset', help='subset of data to use', default='yfcc100m_subset_data.tsv')
    args = parser.parse_args()

    return args


def main():
    args = parse_args()

    files = locate_files(FILES, args.input_dir)
    # download DB file with AWS tools
    download_db(files)

    fullset_name = 'yfcc100m_dataset'
    subset_name = 'yfcc14m_dataset'
    conn = sqlite3.connect(osp.join(args.input_dir, 'yfcc100m_dataset.sql'))
    # get column names
    # some settings that hopefully speed up the queries
    # conn.execute(f'PRAGMA query_only = YES')
    conn.execute(f'PRAGMA journal_mode = OFF')
    conn.execute(f'PRAGMA locking_mode = EXCLUSIVE')
    conn.execute(f'PRAGMA page_size = 4096')
    conn.execute(f'PRAGMA mmap_size = {4*1024*1024}')
    conn.execute(f'PRAGMA cache_size = 10000')

    print('reading subset data')
    subset_df = pd.read_csv(args.subset, sep='\t', usecols=[1, 2], names=['photoid', 'photo_hash'], index_col='photoid')
    subset_df.to_sql(subset_name, con=conn, if_exists='replace')

    print('overwriting with subset')
    select_query = f'select {fullset_name}.*, {subset_name}.photo_hash from {fullset_name} inner join {subset_name} on {fullset_name}.photoid = {subset_name}.photoid'
    new_name = 'yfcc100m_dataset_new'
    print('creating new table')
    conn.execute(f'drop table if exists {new_name}')
    conn.execute(' '.join([f'create table {new_name} as ', select_query]))
    print(f'droping {fullset_name}')
    conn.execute(f'drop table if exists {fullset_name}')
    print(f'droping {subset_name}')
    conn.execute(f'drop table if exists {subset_name}')
    print(f'renaming {new_name} to {fullset_name}')
    conn.execute(f'alter table {new_name} rename to {fullset_name}')
    print('vacuuming db')
    conn.execute('vacuum')

    print(f'Loading dataframe from SQL')
    anno_df = pd.read_sql(f'select * from {fullset_name}', con=conn)
    print(f'Loaded dataframe from SQL: \n{anno_df.head()}')
    print(f'Length: \n{len(anno_df)}')
    print(f'generating filepath')
    anno_df['file'] = anno_df['photo_hash'].parallel_map(key2path)
    anno_df['caption'] = anno_df['description'].parallel_map(clean_caption)
    anno_df = anno_df[['file', 'caption']]
    print(f'Generated dataframe: \n{anno_df.head()}')

    print('saving subset as tsv')
    os.makedirs(args.output_dir, exist_ok=True)
    anno_df.to_csv(osp.join(args.output_dir, 'yfcc14m_dataset.tsv'), sep='\t', index=False)
    conn.close()


if __name__ == '__main__':
    main()