# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


import os
import argparse
import pandas as pd
import sys


WORKDIR_ROOT = os.environ.get('WORKDIR_ROOT', None)

if WORKDIR_ROOT is None or  not WORKDIR_ROOT.strip():
    print('please specify your working directory root in OS environment variable WORKDIR_ROOT. Exitting..."')
    sys.exit(-1)

def load_langs(path):
    with open(path) as fr:
        langs = [l.strip() for l in fr]
    return langs



def load_sentences(raw_data, split, direction):
    src, tgt = direction.split('-')
    src_path = f"{raw_data}/{split}.{direction}.{src}"
    tgt_path = f"{raw_data}/{split}.{direction}.{tgt}"
    if os.path.exists(src_path) and os.path.exists(tgt_path):
        return [(src, open(src_path).read().splitlines()), (tgt, open(tgt_path).read().splitlines())]
    else:
        return []

def swap_direction(d):
    src, tgt = d.split('-')
    return f'{tgt}-{src}'

def get_all_test_data(raw_data, directions, split='test'):
    test_data = [ 
        x
        for dd in directions
        for d in [dd, swap_direction(dd)]
        for x in load_sentences(raw_data, split, d)
    ]
    # all_test_data = {s for _, d in test_data for s in d}
    all_test_data = {}
    for lang, d in test_data:
        for s in d:
            s = s.strip()
            lgs = all_test_data.get(s, set())
            lgs.add(lang)
            all_test_data[s] = lgs
    return all_test_data, test_data


def check_train_sentences(src_path, tgt_path, direction, all_test_data, mess_up_train={}):
    # src, tgt = direction.split('-')
    print(f'check training data for {direction} in {src_path} and {tgt_path}')
    size = 0
    overlapped_size_counted_dup = 0
    if not os.path.exists(tgt_path) or not os.path.exists(src_path):
        return mess_up_train, size, overlapped_size_counted_dup

    with open(src_path) as f, open(tgt_path) as g:
        for src_line, tgt_line in zip(f, g):
            s = src_line.strip()
            t = tgt_line.strip()
            size += 1
            if  s in all_test_data:
                langs = mess_up_train.get(s, set())
                langs.add(direction)
                mess_up_train[s] = langs
                overlapped_size_counted_dup += 1
            if t in all_test_data:
                langs = mess_up_train.get(t, set())
                langs.add(direction)
                mess_up_train[t] = langs 
                overlapped_size_counted_dup += 1
    print(f'{direction}: size={size}, overlapped={overlapped_size_counted_dup}')
    return mess_up_train, size, overlapped_size_counted_dup

def check_train_all(raw_data, directions, all_test_data):
    mess_up_train = {}
    data_sizes = {}
    # raw_data = '~chau/data-bin/MineBART/multilingual_mined_100M/en_XX/et_EE-en_XX/all.{en_XX, et_EE}'
    print(f'checking training data againsts # {len(all_test_data)} sentences')
    print(f'example test data: ', [s for i, s in enumerate(all_test_data.keys()) if i < 10])
    for direction in directions:
        src, tgt = direction.split('-')
        path = f'{raw_data}/en_XX/{direction}/all'
        src_path = f'{path}.{src}'
        tgt_path = f'{path}.{tgt}'
        print(f'checking {src_path} {tgt_path}')
        _, size, overlapped_size_counted_dup = check_train_sentences(src_path, tgt_path, direction, all_test_data, mess_up_train)
        data_sizes[direction] = (size, overlapped_size_counted_dup)
    return mess_up_train, data_sizes




def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--folder", type=str, required=True,
                        help="the data folder ")
    parser.add_argument("--test-data", type=str, required=True,
                        help="the test data folder ")                        
    parser.add_argument('--directions', type=str, default=None, required=False)

    args = parser.parse_args()    
    directions = args.directions.split(',')
    directions = sorted(set(directions))

    results = []
    # print(f'checking where {args.split} split data are in training')
    # print(f'direction\tcommon_count\tsrc common\ttgt common\tfrom_size\tto_size')
    raw_data = args.folder
    all_test_data, test_data = get_all_test_data(args.test_data, directions, split='test')
    mess_up_train, data_sizes = check_train_all(raw_data, directions, all_test_data)
    print(data_sizes)


if __name__ == "__main__":
    main()