#!/usr/bin/env python3
#
# Converts unencoded extractor data to CSV files.
#
# Example:
#   python3 src/extractor/convert_to_csv.py ../data ../data_processed

import os
import sys
import pandas as pd

data_subsets = [2048, 4096]

def get_file_list(base_dir):
    data = {}
    for sub_dir in os.listdir(base_dir):
        if not os.path.isdir(os.path.join(base_dir, sub_dir)):
            continue
        for _entry in os.listdir(os.path.join(base_dir, sub_dir)):
            (_base, _ext) = os.path.splitext(os.path.basename(_entry))
            data.setdefault(os.path.join(base_dir, sub_dir, _base), {})[_ext] = os.path.join(base_dir, sub_dir, _entry)

    # filter for only entries that contain all the patch data we need
    data = {k: v for k,v in data.items()
            if (".desc" in v and ".patch" in v and ".raw_i" in v and ".raw_o" in v and
		".raw_c" in v)}

    return data


def main():
    data_dir = sys.argv[1]
    output_dir = sys.argv[2]

    dfs = []
    for (base,file) in get_file_list(data_dir).items():
        dfs.append(pd.DataFrame({'ID': [base],
                                 'Description': [open(file['.desc']).read()],
                                 'Patch': [open(file['.patch']).read()],
                                 'Input': [open(file['.raw_i']).read()],
                                 'Output': [open(file['.raw_o']).read()],
				 'Context': [open(file['.raw_c']).read()]}))
    dfs = pd.concat(dfs, ignore_index=True)

    for _size in [str(x) for x in data_subsets + ["full"]]:
        _dfs = dfs
        if _size != 'full':
            _dfs = _dfs.head(int(_size))

        _size_dir = os.path.join(output_dir, _size)
        os.makedirs(_size_dir, exist_ok=True)

        # split data
        _train = _dfs.sample(frac=0.7) # train 70%
        _eval = _dfs.drop(_train.index)
        _temp = _eval.sample(frac=0.666666) # eval 20%
        _test = _eval.drop(_temp.index) # test 10%
        _eval = _temp

        _dfs.to_csv(os.path.join(output_dir, _size, "dataset.csv"))
        _train.to_csv(os.path.join(output_dir, _size, "train.csv"))
        _eval.to_csv(os.path.join(output_dir, _size, "evaluation.csv"))
        _test.to_csv(os.path.join(output_dir, _size, "test.csv"))

        print(f"Dataset Samples ({os.path.join(output_dir, _size)}):")
        print(f"  total: {len(_dfs)}")
        print(f"  train: {len(_train)}")
        print(f"  eval:  {len(_eval)}")
        print(f"  test:  {len(_test)}\n")


if __name__ == '__main__':
    main()
